feat: add local snapshots management commands (#16067)

Co-authored-by: Marko <marbar3778@yahoo.com>

client/snapshot/cmd.go

@@ -0,0 +1,24 @@
+package snapshot
+
+import (
+	servertypes "github.com/cosmos/cosmos-sdk/server/types"
+	"github.com/spf13/cobra"
+)
+
+// Cmd returns the snapshots group command
+func Cmd(appCreator servertypes.AppCreator) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "snapshots",
+		Short: "Manage local snapshots",
+		Long:  "Manage local snapshots",
+	}
+	cmd.AddCommand(
+		ListSnapshotsCmd,
+		RestoreSnapshotCmd(appCreator),
+		ExportSnapshotCmd(appCreator),
+		DumpArchiveCmd(),
+		LoadArchiveCmd(),
+		DeleteSnapshotCmd(),
+	)
+	return cmd
+}

client/snapshot/delete.go

@@ -0,0 +1,35 @@
+package snapshot
+
+import (
+	"strconv"
+
+	"github.com/cosmos/cosmos-sdk/server"
+	"github.com/spf13/cobra"
+)
+
+func DeleteSnapshotCmd() *cobra.Command {
+	return &cobra.Command{
+		Use:   "delete <height> <format>",
+		Short: "Delete a local snapshot",
+		Args:  cobra.ExactArgs(2),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := server.GetServerContextFromCmd(cmd)
+
+			height, err := strconv.ParseUint(args[0], 10, 64)
+			if err != nil {
+				return err
+			}
+			format, err := strconv.ParseUint(args[1], 10, 32)
+			if err != nil {
+				return err
+			}
+
+			snapshotStore, err := server.GetSnapshotStore(ctx.Viper)
+			if err != nil {
+				return err
+			}
+
+			return snapshotStore.Delete(height, uint32(format))
+		},
+	}
+}

client/snapshot/dump.go

@@ -0,0 +1,119 @@
+package snapshot
+
+import (
+	"archive/tar"
+	"compress/gzip"
+	"fmt"
+	"io"
+	"os"
+	"strconv"
+
+	"github.com/cosmos/cosmos-sdk/server"
+	"github.com/spf13/cobra"
+)
+
+// DumpArchiveCmd returns a command to dump the snapshot as portable archive format
+func DumpArchiveCmd() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "dump <height> <format>",
+		Short: "Dump the snapshot as portable archive format",
+		Args:  cobra.ExactArgs(2),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := server.GetServerContextFromCmd(cmd)
+			snapshotStore, err := server.GetSnapshotStore(ctx.Viper)
+			if err != nil {
+				return err
+			}
+
+			output, err := cmd.Flags().GetString("output")
+			if err != nil {
+				return err
+			}
+
+			height, err := strconv.ParseUint(args[0], 10, 64)
+			if err != nil {
+				return err
+			}
+			format, err := strconv.ParseUint(args[1], 10, 32)
+			if err != nil {
+				return err
+			}
+
+			if output == "" {
+				output = fmt.Sprintf("%d-%d.tar.gz", height, format)
+			}
+
+			snapshot, err := snapshotStore.Get(height, uint32(format))
+			if err != nil {
+				return err
+			}
+
+			bz, err := snapshot.Marshal()
+			if err != nil {
+				return err
+			}
+
+			fp, err := os.Create(output)
+			if err != nil {
+				return err
+			}
+			defer fp.Close()
+
+			// since the chunk files are already compressed, we just use fastest compression here
+			gzipWriter, err := gzip.NewWriterLevel(fp, gzip.BestSpeed)
+			if err != nil {
+				return err
+			}
+			tarWriter := tar.NewWriter(gzipWriter)
+			if err := tarWriter.WriteHeader(&tar.Header{
+				Name: SnapshotFileName,
+				Mode: 0o644,
+				Size: int64(len(bz)),
+			}); err != nil {
+				return fmt.Errorf("failed to write snapshot header to tar: %w", err)
+			}
+			if _, err := tarWriter.Write(bz); err != nil {
+				return fmt.Errorf("failed to write snapshot to tar: %w", err)
+			}
+
+			for i := uint32(0); i < snapshot.Chunks; i++ {
+				path := snapshotStore.PathChunk(height, uint32(format), i)
+				file, err := os.Open(path)
+				if err != nil {
+					return fmt.Errorf("failed to open chunk file %s: %w", path, err)
+				}
+
+				st, err := file.Stat()
+				if err != nil {
+					return fmt.Errorf("failed to stat chunk file %s: %w", path, err)
+				}
+
+				if err := tarWriter.WriteHeader(&tar.Header{
+					Name: strconv.FormatUint(uint64(i), 10),
+					Mode: 0o644,
+					Size: st.Size(),
+				}); err != nil {
+					return fmt.Errorf("failed to write chunk header to tar: %w", err)
+				}
+
+				if _, err := io.Copy(tarWriter, file); err != nil {
+					return fmt.Errorf("failed to write chunk to tar: %w", err)
+				}
+			}
+
+			if err := tarWriter.Close(); err != nil {
+				return fmt.Errorf("failed to close tar writer: %w", err)
+			}
+
+			if err := gzipWriter.Close(); err != nil {
+				return fmt.Errorf("failed to close gzip writer: %w", err)
+			}
+
+			return fp.Close()
+		},
+	}
+
+	cmd.Flags().StringP("output", "o", "", "output file")
+
+	return cmd
+}

