feat: add cleanup test, and isolate id mapping code and cleanup code

blob/local_store.go

@@ -8,6 +8,8 @@ import (
 	"os"
 	"path/filepath"
 
+	"strings"
+
 	"github.com/gammazero/fsutil/disk"
 )
 
@@ -128,3 +130,52 @@ func (l *LocalStore) Describe(ctx context.Context, id ID) (*Descriptor, error) {
 		ModificationTime: stat.ModTime(),
 	}, nil
 }
+
+// Lists all locally stored blob IDs. Files in the local store directory which
+// do not end in ".bin", or which cannot be parsed into a blob ID, will not be
+// included.
+func (l *LocalStore) List(ctx context.Context) ([]ID, error) {
+	var ids []ID
+
+	dir, err := os.ReadDir(l.dir)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read local store directory: %w", err)
+	}
+
+	for _, binFileEntry := range dir {
+		if binFileEntry.IsDir() {
+			// Do not include directories
+			continue
+		}
+
+		idString, isBin := strings.CutSuffix(binFileEntry.Name(), ".bin")
+		if !isBin {
+			// Do not include files that don't end in .bin
+			continue
+		}
+
+		var id ID
+		if err := id.Decode(idString); err != nil {
+			// Do not include files that cannot be parsed into a blob ID
+			continue
+		}
+
+		ids = append(ids, ID(id))
+	}
+
+	return ids, nil
+}
+
+// Removes the blob. Errors with ErrBlobNotFound if the blob does not exist.
+func (l *LocalStore) Remove(ctx context.Context, id ID) error {
+	binFileName := id.String() + ".bin"
+	if err := os.Remove(filepath.Join(l.dir, binFileName)); err != nil {
+		if errors.Is(err, os.ErrNotExist) {
+			return ErrBlobNotFound
+		}
+
+		return fmt.Errorf("failed to remove bin file '%s': %w", binFileName, err)
+	}
+
+	return nil
+}

integration/singularity/cleanup.go

@@ -0,0 +1,118 @@
+package singularity
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/filecoin-project/motion/blob"
+)
+
+type cleanupSchedulerConfig struct {
+	interval time.Duration
+}
+
+// This is run by the cleanup scheduler to determine whether to clean up a local
+// file.
+type cleanupReadyCallback func(ctx context.Context, blobID blob.ID) (bool, error)
+
+type cleanupScheduler struct {
+	cfg          cleanupSchedulerConfig
+	local        *blob.LocalStore
+	cleanupReady cleanupReadyCallback
+	closing      chan struct{}
+	closed       sync.WaitGroup
+}
+
+func newCleanupScheduler(
+	cfg cleanupSchedulerConfig,
+	local *blob.LocalStore,
+	cleanupReady cleanupReadyCallback,
+) *cleanupScheduler {
+	return &cleanupScheduler{
+		cfg:          cfg,
+		local:        local,
+		cleanupReady: cleanupReady,
+		closing:      make(chan struct{}),
+	}
+}
+
+func (cs *cleanupScheduler) start(ctx context.Context) {
+	cs.closed.Add(1)
+
+	go func() {
+		defer cs.closed.Done()
+
+		ctx, cancel := context.WithCancel(context.Background())
+		go func() {
+			<-cs.closing
+			cancel()
+		}()
+
+		ticker := time.NewTicker(cs.cfg.interval)
+		defer ticker.Stop()
+
+		// Run once immediately on startup
+		cs.cleanup(ctx)
+
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-ticker.C:
+				cs.cleanup(ctx)
+			}
+		}
+	}()
+}
+
+func (cs *cleanupScheduler) stop(ctx context.Context) error {
+	close(cs.closing)
+
+	done := make(chan struct{})
+	go func() {
+		cs.closed.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-done:
+		return nil
+	}
+}
+
+func (cs *cleanupScheduler) cleanup(ctx context.Context) error {
+	logger.Info("Starting cleanup")
+
+	ids, err := cs.local.List(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to list local blob IDs: %w", err)
+	}
+
+	var removals []blob.ID
+	for _, id := range ids {
+		cleanupReady, err := cs.cleanupReady(ctx, id)
+		if err != nil {
+			logger.Warnw("failed to check if blob is ready for cleanup, skipping for this cleanup cycle", "err", err)
+			continue
+		}
+		if cleanupReady {
+			removals = append(removals, id)
+		}
+	}
+
+	for _, blobID := range removals {
+		cs.local.Remove(ctx, blobID)
+	}
+
+	if len(removals) > 0 {
+		logger.Infow("Cleaned up unneeded local files", "count", len(removals))
+	} else {
+		logger.Info("Did not find any local files to clean up")
+	}
+
+	return nil
+}

integration/singularity/cleanup_test.go

