From 87801b585508455835bc5e57b98e2e6a32a3a955 Mon Sep 17 00:00:00 2001
From: Steven Allen <steven@stebalien.com>
Date: Fri, 27 Oct 2023 11:33:18 -0700
Subject: [PATCH] feat: shed: fix blockstore prune (#11197)

1. Implement a pass-through for garbage collection on the idstore.
2. Fix the `lotus-shed state-prune` command.

NOTE: the new performance of running a full prune will be significantly
less as this version doesn't batch. However, it should now actually
_work_ and most users will be using the splitstore anyways.
---
 blockstore/idstore.go     | 14 ++++++++++++++
 cmd/lotus-shed/pruning.go | 36 ++++++++++++++----------------------
 2 files changed, 28 insertions(+), 22 deletions(-)

diff --git a/blockstore/idstore.go b/blockstore/idstore.go
index fb575dca7..a10aee575 100644
--- a/blockstore/idstore.go
+++ b/blockstore/idstore.go
@@ -183,3 +183,17 @@ func (b *idstore) Close() error {
 func (b *idstore) Flush(ctx context.Context) error {
 	return b.bs.Flush(ctx)
 }
+
+func (b *idstore) CollectGarbage(ctx context.Context, options ...BlockstoreGCOption) error {
+	if bs, ok := b.bs.(BlockstoreGC); ok {
+		return bs.CollectGarbage(ctx, options...)
+	}
+	return xerrors.Errorf("not supported")
+}
+
+func (b *idstore) GCOnce(ctx context.Context, options ...BlockstoreGCOption) error {
+	if bs, ok := b.bs.(BlockstoreGCOnce); ok {
+		return bs.GCOnce(ctx, options...)
+	}
+	return xerrors.Errorf("not supported")
+}
diff --git a/cmd/lotus-shed/pruning.go b/cmd/lotus-shed/pruning.go
index 275f3bc0a..c0bd453b1 100644
--- a/cmd/lotus-shed/pruning.go
+++ b/cmd/lotus-shed/pruning.go
@@ -12,7 +12,7 @@ import (
 
 	"github.com/filecoin-project/go-state-types/abi"
 
-	badgerbs "github.com/filecoin-project/lotus/blockstore/badger"
+	"github.com/filecoin-project/lotus/blockstore"
 	"github.com/filecoin-project/lotus/chain/consensus/filcns"
 	"github.com/filecoin-project/lotus/chain/store"
 	"github.com/filecoin-project/lotus/node/repo"
@@ -144,13 +144,6 @@ var stateTreePruneCmd = &cli.Command{
 			}
 		}()
 
-		// After migrating to native blockstores, this has been made
-		// database-specific.
-		badgbs, ok := bs.(*badgerbs.Blockstore)
-		if !ok {
-			return fmt.Errorf("only badger blockstores are supported")
-		}
-
 		mds, err := lkrepo.Datastore(context.Background(), "/metadata")
 		if err != nil {
 			return err
@@ -160,8 +153,12 @@ var stateTreePruneCmd = &cli.Command{
 		const DiscardRatio = 0.2
 		if cctx.Bool("only-ds-gc") {
 			fmt.Println("running datastore gc....")
+			gbs, ok := bs.(blockstore.BlockstoreGCOnce)
+			if !ok {
+				return xerrors.Errorf("blockstore %T does not support GC", bs)
+			}
 			for i := 0; i < cctx.Int("gc-count"); i++ {
-				if err := badgbs.DB().RunValueLogGC(DiscardRatio); err != nil {
+				if err := gbs.GCOnce(ctx, blockstore.WithThreshold(DiscardRatio)); err != nil {
 					return xerrors.Errorf("datastore GC failed: %w", err)
 				}
 			}
@@ -208,13 +205,6 @@ var stateTreePruneCmd = &cli.Command{
 			return nil
 		}
 
-		b := badgbs.DB().NewWriteBatch()
-		defer b.Cancel()
-
-		markForRemoval := func(c cid.Cid) error {
-			return b.Delete(badgbs.StorageKey(nil, c))
-		}
-
 		keys, err := bs.AllKeysChan(context.Background())
 		if err != nil {
 			return xerrors.Errorf("failed to query blockstore: %w", err)
@@ -225,12 +215,12 @@ var stateTreePruneCmd = &cli.Command{
 		var deleteCount int
 		var goodHits int
 		for k := range keys {
-			if goodSet.HasRaw(k.Bytes()) {
+			if goodSet.Has(k) {
 				goodHits++
 				continue
 			}
 
-			if err := markForRemoval(k); err != nil {
+			if err := bs.DeleteBlock(ctx, k); err != nil {
 				return fmt.Errorf("failed to remove cid %s: %w", k, err)
 			}
 
@@ -243,13 +233,15 @@ var stateTreePruneCmd = &cli.Command{
 			}
 		}
 
-		if err := b.Flush(); err != nil {
-			return xerrors.Errorf("failed to flush final batch delete: %w", err)
+		fmt.Println("running datastore gc....")
+		gbs, ok := bs.(blockstore.BlockstoreGCOnce)
+		if !ok {
+			fmt.Println("gc not supported...")
+			return nil
 		}
 
-		fmt.Println("running datastore gc....")
 		for i := 0; i < cctx.Int("gc-count"); i++ {
-			if err := badgbs.DB().RunValueLogGC(DiscardRatio); err != nil {
+			if err := gbs.GCOnce(ctx, blockstore.WithThreshold(DiscardRatio)); err != nil {
 				return xerrors.Errorf("datastore GC failed: %w", err)
 			}
 		}