go/oasis-node/cmd/storage: Add create checkpoint cmd

Author: martintomazic

go/oasis-node/cmd/storage/checkpoint.go

@@ -0,0 +1,324 @@
+package storage
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"os"
+	"path/filepath"
+
+	cmtCfg "github.com/cometbft/cometbft/config"
+	cmtNode "github.com/cometbft/cometbft/node"
+	"github.com/spf13/cobra"
+
+	"github.com/oasisprotocol/oasis-core/go/common"
+	"github.com/oasisprotocol/oasis-core/go/config"
+	cmtCommon "github.com/oasisprotocol/oasis-core/go/consensus/cometbft/common"
+	cmtDB "github.com/oasisprotocol/oasis-core/go/consensus/cometbft/db"
+	cmdCommon "github.com/oasisprotocol/oasis-core/go/oasis-node/cmd/common"
+	"github.com/oasisprotocol/oasis-core/go/storage/mkvs/checkpoint"
+	"github.com/oasisprotocol/oasis-core/go/storage/mkvs/db/api"
+	"github.com/oasisprotocol/oasis-core/go/storage/mkvs/node"
+)
+
+const (
+	consensusSubdir = "consensus"
+	runtimesSubdir  = "runtimes"
+)
+
+func newCheckpointCmd() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "checkpoint",
+		Short: "create and restore storage checkpoints",
+		PersistentPreRunE: func(_ *cobra.Command, args []string) error {
+			if err := cmdCommon.Init(); err != nil {
+				cmdCommon.EarlyLogAndExit(err)
+			}
+			running, err := cmdCommon.IsNodeRunning()
+			if err != nil {
+				return fmt.Errorf("failed to ensure the node is not running: %w", err)
+			}
+			if running {
+				return fmt.Errorf("checkpoint operations can only be done when the node is not running")
+			}
+			return nil
+		},
+	}
+
+	cmd.AddCommand(newCreateCmd())
+	cmd.AddCommand(newImportCmd())
+
+	return cmd
+}
+
+func newCreateCmd() *cobra.Command {
+	var (
+		height    uint64
+		runtimeID string
+		round     uint64
+		outDir    string
+	)
+
+	cmd := &cobra.Command{
+		Use:   "create",
+		Args:  cobra.NoArgs,
+		Short: "create storage checkpoints",
+		RunE: func(cmd *cobra.Command, args []string) error {
+			createConsensusCp := func() error {
+				ndb, close, err := openConsensusNodeDB(cmdCommon.DataDir())
+				if err != nil {
+					return fmt.Errorf("failed to open consensus state DB: %w", err)
+				}
+				defer close()
+
+				consensusOutDir := filepath.Join(outDir, consensusSubdir)
+				return createConsensusCheckpoint(cmd.Context(), ndb, height, consensusOutDir)
+			}
+
+			createRuntimeCps := func() error {
+				var ns common.Namespace
+				if err := ns.UnmarshalHex(runtimeID); err != nil {
+					return fmt.Errorf("malformed source runtime ID: %q: %w", ns, err)
+				}
+
+				ndb, err := openRuntimeStateDB(cmdCommon.DataDir(), ns)
+				if err != nil {
+					return fmt.Errorf("failed to open runtime state DB: %w", err)
+				}
+				defer ndb.Close()
+
+				rtOutDir := filepath.Join(outDir, runtimesSubdir, ns.Hex())
+				return createRuntimeCheckpoints(cmd.Context(), ndb, round, rtOutDir)
+
+			}
+
+			if height != 0 { // TODO handle zero value vs not set correctly.
+				if err := createConsensusCp(); err != nil {
+					return fmt.Errorf("failed to create consensus checkpoint (height: %d): %w", height, err)
+				}
+			}
+
+			if runtimeID != "" {
+				if err := createRuntimeCps(); err != nil {
+					return fmt.Errorf("failed to create checkpoints (runtime: %s, round: %d): %w", runtimeID, round, err)
+				}
+			}
+
+			return nil
+		},
+	}
+
+	cmd.Flags().Uint64Var(&height, "height", 0, "consensus height")
+	cmd.Flags().StringVar(&runtimeID, "runtime", "", "hex encoded runtime ID")
+	cmd.Flags().Uint64Var(&round, "round", 0, "round for which checkpoints will be created")
+	cmd.Flags().StringVar(&outDir, "output-dir", "", "output directory")
+	cmd.MarkFlagRequired("output-dir")
+
+	return cmd
+}
+
+func newImportCmd() *cobra.Command {
+	var inputDir string
+
+	cmd := &cobra.Command{
+		Use:   "import",
+		Args:  cobra.NoArgs,
+		Short: "import storage checkpoints",
+		RunE: func(cmd *cobra.Command, args []string) error {
+
+			// Consensus checkpoint:
+			consensusInputDir := filepath.Join(inputDir, consensusSubdir)
+			entries, err := os.ReadDir(consensusInputDir)
+			switch {
+			case err == nil:
+				if err := restoreConsensusCp(cmd.Context(), cmdCommon.DataDir(), consensusInputDir); err != nil {
+					return fmt.Errorf("failed to import consensus checkpoint: %w", err)
+				}
+			case os.IsNotExist(err):
+				// No consensus checkpoints to restore
+			default:
+				return fmt.Errorf("failed to stat: %w", err)
+			}
+
+			// Runtime checkpoints:
+			runtimeInputDir := filepath.Join(inputDir, runtimesSubdir)
+			entries, err = os.ReadDir(runtimeInputDir)
+			switch {
+			case err == nil:
+				for _, ns := range entries {
+					if err := restoreRuntimeCps(cmd.Context(), runtimeInputDir, ns.Name()); err != nil {
+						return fmt.Errorf("failed to import checkpoints (runtime: %s): %w", ns, err)
+					}
+				}
+			case os.IsNotExist(err):
+				// No runtime checkpoints to restore
+			default:
+				return fmt.Errorf("failed to stat: %w", err)
+			}
+
+			return nil
+		},
+	}
+
+	cmd.Flags().StringVar(&inputDir, "input-dir", "", "directory with checkpoints to import")
+	cmd.MarkFlagRequired("input-dir")
+
+	return cmd
+}
+
+func createConsensusCheckpoint(ctx context.Context, ndb api.NodeDB, height uint64, outputDir string) error {
+	roots, err := ndb.GetRootsForVersion(height)
+	if err != nil {
+		return fmt.Errorf("failed to get roots for height %d: %w", height, err)
+	}
+	if len(roots) == 0 {
+		return fmt.Errorf("empty roots")
+	}
+	return createCheckpoints(ctx, ndb, roots, outputDir)
+}
+
+func createRuntimeCheckpoints(ctx context.Context, ndb api.NodeDB, round uint64, outputDir string) error {
+	roots, err := ndb.GetRootsForVersion(round)
+	if err != nil {
+		return fmt.Errorf("failed to get roots for round %d: %w", round, err)
+	}
+	if lenRoots := len(roots); lenRoots != 2 {
+		return fmt.Errorf("unexpected number of roots: got %d, want %d", lenRoots, 2)
+	}
+	return createCheckpoints(ctx, ndb, roots, outputDir)
+}
+
+func createCheckpoints(ctx context.Context, ndb api.NodeDB, roots []node.Root, outputDir string) error {
+	creator, err := checkpoint.NewFileCreator(outputDir, ndb)
+	if err != nil {
+		return fmt.Errorf("failed to create checkpoint file creator: %w", err)
+	}
+
+	for _, root := range roots {
+		_, err := creator.CreateCheckpoint(ctx, root, 8*1024*1024, 8) // 8MiB chunks
+		if err != nil {
+			return fmt.Errorf("failed to create checkpoint (rootType: %s): %w", root.Type, err)
+		}
+	}
+
+	return nil
+}
+
+func restoreConsensusCp(ctx context.Context, dataDir, inputDir string) error {
+	ndb, close, err := openConsensusNodeDB(cmdCommon.DataDir())
+	if err != nil {
+		return fmt.Errorf("failed to open consensus state DB: %w", err)
+	}
+	defer close()
+
+	if _, ok := ndb.GetLatestVersion(); ok {
+		return fmt.Errorf("state db not empty")
+	}
+
+	provider, err := checkpoint.NewFileCreator(inputDir, nil) // ndb = nil since we will only restore checkpoints.
+	if err != nil {
+		return fmt.Errorf("failed to create checkpoint file creator: %w", err)
+	}
+
+	cps, err := provider.GetCheckpoints(ctx, &checkpoint.GetCheckpointsRequest{Version: 1})
+	if err != nil {
+		return fmt.Errorf("failed to read checkpoints: %w", err)
+	}
+
+	if lenCps := len(cps); lenCps != 1 {
+		return fmt.Errorf("unexpected number of checkpoints: got %d, want 1", lenCps)
+	}
+
+	if err := restoreCheckpoints(ctx, provider, ndb, cps); err != nil {
+		return fmt.Errorf("failed to restore checkpoint: %w", err)
+	}
+
+	// TODO: This is just to show how CometBFT is meant to synchronize block and state stores
+	// when the state sync is done offline (aka import checkpoint). Obviously we use our custom
+	// genesis/doc provider and light clients. In theory this shows we could make bootstrap of
+	// oasis node from the checkpoints completely trustless!
+	//
+	// In practice I plan to write our own version of `BootstrapUntrustedState`, where the node
+	// creating checkpoints, so also dump bootstrap metadata.
+	cmtConfig := cmtCfg.DefaultConfig()
+	cmtConfig.SetRoot(filepath.Join(cmdCommon.DataDir(), cmtCommon.StateDir))
+	dbProvider, err := cmtDB.Provider()
+	if err != nil {
+		return fmt.Errorf("failed to obtain db provider: %w", err)
+	}
+	cmtConfig.Genesis = config.GlobalConfig.Genesis.File
+	cmtConfig.StateSync.RPCServers = config.GlobalConfig.Consensus.Providers
+	cmtConfig.StateSync.TrustPeriod = config.GlobalConfig.Consensus.LightClient.Trust.Period
+	cmtConfig.StateSync.TrustHeight = int64(config.GlobalConfig.Consensus.LightClient.Trust.Height)
+	cmtConfig.StateSync.TrustHash = config.GlobalConfig.Consensus.LightClient.Trust.Hash
+
+	root := cps[0].Root
+	if err := cmtNode.BootstrapState(ctx, cmtConfig, dbProvider, root.Version, root.Hash[:]); err != nil {
+		return fmt.Errorf("failed to bootstrap cometbft dbs from the state DB: %w", err)
+	}
+
+	return nil
+}
+
+func restoreRuntimeCps(ctx context.Context, inputDir, namespace string) error {
+	var ns common.Namespace
+	if err := ns.UnmarshalHex(namespace); err != nil {
+		return fmt.Errorf("malformed source runtime ID: %q: %w", ns, err)
+	}
+
+	ndb, err := openRuntimeStateDB(cmdCommon.DataDir(), ns)
+	if err != nil {
+		return err
+	}
+	defer ndb.Close()
+
+	if _, ok := ndb.GetLatestVersion(); ok {
+		return fmt.Errorf("state db not empty")
+	}
+
+	cpDir := filepath.Join(inputDir, namespace)
+	provider, err := checkpoint.NewFileCreator(cpDir, nil) // ndb = nil since we will only restore checkpoints.
+	if err != nil {
+		return fmt.Errorf("failed to create checkpoint file creator: %w", err)
+	}
+
+	cps, err := provider.GetCheckpoints(ctx, &checkpoint.GetCheckpointsRequest{Version: 1, Namespace: ns})
+	if err != nil {
+		return fmt.Errorf("failed to read checkpoints: %w", err)
+	}
+	if err := restoreCheckpoints(ctx, provider, ndb, cps); err != nil {
+		return fmt.Errorf("failed to restore checkpoint: %w", err)
+	}
+
+	return nil
+}
+
+func restoreCheckpoints(ctx context.Context, provider checkpoint.ChunkProvider, ndb api.NodeDB, cps []*checkpoint.Metadata) error {
+	var roots []node.Root
+	for _, cp := range cps {
+		if err := ndb.StartMultipartInsert(cp.Root.Version); err != nil {
+			return fmt.Errorf("failed to start multipart insert: %w", err)
+		}
+		defer ndb.AbortMultipartInsert()
+
+		for idx, _ := range cp.Chunks {
+			chunk, err := cp.GetChunkMetadata(uint64(idx))
+			if err != nil {
+				return fmt.Errorf("failed to get chunk metadata: %w", err)
+			}
+			var buf bytes.Buffer
+			if err := provider.GetCheckpointChunk(ctx, chunk, &buf); err != nil {
+				return fmt.Errorf("failed to read checkpoint chunk (idx: %d): %w", idx, err)
+			}
+			if err := checkpoint.RestoreChunk(ctx, ndb, chunk, &buf); err != nil {
+				return fmt.Errorf("failed to restore chunk: %w", err)
+			}
+		}
+		roots = append(roots, cp.Root)
+	}
+	if err := ndb.Finalize(roots); err != nil {
+		return fmt.Errorf("failed to finalize: %w", err) // TODO maybe you need to force sync?
+	}
+
+	return nil
+}

