turbo-api: Allow apps based on turbo-geth to add custom buckets to the chaindata (#1063)

cmd/tgcustom/main.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"os"
 
+	"github.com/ledgerwatch/turbo-geth/common/dbutils"
 	"github.com/ledgerwatch/turbo-geth/eth/stagedsync"
 	"github.com/ledgerwatch/turbo-geth/eth/stagedsync/stages"
 	"github.com/ledgerwatch/turbo-geth/log"
@@ -19,6 +20,10 @@ var flag = cli.StringFlag{
 	Value: "default-value",
 }
 
+const (
+	customBucketName = "ZZZ_0x0F_CUSTOM_BUCKET"
+)
+
 func main() {
 	app := turbocli.MakeApp(runTurboGeth, append(turbocli.DefaultFlags, flag))
 	if err := app.Run(os.Args); err != nil {
@@ -38,12 +43,16 @@ func syncStages(ctx *cli.Context) stagedsync.StageBuilders {
 					Description: "Custom Stage",
 					ExecFunc: func(s *stagedsync.StageState, _ stagedsync.Unwinder) error {
 						fmt.Println("hello from the custom stage", ctx.String(flag.Name))
+						val, err := world.DB.Get(customBucketName, []byte("test"))
+						fmt.Println("val", string(val), "err", err)
+						world.DB.Put(customBucketName, []byte("test"), []byte(ctx.String(flag.Name))) //nolint:errcheck
 						s.Done()
 						return nil
 					},
 					UnwindFunc: func(u *stagedsync.UnwindState, s *stagedsync.StageState) error {
 						fmt.Println("hello from the custom stage unwind", ctx.String(flag.Name))
-						return nil
+						world.DB.Delete(customBucketName, []byte("test")) //nolint:errcheck
+						return u.Done(world.DB)
 					},
 				}
 			},
@@ -57,7 +66,12 @@ func runTurboGeth(ctx *cli.Context) {
 		stagedsync.DefaultUnwindOrder(),
 	)
 
-	tg := node.New(ctx, sync, node.Params{})
+	tg := node.New(ctx, sync, node.Params{
+		CustomBuckets: map[string]dbutils.BucketConfigItem{
+			customBucketName: {},
+		},
+	})
+
 	err := tg.Serve()
 
 	if err != nil {

common/dbutils/bucket.go

@@ -268,10 +268,35 @@ var BucketsConfigs = BucketsCfg{
 	//},
 }
 
-func init() {
+func sortBuckets() {
 	sort.SliceStable(Buckets, func(i, j int) bool {
 		return strings.Compare(Buckets[i], Buckets[j]) < 0
 	})
+}
+
+func DefaultBuckets() BucketsCfg {
+	return BucketsConfigs
+}
+
+func UpdateBucketsList(newBucketCfg BucketsCfg) {
+	newBuckets := make([]string, 0)
+	for k, v := range newBucketCfg {
+		if !v.IsDeprecated {
+			newBuckets = append(newBuckets, k)
+		}
+	}
+	Buckets = newBuckets
+	BucketsConfigs = newBucketCfg
+
+	reinit()
+}
+
+func init() {
+	reinit()
+}
+
+func reinit() {
+	sortBuckets()
 
 	for _, name := range Buckets {
 		_, ok := BucketsConfigs[name]

turbo/node/buckets.go

@@ -0,0 +1,24 @@
+package node
+
+import (
+	"fmt"
+
+	"github.com/ledgerwatch/turbo-geth/common/dbutils"
+)
+
+func prepareBuckets(customBuckets dbutils.BucketsCfg) {
+	if len(customBuckets) == 0 {
+		return
+	}
+
+	currentBuckets := dbutils.DefaultBuckets()
+
+	for k, v := range customBuckets {
+		if _, ok := currentBuckets[k]; ok {
+			panic(fmt.Errorf("overriding existing buckets is not supported (bucket key=%s)", k))
+		}
+		currentBuckets[k] = v
+	}
+
+	dbutils.UpdateBucketsList(currentBuckets)
+}

turbo/node/node.go

@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/ledgerwatch/turbo-geth/cmd/utils"
+	"github.com/ledgerwatch/turbo-geth/common/dbutils"
 	"github.com/ledgerwatch/turbo-geth/eth"
 	"github.com/ledgerwatch/turbo-geth/eth/stagedsync"
 	"github.com/ledgerwatch/turbo-geth/log"
@@ -42,11 +43,17 @@ func (tg *TurboGethNode) run() {
 }
 
 type Params struct {
-	GitDate   string
-	GitCommit string
+	GitDate       string
+	GitCommit     string
+	CustomBuckets dbutils.BucketsCfg
 }
 
-func New(ctx *cli.Context, sync *stagedsync.StagedSync, p Params) *TurboGethNode {
+func New(
+	ctx *cli.Context,
+	sync *stagedsync.StagedSync,
+	p Params,
+) *TurboGethNode {
+	prepareBuckets(p.CustomBuckets)
 	prepare(ctx)
 	nodeConfig := makeNodeConfig(ctx, p)
 	node := makeConfigNode(nodeConfig)