From f1a2c24563be19f4b28d72c1dcfaf5275d5fb087 Mon Sep 17 00:00:00 2001
From: "Brian C. Lane" <bcl@redhat.com>
Date: Wed, 1 May 2024 18:02:31 -0700
Subject: [PATCH] worker: Add CleanupArtifacts function

This removes all artifact directories, and their contents, if there
isn't an associated Job. This is used to clean up local artifacts after
the compose job has been deleted.

Related: RHEL-60120
---
 internal/worker/server.go      | 32 +++++++++++++
 internal/worker/server_test.go | 83 +++++++++++++++++++++++++++++++++-
 2 files changed, 113 insertions(+), 2 deletions(-)

diff --git a/internal/worker/server.go b/internal/worker/server.go
index eb4ac04b3..4969b482d 100644
--- a/internal/worker/server.go
+++ b/internal/worker/server.go
@@ -353,6 +353,38 @@ func (s *Server) AllRootJobIDs(ctx context.Context) ([]uuid.UUID, error) {
 	return s.jobs.AllRootJobIDs(ctx)
 }
 
+// CleanupArtifacts removes worker artifact directories that do not have matching jobs
+// The UUID used for the artifact directory is the same as for the job that created it
+func (s *Server) CleanupArtifacts() error {
+	artifacts, err := os.ReadDir(s.config.ArtifactsDir)
+	if err != nil {
+		return err
+	}
+
+	for _, d := range artifacts {
+		if !d.IsDir() {
+			continue
+		}
+		id, err := uuid.Parse(d.Name())
+		if err != nil {
+			continue
+		}
+
+		// Is there a job with this UUID?
+		if _, _, _, _, err := s.jobs.Job(id); err != nil {
+			// No associated job, it is safe to remove the unused artifact directory
+			// and everything under it, and the ComposeRequest (if it exists)
+			_ = os.Remove(path.Join(s.config.ArtifactsDir, "ComposeRequest", id.String()+".json"))
+			err = os.RemoveAll(path.Join(s.config.ArtifactsDir, id.String()))
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}
+
 func (s *Server) OSBuildJobInfo(id uuid.UUID, result *OSBuildJobResult) (*JobInfo, error) {
 	jobInfo, err := s.jobInfo(id, result)
 	if err != nil {
diff --git a/internal/worker/server_test.go b/internal/worker/server_test.go
index 374742067..5d62d38d7 100644
--- a/internal/worker/server_test.go
+++ b/internal/worker/server_test.go
@@ -3,6 +3,7 @@ package worker_test
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"net/http"
 	"os"
@@ -32,7 +33,14 @@ import (
 )
 
 func newTestServer(t *testing.T, tempdir string, config worker.Config, acceptArtifacts bool) *worker.Server {
-	q, err := fsjobqueue.New(tempdir)
+	// NOTE: jobs and artifacts directories need to be next to each other. artifacts inside the fsjobqueue
+	// directory makes it crash when trying to list all the jobs because it isn't a UUID.
+	jobsDir := path.Join(tempdir, "jobs")
+	err := os.Mkdir(jobsDir, 0755)
+	if err != nil && !os.IsExist(err) {
+		t.Fatalf("cannot create jobs directory %s: %v", jobsDir, err)
+	}
+	q, err := fsjobqueue.New(jobsDir)
 	if err != nil {
 		t.Fatalf("error creating fsjobqueue: %v", err)
 	}
@@ -41,7 +49,7 @@ func newTestServer(t *testing.T, tempdir string, config worker.Config, acceptArt
 		artifactsDir := path.Join(tempdir, "artifacts")
 		err := os.Mkdir(artifactsDir, 0755)
 		if err != nil && !os.IsExist(err) {
-			t.Fatalf("cannot create state directory %s: %v", artifactsDir, err)
+			t.Fatalf("cannot create artifacts directory %s: %v", artifactsDir, err)
 		}
 		config.ArtifactsDir = artifactsDir
 	}
@@ -1607,3 +1615,74 @@ func TestJobHeartbeats(t *testing.T) {
 	require.Equal(t, float64(0), promtest.ToFloat64(prometheus.PendingJobs))
 	require.Equal(t, float64(0), promtest.ToFloat64(prometheus.RunningJobs))
 }
+
+func makeFakeArtifact(tempdir string, id uuid.UUID, filename string) error {
+	d := path.Join(tempdir, "artifacts", id.String())
+	err := os.Mkdir(d, 0755)
+	if err != nil {
+		return err
+	}
+	if len(filename) > 0 {
+		p := path.Join(d, filename)
+		fp, err := os.Create(p)
+		if err != nil {
+			return err
+		}
+		return fp.Close()
+	}
+
+	return nil
+}
+
+func artifactUUIDExists(tempdir string, id uuid.UUID) bool {
+	dir := path.Join(tempdir, "artifacts", id.String())
+	_, err := os.Stat(dir)
+	return !errors.Is(err, os.ErrNotExist)
+}
+
+func TestCleanupArtifacts(t *testing.T) {
+	distroStruct := newTestDistro(t)
+	arch, err := distroStruct.GetArch(test_distro.TestArchName)
+	if err != nil {
+		t.Fatalf("error getting arch from distro: %v", err)
+	}
+	imageType, err := arch.GetImageType(test_distro.TestImageTypeName)
+	if err != nil {
+		t.Fatalf("error getting image type from arch: %v", err)
+	}
+	manifest, _, err := imageType.Manifest(nil, distro.ImageOptions{Size: imageType.Size(0)}, nil, nil)
+	if err != nil {
+		t.Fatalf("error creating osbuild manifest: %v", err)
+	}
+	tempdir := t.TempDir()
+	server := newTestServer(t, tempdir, defaultConfig, true)
+	mf, err := manifest.Serialize(nil, nil, nil, nil)
+	if err != nil {
+		t.Fatalf("error creating osbuild manifest: %v", err)
+	}
+
+	jobID, err := server.EnqueueOSBuild(arch.Name(), &worker.OSBuildJob{Manifest: mf}, "")
+	require.NoError(t, err)
+
+	// Make a fake artifact for the existing jobid
+	err = makeFakeArtifact(tempdir, jobID, "not-a-real.img")
+	require.Nil(t, err)
+
+	// Make an artifact directory for a job that is already gone
+	lostJobID := uuid.New()
+	err = makeFakeArtifact(tempdir, lostJobID, "not-a-real.img")
+	require.Nil(t, err)
+
+	// Make an artifact directory with no files
+	emptyJobID := uuid.New()
+	err = makeFakeArtifact(tempdir, emptyJobID, "")
+	require.Nil(t, err)
+
+	// This should remove lostJobID and not jobID
+	err = server.CleanupArtifacts()
+	require.Nil(t, err)
+
+	assert.True(t, artifactUUIDExists(tempdir, jobID))
+	assert.False(t, artifactUUIDExists(tempdir, lostJobID))
+	assert.False(t, artifactUUIDExists(tempdir, emptyJobID))
+}