client/snapshot/export.go

@@ -0,0 +1,54 @@
+package snapshot
+
+import (
+	"fmt"
+
+	"cosmossdk.io/log"
+	"github.com/cosmos/cosmos-sdk/server"
+	servertypes "github.com/cosmos/cosmos-sdk/server/types"
+	"github.com/spf13/cobra"
+)
+
+// ExportSnapshotCmd returns a command to take a snapshot of the application state
+func ExportSnapshotCmd(appCreator servertypes.AppCreator) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "export",
+		Short: "Export app state to snapshot store",
+		Args:  cobra.NoArgs,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := server.GetServerContextFromCmd(cmd)
+
+			height, err := cmd.Flags().GetInt64("height")
+			if err != nil {
+				return err
+			}
+
+			home := ctx.Config.RootDir
+			db, err := openDB(home, server.GetAppDBBackend(ctx.Viper))
+			if err != nil {
+				return err
+			}
+			logger := log.NewLogger(cmd.OutOrStdout())
+			app := appCreator(logger, db, nil, ctx.Viper)
+
+			if height == 0 {
+				height = app.CommitMultiStore().LastCommitID().Version
+			}
+
+			fmt.Printf("Exporting snapshot for height %d\n", height)
+
+			sm := app.SnapshotManager()
+			snapshot, err := sm.Create(uint64(height))
+			if err != nil {
+				return err
+			}
+
+			fmt.Printf("Snapshot created at height %d, format %d, chunks %d\n", snapshot.Height, snapshot.Format, snapshot.Chunks)
+			return nil
+		},
+	}
+
+	cmd.Flags().Int64("height", 0, "Height to export, default to latest state height")
+
+	return cmd
+}

client/snapshot/list.go

@@ -0,0 +1,30 @@
+package snapshot
+
+import (
+	"fmt"
+
+	"github.com/cosmos/cosmos-sdk/server"
+	"github.com/spf13/cobra"
+)
+
+// ListSnapshotsCmd returns the command to list local snapshots
+var ListSnapshotsCmd = &cobra.Command{
+	Use:   "list",
+	Short: "List local snapshots",
+	RunE: func(cmd *cobra.Command, args []string) error {
+		ctx := server.GetServerContextFromCmd(cmd)
+		snapshotStore, err := server.GetSnapshotStore(ctx.Viper)
+		if err != nil {
+			return err
+		}
+		snapshots, err := snapshotStore.List()
+		if err != nil {
+			return fmt.Errorf("failed to list snapshots: %w", err)
+		}
+		for _, snapshot := range snapshots {
+			fmt.Println("height:", snapshot.Height, "format:", snapshot.Format, "chunks:", snapshot.Chunks)
+		}
+
+		return nil
+	},
+}

client/snapshot/load.go