go/oasis-node/cmd/storage/storage.go

@@ -605,5 +605,6 @@ func Register(parentCmd *cobra.Command) {
 	storageCmd.AddCommand(storageCompactCmd)
 	storageCmd.AddCommand(pruneCmd)
 	storageCmd.AddCommand(newInspectCmd())
+	storageCmd.AddCommand(newCheckpointCmd())
 	parentCmd.AddCommand(storageCmd)
 }

go/storage/mkvs/checkpoint/chunk.go

@@ -262,7 +262,7 @@ func writeChunk(proof *syncer.Proof, w io.Writer) (hash.Hash, error) {
 	return hb.Build(), nil
 }
 
-func restoreChunk(ctx context.Context, ndb db.NodeDB, chunk *ChunkMetadata, r io.Reader) error {
+func RestoreChunk(ctx context.Context, ndb db.NodeDB, chunk *ChunkMetadata, r io.Reader) error {
 	hb := hash.NewBuilder()
 	tr := io.TeeReader(r, hb)
 	sr := snappy.NewReader(tr)

go/storage/mkvs/checkpoint/restorer.go

@@ -83,7 +83,7 @@ func (rs *restorer) RestoreChunk(ctx context.Context, idx uint64, r io.Reader) (
 		return false, err
 	}
 
-	err = restoreChunk(ctx, rs.ndb, chunk, r)
+	err = RestoreChunk(ctx, rs.ndb, chunk, r)
 	switch {
 	case err == nil:
 	case errors.Is(err, ErrChunkProofVerificationFailed):