@@ -0,0 +1,113 @@
+package singularity
+
+import (
+	"bytes"
+	"context"
+	"crypto/rand"
+	"io"
+	"os"
+	"path/filepath"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/filecoin-project/motion/blob"
+	"github.com/stretchr/testify/require"
+)
+
+func TestCleanupScheduler(t *testing.T) {
+	localDir := filepath.Join(t.TempDir(), "motion-cleanup-test")
+	require.NoError(t, os.MkdirAll(localDir, 0777))
+	local := blob.NewLocalStore(localDir)
+
+	var lock sync.Mutex
+	var cleanupNotReadyBlobs []blob.ID
+	var cleanupReadyBlobs []blob.ID
+
+	// Pushes a new blob and optionally add it to the shouldRemove array for the
+	// cleanup test
+	add := func(shouldRemove bool) {
+		data := make([]byte, 16)
+		_, err := rand.Reader.Read(data)
+		require.NoError(t, err)
+
+		desc, err := local.Put(context.Background(), io.NopCloser(bytes.NewReader(data)))
+		require.NoError(t, err)
+
+		if shouldRemove {
+			cleanupReadyBlobs = append(cleanupReadyBlobs, desc.ID)
+		} else {
+			cleanupNotReadyBlobs = append(cleanupNotReadyBlobs, desc.ID)
+		}
+	}
+
+	// Intersperse ready and non-ready blobs
+	for i := 0; i < 100; i++ {
+		add(false)
+		add(true)
+	}
+
+	// Returns true if the value is present in the shouldRemove slice
+	callback := func(ctx context.Context, blobID blob.ID) (bool, error) {
+		lock.Lock()
+		defer lock.Unlock()
+		for _, other := range cleanupReadyBlobs {
+			if blobID == other {
+				return true, nil
+			}
+		}
+		return false, nil
+	}
+
+	cfg := cleanupSchedulerConfig{
+		interval: time.Second,
+	}
+
+	// Before starting the cleanup scheduler, make sure all blobs are present
+	listBefore, err := local.List(context.Background())
+	t.Logf("length before: %v", len(listBefore))
+	require.NoError(t, err)
+	lock.Lock()
+	for _, blob := range cleanupNotReadyBlobs {
+		require.Contains(t, listBefore, blob)
+	}
+	for _, blob := range cleanupReadyBlobs {
+		require.Contains(t, listBefore, blob)
+	}
+	lock.Unlock()
+
+	// Start cleanup scheduler and check that all cleanup ready blobs are
+	// already removed before the first cleanup tick, since there should be one
+	// iteration immediately on startup
+	cleanupScheduler := newCleanupScheduler(cfg, local, callback)
+	cleanupScheduler.start(context.Background())
+
+	time.Sleep(cfg.interval / 2)
+
+	listAfterStart, err := local.List(context.Background())
+	t.Logf("length after start: %v", len(listAfterStart))
+	require.NoError(t, err)
+	lock.Lock()
+	for _, blob := range cleanupNotReadyBlobs {
+		require.Contains(t, listAfterStart, blob)
+	}
+	for _, blob := range cleanupReadyBlobs {
+		require.NotContains(t, listAfterStart, blob)
+	}
+	lock.Unlock()
+
+	// Add the rest of the blobs to cleanup ready, wait for 1 more tick, and
+	// make sure that no more blobs are left
+	lock.Lock()
+	cleanupReadyBlobs = append(cleanupReadyBlobs, cleanupNotReadyBlobs...)
+	lock.Unlock()
+
+	time.Sleep(cfg.interval)
+
+	listAfterTick, err := local.List(context.Background())
+	t.Logf("length after tick: %v", len(listAfterTick))
+	require.NoError(t, err)
+	require.Empty(t, listAfterTick)
+
+	require.NoError(t, cleanupScheduler.stop(context.Background()))
+}

integration/singularity/id_map.go

@@ -0,0 +1,85 @@
+package singularity
+
+import (
+	"errors"
+	"fmt"
+	"os"
+	"path/filepath"
+	"strconv"
+
+	"github.com/filecoin-project/motion/blob"
+)
+
+type idMap struct {
+	dir string
+}
+
+func newIDMap(dir string) *idMap {
+	return &idMap{
+		dir: dir,
+	}
+}
+
+func (im *idMap) path(blobID blob.ID) string {
+	return filepath.Join(im.dir, blobID.String()+".id")
+}
+
+// Inserts a blob ID to Singularity ID mapping.
+func (im *idMap) insert(blobID blob.ID, singularityID int64) error {
+	idFile, err := os.CreateTemp(im.dir, "motion_local_store_*.id.temp")
+	if err != nil {
+		return fmt.Errorf("failed to create temporary file: %w", err)
+	}
+	defer func() {
+		if err := idFile.Close(); err != nil {
+			logger.Debugw("Failed to close temporary file", "err", err)
+		}
+	}()
+	_, err = idFile.Write([]byte(strconv.FormatUint(uint64(singularityID), 10)))
+	if err != nil {
+		if err := os.Remove(idFile.Name()); err != nil {
+			logger.Debugw("Failed to remove temporary file", "path", idFile.Name(), "err", err)
+		}
+		return fmt.Errorf("failed to write ID file: %w", err)
+	}
+	if err = os.Rename(idFile.Name(), im.path(blobID)); err != nil {
+		return fmt.Errorf("failed to move ID file to store: %w", err)
+	}
+
+	return nil
+}
+
+// Maps blob ID to Singularity ID. Returns blob.ErrBlobNotFound if no mapping
+// exists.
+func (im *idMap) get(blobID blob.ID) (int64, error) {
+	fileIDString, err := os.ReadFile(im.path(blobID))
+	if err != nil {
+		if errors.Is(err, os.ErrNotExist) {
+			return 0, blob.ErrBlobNotFound
+		}
+
+		return 0, fmt.Errorf("could not read ID file: %w", err)
+	}
+
+	fileID, err := strconv.ParseUint(string(fileIDString), 10, 64)
+	if err != nil {
+		return 0, fmt.Errorf("could not parse Singularity file ID '%s' of ID file for blob '%s': %w", fileIDString, blobID, err)
+	}
+
+	return int64(fileID), nil
+}
+
+// TODO: currently commented to silence unused warning
+// // Removes blob ID to Singularity ID mapping. If no ID file existed,
+// // blob.ErrBlobNotFound will be returned.
+// func (im *idMap) remove(blobID blob.ID) error {
+// 	if err := os.Remove(im.path(blobID)); err != nil {
+// 		if errors.Is(err, os.ErrNotExist) {
+// 			return blob.ErrBlobNotFound
+// 		}
+
+// 		return fmt.Errorf("could not read ID file: %w", err)
+// 	}
+
+// 	return nil
+// }