From c19f25368043129ca5ad3a40ff58f1df1e1c6434 Mon Sep 17 00:00:00 2001
From: Nate <me@n8m.us>
Date: Fri, 6 Dec 2024 22:21:52 -0800
Subject: [PATCH] add volume manager prune test

---
 host/storage/storage_test.go | 85 ++++++++++++++++++++++++++++++++++++
 persist/sqlite/sectors.go    |  2 +-
 2 files changed, 86 insertions(+), 1 deletion(-)

diff --git a/host/storage/storage_test.go b/host/storage/storage_test.go
index c83db1fd..63cfe469 100644
--- a/host/storage/storage_test.go
+++ b/host/storage/storage_test.go
@@ -1138,6 +1138,91 @@ func TestSectorCache(t *testing.T) {
 	}
 }
 
+func TestStoragePrune(t *testing.T) {
+	const sectors = 10
+	dir := t.TempDir()
+
+	// create the database
+	log := zaptest.NewLogger(t)
+	db, err := sqlite.OpenDatabase(filepath.Join(dir, "hostd.db"), log.Named("sqlite"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer db.Close()
+
+	// initialize the storage manager
+	vm, err := storage.NewVolumeManager(db, storage.WithLogger(log.Named("volumes")), storage.WithCacheSize(0), storage.WithPruneInterval(time.Second))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer vm.Close()
+
+	result := make(chan error, 1)
+	volumeFilePath := filepath.Join(t.TempDir(), "hostdata.dat")
+	vol, err := vm.AddVolume(context.Background(), volumeFilePath, sectors, result)
+	if err != nil {
+		t.Fatal(err)
+	} else if err := <-result; err != nil {
+		t.Fatal(err)
+	}
+
+	assertUsedSectors := func(t *testing.T, used uint64) {
+		t.Helper()
+
+		time.Sleep(time.Second) // prune interval
+		volume, err := vm.Volume(vol.ID)
+		if err != nil {
+			t.Fatal(err)
+		} else if volume.UsedSectors != used {
+			t.Fatalf("expected %v used sectors, got %v", used, volume.UsedSectors)
+		}
+
+		m, err := db.Metrics(time.Now())
+		if err != nil {
+			t.Fatal(err)
+		} else if m.Storage.PhysicalSectors != used {
+			t.Fatalf("expected %v used sectors, got %v", used, m.Storage.PhysicalSectors)
+		}
+	}
+
+	assertUsedSectors(t, 0)
+
+	storeRandomSector := func(t *testing.T, expiration uint64) types.Hash256 {
+		t.Helper()
+
+		var sector [rhp2.SectorSize]byte
+		if _, err := frand.Read(sector[:256]); err != nil {
+			t.Fatal("failed to generate random sector:", err)
+		}
+		root := rhp2.SectorRoot(&sector)
+		if err = vm.StoreSector(root, &sector, expiration); err != nil {
+			t.Fatal("failed to store sector:", err)
+		}
+		return root
+	}
+
+	roots := make([]types.Hash256, 0, sectors)
+	// fill the volume
+	for i := 0; i < cap(roots); i++ {
+		storeRandomSector(t, uint64(i+1))
+	}
+
+	// ensure the sectors are not pruned immediately
+	assertUsedSectors(t, sectors)
+
+	// expire half the sectors
+	if err := db.ExpireTempSectors(5); err != nil {
+		t.Fatal(err)
+	}
+	assertUsedSectors(t, 5)
+
+	// expire the remaining sectors
+	if err := db.ExpireTempSectors(10); err != nil {
+		t.Fatal(err)
+	}
+	assertUsedSectors(t, 0)
+}
+
 func BenchmarkVolumeManagerWrite(b *testing.B) {
 	dir := b.TempDir()
 
diff --git a/persist/sqlite/sectors.go b/persist/sqlite/sectors.go
index 85baa4f3..11ad343f 100644
--- a/persist/sqlite/sectors.go
+++ b/persist/sqlite/sectors.go
@@ -114,7 +114,7 @@ WHERE ss.sector_root=$1`, encode(root)).Scan(&sectorID)
 		} else if err := incrementNumericStat(tx, metricTempSectors, 1, time.Now()); err != nil {
 			return fmt.Errorf("failed to update metric: %w", err)
 		}
-		_, err = tx.Exec(`INSERT INTO temp_storage_sector_roots (sector_id, expiration) VALUES ($1, $2)`, sectorID, expiration)
+		_, err = tx.Exec(`INSERT INTO temp_storage_sector_roots (sector_id, expiration_height) VALUES ($1, $2)`, sectorID, expiration)
 		return err
 	})
 }