@@ -0,0 +1,113 @@
+package snapshot
+
+import (
+	"archive/tar"
+	"bytes"
+	"compress/gzip"
+	"fmt"
+	"io"
+	"os"
+	"reflect"
+	"strconv"
+
+	"github.com/cosmos/cosmos-sdk/server"
+	"github.com/spf13/cobra"
+
+	snapshottypes "cosmossdk.io/store/snapshots/types"
+)
+
+const SnapshotFileName = "_snapshot"
+
+// LoadArchiveCmd load a portable archive format snapshot into snapshot store
+func LoadArchiveCmd() *cobra.Command {
+	return &cobra.Command{
+		Use:   "load <archive-file>",
+		Short: "Load a snapshot archive file into snapshot store",
+		Args:  cobra.ExactArgs(1),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := server.GetServerContextFromCmd(cmd)
+			snapshotStore, err := server.GetSnapshotStore(ctx.Viper)
+			if err != nil {
+				return err
+			}
+
+			path := args[0]
+			fp, err := os.Open(path)
+			if err != nil {
+				return fmt.Errorf("failed to open archive file: %w", err)
+			}
+			reader, err := gzip.NewReader(fp)
+			if err != nil {
+				return fmt.Errorf("failed to create gzip reader: %w", err)
+			}
+
+			var snapshot snapshottypes.Snapshot
+			tr := tar.NewReader(reader)
+			if err != nil {
+				return fmt.Errorf("failed to create tar reader: %w", err)
+			}
+
+			hdr, err := tr.Next()
+			if err != nil {
+				return fmt.Errorf("failed to read snapshot file header: %w", err)
+			}
+			if hdr.Name != SnapshotFileName {
+				return fmt.Errorf("invalid archive, expect file: snapshot, got: %s", hdr.Name)
+			}
+			bz, err := io.ReadAll(tr)
+			if err != nil {
+				return fmt.Errorf("failed to read snapshot file: %w", err)
+			}
+			if err := snapshot.Unmarshal(bz); err != nil {
+				return fmt.Errorf("failed to unmarshal snapshot: %w", err)
+			}
+
+			// make sure the channel is unbuffered, because the tar reader can't do concurrency
+			chunks := make(chan io.ReadCloser)
+			quitChan := make(chan *snapshottypes.Snapshot)
+			go func() {
+				defer close(quitChan)
+
+				savedSnapshot, err := snapshotStore.Save(snapshot.Height, snapshot.Format, chunks)
+				if err != nil {
+					fmt.Println("failed to save snapshot", err)
+					return
+				}
+				quitChan <- savedSnapshot
+			}()
+
+			for i := uint32(0); i < snapshot.Chunks; i++ {
+				hdr, err = tr.Next()
+				if err != nil {
+					if err == io.EOF {
+						break
+					}
+					return err
+				}
+
+				if hdr.Name != strconv.FormatInt(int64(i), 10) {
+					return fmt.Errorf("invalid archive, expect file: %d, got: %s", i, hdr.Name)
+				}
+
+				bz, err := io.ReadAll(tr)
+				if err != nil {
+					return fmt.Errorf("failed to read chunk file: %w", err)
+				}
+				chunks <- io.NopCloser(bytes.NewReader(bz))
+			}
+			close(chunks)
+
+			savedSnapshot := <-quitChan
+			if savedSnapshot == nil {
+				return fmt.Errorf("failed to save snapshot")
+			}
+
+			if !reflect.DeepEqual(&snapshot, savedSnapshot) {
+				_ = snapshotStore.Delete(snapshot.Height, snapshot.Format)
+				return fmt.Errorf("invalid archive, the saved snapshot is not equal to the original one")
+			}
+
+			return nil
+		},
+	}
+}

client/snapshot/restore.go

@@ -0,0 +1,52 @@
+package snapshot
+
+import (
+	"path/filepath"
+	"strconv"
+
+	"cosmossdk.io/log"
+	"github.com/spf13/cobra"
+
+	dbm "github.com/cosmos/cosmos-db"
+	"github.com/cosmos/cosmos-sdk/server"
+	servertypes "github.com/cosmos/cosmos-sdk/server/types"
+)
+
+// RestoreSnapshotCmd returns a command to restore a snapshot
+func RestoreSnapshotCmd(appCreator servertypes.AppCreator) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "restore <height> <format>",
+		Short: "Restore app state from local snapshot",
+		Long:  "Restore app state from local snapshot",
+		Args:  cobra.ExactArgs(2),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := server.GetServerContextFromCmd(cmd)
+
+			height, err := strconv.ParseUint(args[0], 10, 64)
+			if err != nil {
+				return err
+			}
+			format, err := strconv.ParseUint(args[1], 10, 32)
+			if err != nil {
+				return err
+			}
+
+			home := ctx.Config.RootDir
+			db, err := openDB(home, server.GetAppDBBackend(ctx.Viper))
+			if err != nil {
+				return err
+			}
+			logger := log.NewLogger(cmd.OutOrStdout())
+			app := appCreator(logger, db, nil, ctx.Viper)
+
+			sm := app.SnapshotManager()
+			return sm.RestoreLocalSnapshot(height, uint32(format))
+		},
+	}
+	return cmd
+}
+
+func openDB(rootDir string, backendType dbm.BackendType) (dbm.DB, error) {
+	dataDir := filepath.Join(rootDir, "data")
+	return dbm.NewDB("application", backendType, dataDir)
+}