Allow for bootstrapping with data from a snapshot

pkg/store/options_bootstrap.go

@@ -35,6 +35,7 @@ const (
 	BootstrapAdminEnvVar                = "BOOTSTRAP_ADMIN"
 	BootstrapVotersEnvVar               = "BOOTSTRAP_VOTERS"
 	BootstrapDefaultNetworkPolicyEnvVar = "BOOTSTRAP_DEFAULT_NETWORK_POLICY"
+	BootstrapRestoreSnapshotEnvVar      = "BOOTSTRAP_RESTORE_SNAPSHOT"
 	ForceBootstrapClusterEnvVar         = "BOOTSTRAP_FORCE"
 )
 
@@ -63,6 +64,8 @@ type BootstrapOptions struct {
 	Voters string `json:"voters,omitempty" yaml:"voters,omitempty" toml:"voters,omitempty"`
 	// DefaultNetworkPolicy is the default network policy to apply to the mesh when bootstraping a new cluster.
 	DefaultNetworkPolicy string `json:"default-network-policy,omitempty" yaml:"default-network-policy,omitempty" toml:"default-network-policy,omitempty"`
+	// RestoreSnapshot is the path to a snapshot to restore from when bootstrapping a new cluster.
+	RestoreSnapshot string `json:"restore-snapshot,omitempty" yaml:"restore-snapshot,omitempty" toml:"restore-snapshot,omitempty"`
 	// Force is the force new bootstrap flag.
 	Force bool `json:"force,omitempty" yaml:"force,omitempty" toml:"force,omitempty"`
 }
@@ -156,6 +159,9 @@ Ports should be in the form of <node-id>=<port>.`)
 	fl.StringVar(&o.DefaultNetworkPolicy, "bootstrap.default-network-policy", util.GetEnvDefault(BootstrapDefaultNetworkPolicyEnvVar, string(NetworkPolicyDeny)),
 		"Default network policy to bootstrap the cluster with.")
 
+	fl.StringVar(&o.RestoreSnapshot, "bootstrap.restore-snapshot", util.GetEnvDefault(BootstrapRestoreSnapshotEnvVar, ""),
+		"Path to a snapshot to restore from when bootstrapping a new cluster.")
+
 	fl.BoolVar(&o.Force, "bootstrap.force", util.GetEnvDefault(ForceBootstrapClusterEnvVar, "false") == "true",
 		"Force bootstrapping a new cluster even if data is present.")
 }

pkg/store/store_bootstrap.go

@@ -23,6 +23,7 @@ import (
 	"math"
 	"net"
 	"net/netip"
+	"os"
 	"strconv"
 	"strings"
 	"time"
@@ -42,14 +43,14 @@ import (
 func (s *store) bootstrap(ctx context.Context) error {
 	version, err := models.GetDBVersion(s.weakData)
 	if err != nil {
-		return fmt.Errorf("get raft schema version: %w", err)
+		return fmt.Errorf("get db schema version: %w", err)
 	}
-	s.log.Info("current raft schema version", slog.Int("version", int(version)))
+	s.log.Info("current db schema version", slog.Int("version", int(version)))
 	if version != 0 {
 		// We have a version, so the cluster is already bootstrapped.
 		s.log.Info("cluster already bootstrapped, migrating schema to latest version")
 		if err = models.MigrateRaftDB(s.weakData); err != nil {
-			return fmt.Errorf("raft db migrate: %w", err)
+			return fmt.Errorf("db migrate: %w", err)
 		}
 		// We rejoin as a voter no matter what
 		s.opts.Mesh.JoinAsVoter = true
@@ -63,12 +64,31 @@ func (s *store) bootstrap(ctx context.Context) error {
 		// Try to rejoin one of the bootstrap servers
 		return s.rejoinBootstrapServer(ctx)
 	}
+	if s.opts.Bootstrap.RestoreSnapshot != "" {
+		s.log.Info("restoring snapshot from file", slog.String("file", s.opts.Bootstrap.RestoreSnapshot))
+		f, err := os.Open(s.opts.Bootstrap.RestoreSnapshot)
+		if err != nil {
+			return fmt.Errorf("open snapshot file: %w", err)
+		}
+		defer f.Close()
+		if err := s.snapshotter.Restore(ctx, f); err != nil {
+			return fmt.Errorf("restore snapshot: %w", err)
+		}
+		s.log.Info("migrating schema to latest version")
+		if err = models.MigrateRaftDB(s.weakData); err != nil {
+			return fmt.Errorf("db migrate: %w", err)
+		}
+		// We're done here, but restore procedure needs to be documented
+		return nil
+	}
 	s.firstBootstrap.Store(true)
+	s.log.Info("migrating schema to latest version")
+	if err = models.MigrateRaftDB(s.weakData); err != nil {
+		return fmt.Errorf("db migrate: %w", err)
+	}
 	if s.opts.Bootstrap.AdvertiseAddress == "" && s.opts.Bootstrap.Servers == "" {
 		s.opts.Bootstrap.AdvertiseAddress = fmt.Sprintf("localhost:%d", s.sl.ListenPort())
 	} else if s.opts.Bootstrap.AdvertiseAddress == "" {
-		// Validate() doesn't allow this on the options
-		// but lets go ahead and support it anyway.
 		s.opts.Bootstrap.AdvertiseAddress = s.opts.Bootstrap.Servers
 	}
 	// There is a chance we are waiting for DNS to resolve.
@@ -140,9 +160,6 @@ func (s *store) bootstrap(ctx context.Context) error {
 		// there were other servers to bootstrap with, then
 		// we might just need to rejoin the cluster.
 		if errors.Is(err, raft.ErrCantBootstrap) {
-			if err = models.MigrateRaftDB(s.weakData); err != nil {
-				return fmt.Errorf("raft db migrate: %w", err)
-			}
 			if s.opts.Bootstrap.Servers != "" {
 				s.log.Info("cluster already bootstrapped, attempting to rejoin as voter")
 				s.opts.Mesh.JoinAsVoter = true
@@ -156,10 +173,6 @@ func (s *store) bootstrap(ctx context.Context) error {
 		}
 		return fmt.Errorf("bootstrap cluster: %w", err)
 	}
-	s.log.Info("migrating raft schema to latest version")
-	if err = models.MigrateRaftDB(s.weakData); err != nil {
-		return fmt.Errorf("raft db migrate: %w", err)
-	}
 	go func() {
 		deadline, ok := ctx.Deadline()
 		var cancel context.CancelFunc

pkg/store/store_open.go

@@ -72,6 +72,7 @@ func (s *store) Open() error {
 	}
 	// If bootstrap and force are set, clear the data directory.
 	if s.opts.Bootstrap.Enabled && s.opts.Bootstrap.Force {
+		log.Warn("force bootstrap enabled, clearing data directory")
 		err = os.RemoveAll(s.opts.Raft.DataDir)
 		if err != nil && !os.IsNotExist(err) {
 			return fmt.Errorf("remove all %q: %w", s.opts.Raft.DataDir, err)