feat(v2): metastore snapshot compression (#4052)

Author: kolesnikovae

pkg/experiment/metastore/fsm/boltdb.go

@@ -19,7 +19,7 @@ const (
 	boltDBCompactedName   = "metastore_compacted.boltdb"
 	boltDBInitialMmapSize = 1 << 30
 
-	boltDBCompactionMaxTxnSize = 64 << 10
+	boltDBCompactionMaxTxnSize = 1 << 20
 )
 
 type boltdb struct {
@@ -127,7 +127,7 @@ func (db *boltdb) restore(snapshot io.Reader) error {
 			}
 		}
 		restored.shutdown()
-		err = os.RemoveAll(restored.path)
+		_ = os.RemoveAll(restored.path)
 	}
 	if err != nil {
 		return fmt.Errorf("failed to restore snapshot: %w", err)

pkg/experiment/metastore/fsm/fsm.go

@@ -3,6 +3,7 @@ package fsm
 import (
 	"context"
 	"encoding/binary"
+	"flag"
 	"fmt"
 	"io"
 	"strconv"
@@ -36,9 +37,22 @@ type StateRestorer interface {
 	Restore(*bbolt.Tx) error
 }
 
+type Config struct {
+	SnapshotCompression string `yaml:"snapshot_compression"`
+	// Where the FSM BoltDB data is located.
+	// Does not have to be a persistent volume.
+	DataDir string `yaml:"data_dir"`
+}
+
+func (cfg *Config) RegisterFlagsWithPrefix(prefix string, f *flag.FlagSet) {
+	f.StringVar(&cfg.SnapshotCompression, prefix+"snapshot-compression", "zstd", "Compression algorithm to use for snapshots. Supported compressions: zstd.")
+	f.StringVar(&cfg.DataDir, prefix+"data-dir", "./data-metastore/data", "Directory to store the data.")
+}
+
 // FSM implements the raft.FSM interface.
 type FSM struct {
 	logger  log.Logger
+	config  Config
 	metrics *metrics
 
 	mu   sync.RWMutex
@@ -54,13 +68,14 @@ type FSM struct {
 
 type handler func(tx *bbolt.Tx, cmd *raft.Log, raw []byte) (proto.Message, error)
 
-func New(logger log.Logger, reg prometheus.Registerer, dir string) (*FSM, error) {
+func New(logger log.Logger, reg prometheus.Registerer, config Config) (*FSM, error) {
 	fsm := FSM{
 		logger:   logger,
+		config:   config,
 		metrics:  newMetrics(reg),
 		handlers: make(map[RaftLogEntryType]handler),
 	}
-	db := newDB(logger, fsm.metrics, dir)
+	db := newDB(logger, fsm.metrics, config.DataDir)
 	if err := db.open(false); err != nil {
 		return nil, err
 	}
@@ -133,29 +148,42 @@ func (fsm *FSM) restore() error {
 // Restore restores the FSM state from a snapshot.
 func (fsm *FSM) Restore(snapshot io.ReadCloser) (err error) {
 	start := time.Now()
-	_ = level.Info(fsm.logger).Log("msg", "restoring snapshot")
+	level.Info(fsm.logger).Log("msg", "restoring snapshot")
 	defer func() {
 		_ = snapshot.Close()
 		fsm.db.metrics.fsmRestoreSnapshotDuration.Observe(time.Since(start).Seconds())
 	}()
+
+	level.Info(fsm.logger).Log("msg", "restoring snapshot")
+	var r *snapshotReader
+	if r, err = newSnapshotReader(snapshot); err != nil {
+		level.Error(fsm.logger).Log("msg", "failed to create snapshot reader", "err", err)
+		return err
+	}
+	// The wrapper never returns errors on Close.
+	defer r.Close()
+
 	// Block all new transactions until we restore the snapshot.
 	// TODO(kolesnikovae): set not-serving service status to not
 	//  block incoming requests.
 	fsm.mu.Lock()
 	defer fsm.mu.Unlock()
 	fsm.txns.Wait()
-	if err = fsm.db.restore(snapshot); err != nil {
-		return fmt.Errorf("failed to restore database from snapshot: %w", err)
+	if err = fsm.db.restore(r); err != nil {
+		level.Error(fsm.logger).Log("msg", "failed to restore database from snapshot", "err", err)
+		return err
 	}
 	// First we need to initialize the state: each restorer is called
 	// synchronously and has exclusive access to the database.
 	if err = fsm.init(); err != nil {
-		return fmt.Errorf("failed to init state at restore: %w", err)
+		level.Error(fsm.logger).Log("msg", "failed to init state at restore", "err", err)
+		return err
 	}
 	// Then we restore the state: each restorer is given its own
 	// transaction and run concurrently with others.
 	if err = fsm.restore(); err != nil {
-		return fmt.Errorf("failed to restore state from snapshot: %w", err)
+		level.Error(fsm.logger).Log("msg", "failed to restore state from snapshot", "err", err)
+		return err
 	}
 	return nil
 }
@@ -267,7 +295,11 @@ func (fsm *FSM) Read(fn func(*bbolt.Tx)) error {
 func (fsm *FSM) Snapshot() (raft.FSMSnapshot, error) {
 	// Snapshot should only capture a pointer to the state, and any
 	// expensive IO should happen as part of FSMSnapshot.Persist.
-	s := snapshot{logger: fsm.logger, metrics: fsm.metrics}
+	s := snapshotWriter{
+		logger:      fsm.logger,
+		metrics:     fsm.metrics,
+		compression: fsm.config.SnapshotCompression,
+	}
 	tx, err := fsm.db.boltdb.Begin(false)
 	if err != nil {
 		return nil, fmt.Errorf("failed to open a transaction for snapshot: %w", err)

pkg/experiment/metastore/fsm/snapshot.go

@@ -1,23 +1,29 @@
 package fsm
 
 import (
+	"bufio"
+	"bytes"
 	"context"
+	"io"
 	"runtime/pprof"
 	"time"
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
 	"github.com/hashicorp/raft"
+	"github.com/klauspost/compress/zstd"
 	"go.etcd.io/bbolt"
 )
 
-type snapshot struct {
-	logger  log.Logger
-	tx      *bbolt.Tx
-	metrics *metrics
+type snapshotWriter struct {
+	logger      log.Logger
+	tx          *bbolt.Tx
+	metrics     *metrics
+	compression string
+	written     int64
 }
 
-func (s *snapshot) Persist(sink raft.SnapshotSink) (err error) {
+func (s *snapshotWriter) Persist(sink raft.SnapshotSink) (err error) {
 	ctx := context.Background()
 	pprof.SetGoroutineLabels(pprof.WithLabels(ctx, pprof.Labels("metastore_op", "persist")))
 	defer pprof.SetGoroutineLabels(ctx)
@@ -39,19 +45,75 @@ func (s *snapshot) Persist(sink raft.SnapshotSink) (err error) {
 		}
 	}()
 
-	level.Info(s.logger).Log("msg", "persisting snapshot")
-	var n int64
-	n, err = s.tx.WriteTo(sink)
-	s.metrics.boltDBPersistSnapshotSize.Set(float64(n))
-	if err != nil {
+	var dst io.Writer
+	// Needed to measure the actual snapshot size written to the sink.
+	w := &writeCounter{Writer: sink}
+	if s.compression == "zstd" {
+		// We do not want to bog down the CPU with compression: we are fine with
+		// slowing down the snapshotting process, as it potentially may affect the
+		// WAL writes if they are located on the same disk.
+		var enc *zstd.Encoder
+		if enc, err = zstd.NewWriter(w, zstd.WithEncoderConcurrency(1)); err != nil {
+			level.Error(s.logger).Log("msg", "failed to create zstd encoder", "err", err)
+			return err
+		}
+		dst = enc
+		defer func() {
+			if err = enc.Close(); err != nil {
+				level.Error(s.logger).Log("msg", "zstd compression failed", "err", err)
+			}
+		}()
+	}
+
+	level.Info(s.logger).Log("msg", "persisting snapshot", "compression", s.compression)
+	if _, err = s.tx.WriteTo(dst); err != nil {
 		level.Error(s.logger).Log("msg", "failed to write snapshot", "err", err)
+		return err
 	}
-	return err
+
+	s.metrics.boltDBPersistSnapshotSize.Set(float64(w.written))
+	return nil
 }
 
-func (s *snapshot) Release() {
+type writeCounter struct {
+	io.Writer
+	written int64
+}
+
+func (w *writeCounter) Write(p []byte) (int, error) {
+	n, err := w.Writer.Write(p)
+	w.written += int64(n)
+	return n, err
+}
+
+func (s *snapshotWriter) Release() {
 	if s.tx != nil {
 		// This is an in-memory rollback, no error expected.
 		_ = s.tx.Rollback()
 	}
 }
+
+type snapshotReader struct {
+	io.ReadCloser
+}
+
+var zstdMagic = []byte{0x28, 0xB5, 0x2F, 0xFD}
+
+func newSnapshotReader(snapshot io.ReadCloser) (*snapshotReader, error) {
+	b := bufio.NewReader(snapshot)
+	magic, err := b.Peek(4)
+	if err != nil {
+		return nil, err
+	}
+
+	s := snapshotReader{ReadCloser: io.NopCloser(b)}
+	if bytes.Equal(magic, zstdMagic) {
+		var dec *zstd.Decoder
+		if dec, err = zstd.NewReader(b); err != nil {
+			return nil, err
+		}
+		s.ReadCloser = dec.IOReadCloser()
+	}
+
+	return &s, nil
+}

pkg/experiment/metastore/metastore.go

@@ -33,9 +33,9 @@ import (
 type Config struct {
 	Address          string             `yaml:"address"`
 	GRPCClientConfig grpcclient.Config  `yaml:"grpc_client_config" doc:"description=Configures the gRPC client used to communicate with the metastore."`
-	DataDir          string             `yaml:"data_dir"`
 	MinReadyDuration time.Duration      `yaml:"min_ready_duration" category:"advanced"`
 	Raft             raft.Config        `yaml:"raft"`
+	FSM              fsm.Config         `yaml:",inline" category:"advanced"`
 	Index            index.Config       `yaml:",inline" category:"advanced"`
 	DLQRecovery      dlq.RecoveryConfig `yaml:",inline" category:"advanced"`
 	Compactor        compactor.Config   `yaml:",inline" category:"advanced"`
@@ -45,10 +45,10 @@ type Config struct {
 func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 	const prefix = "metastore."
 	f.StringVar(&cfg.Address, prefix+"address", "localhost:9095", "")
-	f.StringVar(&cfg.DataDir, prefix+"data-dir", "./data-metastore/data", "")
 	f.DurationVar(&cfg.MinReadyDuration, prefix+"min-ready-duration", 15*time.Second, "Minimum duration to wait after the internal readiness checks have passed but before succeeding the readiness endpoint. This is used to slowdown deployment controllers (eg. Kubernetes) after an instance is ready and before they proceed with a rolling update, to give the rest of the cluster instances enough time to receive some (DNS?) updates.")
 	cfg.GRPCClientConfig.RegisterFlagsWithPrefix(prefix+"grpc-client-config", f)
 	cfg.Raft.RegisterFlagsWithPrefix(prefix+"raft.", f)
+	cfg.FSM.RegisterFlagsWithPrefix(prefix, f)
 	cfg.Compactor.RegisterFlagsWithPrefix(prefix, f)
 	cfg.Scheduler.RegisterFlagsWithPrefix(prefix, f)
 	cfg.Index.RegisterFlagsWithPrefix(prefix, f)
@@ -117,9 +117,7 @@ func New(
 	}
 
 	var err error
-
-	m.fsm, err = fsm.New(m.logger, m.reg, m.config.DataDir)
-	if err != nil {
+	if m.fsm, err = fsm.New(m.logger, m.reg, m.config.FSM); err != nil {
 		return nil, fmt.Errorf("failed to initialize store: %w", err)
 	}
 

pkg/experiment/metastore/test/create.go

@@ -53,7 +53,7 @@ func NewMetastoreSet(t *testing.T, cfg *metastore.Config, n int, bucket objstore
 		icfg := *cfg
 		icfg.MinReadyDuration = 0
 		icfg.Address = grpcAddresses[i]
-		icfg.DataDir = t.TempDir()
+		icfg.FSM.DataDir = t.TempDir()
 		icfg.Raft.ServerID = raftIds[i]
 		icfg.Raft.Dir = t.TempDir()
 		icfg.Raft.AdvertiseAddress = raftAddresses[i]