[Raft] Add Raft persistence

.gitignore

@@ -135,3 +135,6 @@ go
 
 # terraform
 .terraform*
+
+# raft
+raft_data/

pkg/aux_/store/raftstore/store.go

@@ -1,6 +1,8 @@
 package raftstore
 
 import (
+	"context"
+
 	"github.com/interuss/dss/pkg/aux_/repos"
 	"github.com/interuss/dss/pkg/raftstore"
 	"go.uber.org/zap"
@@ -9,6 +11,6 @@ import (
 // repo is a full implementation of aux_.repos.Repository for Raft-based storage.
 type repo struct{}
 
-func Init(logger *zap.Logger) (*raftstore.Store[repos.Repository], error) {
-	return raftstore.Init[repos.Repository](logger, func() repos.Repository { return &repo{} })
+func Init(ctx context.Context, logger *zap.Logger) (*raftstore.Store[repos.Repository], error) {
+	return raftstore.Init[repos.Repository](ctx, logger, func() repos.Repository { return &repo{} })
 }

pkg/aux_/store/store.go

@@ -23,7 +23,7 @@ func Init(ctx context.Context, logger *zap.Logger, withCheckCron bool) (Store, e
 	case params.SQLStoreType:
 		return auxsqlstore.Init(ctx, logger, withCheckCron)
 	case params.RaftStoreType:
-		return auxraftstore.Init(logger)
+		return auxraftstore.Init(ctx, logger)
 	default:
 		return nil, stacktrace.NewError("Unsupported store type %q for aux", storeType)
 	}

pkg/raftstore/consensus/consensus.go

@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"net/url"
 
+	"github.com/interuss/dss/pkg/logging"
 	"github.com/interuss/stacktrace"
 	"go.etcd.io/etcd/client/pkg/v3/types"
 	"go.etcd.io/etcd/server/v3/etcdserver/api/rafthttp"
@@ -16,11 +17,14 @@ import (
 	"go.uber.org/zap"
 )
 
-const defaultClusterID uint64 = 1
+const (
+	defaultClusterID uint64 = 1
+)
 
 type Consensus struct {
-	node        raft.Node
-	raftStorage *raft.MemoryStorage
+	logger *zap.Logger
+
+	node raft.Node
 
 	transport *rafthttp.Transport
 	server    *http.Server
@@ -29,29 +33,37 @@ type Consensus struct {
 	errorC  chan error
 }
 
-func NewConsensus(logger *zap.Logger, nodeID uint64, peers map[uint64]*url.URL) (*Consensus, error) {
+func NewConsensus(ctx context.Context, logger *zap.Logger, nodeID uint64, peers map[uint64]*url.URL, dataDir string, snapshotCatchupEntries uint64) (*Consensus, error) {
+	storage, _, err := newStorage(ctx, logger.With(zap.String("component", "storage")), dataDir, nodeID, snapshotCatchupEntries)
+	if err != nil {
+		return nil, stacktrace.Propagate(err, "failed to initialize storage")
+	}
+
 	consensus := &Consensus{
-		errorC: make(chan error, 1),
+		logger: logging.WithValuesFromContext(ctx, logger),
+
+		storage: storage,
+		errorC:  make(chan error, 1),
 	}
 
-	err := consensus.initTransport(logger.With(zap.String("component", "transport")), nodeID, defaultClusterID, peers)
+	err = consensus.initTransport(ctx, nodeID, defaultClusterID, peers)
 	if err != nil {
 		return nil, stacktrace.Propagate(err, "failed to initialize transport")
 	}
 
 	return consensus, nil
 }
 
-func (c *Consensus) initTransport(logger *zap.Logger, nodeID uint64, clusterID uint64, peers map[uint64]*url.URL) error {
+func (c *Consensus) initTransport(ctx context.Context, nodeID uint64, clusterID uint64, peers map[uint64]*url.URL) error {
 	nodeIDStr := fmt.Sprintf("%d", nodeID)
 
 	transport := &rafthttp.Transport{
-		Logger:      logger,
+		Logger:      logging.WithValuesFromContext(ctx, c.logger.With(zap.String("component", "transport"))),
 		ID:          types.ID(nodeID),
 		ClusterID:   types.ID(clusterID),
 		Raft:        c,
 		ServerStats: v2stats.NewServerStats(nodeIDStr, nodeIDStr),
-		LeaderStats: v2stats.NewLeaderStats(logger, nodeIDStr),
+		LeaderStats: v2stats.NewLeaderStats(c.logger, nodeIDStr),
 		ErrorC:      c.errorC,
 	}
 
@@ -82,7 +94,7 @@ func (c *Consensus) initTransport(logger *zap.Logger, nodeID uint64, clusterID u
 	go func() {
 		err := c.server.ListenAndServe()
 		if err != nil && !errors.Is(err, http.ErrServerClosed) {
-			logger.Error("http server error", zap.Error(err))
+			c.logger.Error("http server error", zap.Error(err))
 			c.errorC <- err
 		}
 	}()
@@ -110,3 +122,8 @@ func (c *Consensus) ReportUnreachable(id uint64) {
 func (c *Consensus) ReportSnapshot(id uint64, status raft.SnapshotStatus) {
 	c.node.ReportSnapshot(id, status)
 }
+
+// RegisterStore allows registering a snapshot provider function for a specific store
+func (c *Consensus) RegisterStore(name string, provider snapshotProvider) {
+	c.storage.registerSnapshotProvider(name, provider)
+}

pkg/raftstore/consensus/storage.go

@@ -1,12 +1,268 @@
 package consensus
 
 import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"os"
+	"path"
+
+	"github.com/interuss/dss/pkg/logging"
+	"github.com/interuss/stacktrace"
+	"go.etcd.io/etcd/client/pkg/v3/fileutil"
 	"go.etcd.io/etcd/server/v3/etcdserver/api/snap"
 	"go.etcd.io/etcd/server/v3/storage/wal"
+	"go.etcd.io/etcd/server/v3/storage/wal/walpb"
+	"go.etcd.io/raft/v3"
+	"go.etcd.io/raft/v3/raftpb"
+	"go.uber.org/zap"
 )
 
-// storage persists the Raft log and snapshots.
+// snapshotProvider is a function that returns the snapshot data to be included in the Raft snapshot.
+// We use a snapshotProvider from each component (scd, rid and aux).
+type snapshotProvider func() ([]byte, error)
+
+// storage persists the Raft log and snapshots and manages the raft in-memory storage.
 type storage struct {
-	wal         *wal.WAL
-	snapshotter *snap.Snapshotter
+	logger *zap.Logger
+
+	*raft.MemoryStorage
+
+	wal *wal.WAL
+
+	snapper   *snap.Snapshotter
+	providers map[string]snapshotProvider
+
+	snapshotCatchUpEntries uint64
+}
+
+// newStorage initializes the storage by loading the latest snapshot and wal entries from the disk
+// and applies them to the Raft memory storage.
+// It returns the initialized storage, a boolean indicating whether the storage was pre-existent or an error.
+func newStorage(ctx context.Context, logger *zap.Logger, dataDir string, nodeID uint64, snapshotCatchUpEntries uint64) (*storage, bool, error) {
+	logger = logging.WithValuesFromContext(ctx, logger)
+
+	// load the latest snapshot
+	snapshotPath := path.Join(dataDir, fmt.Sprintf("snapshot_%d", nodeID))
+	if !fileutil.Exist(snapshotPath) {
+		err := os.MkdirAll(snapshotPath, 0o750)
+		if err != nil {
+			return nil, false, stacktrace.Propagate(err, "failed to create directory for snapshot storage at: %s", snapshotPath)
+		}
+	}
+
+	walPath := path.Join(dataDir, fmt.Sprintf("wal_%d", nodeID))
+
+	snapper := snap.New(logger, snapshotPath)
+
+	var (
+		w        *wal.WAL
+		err      error
+		snapshot *raftpb.Snapshot
+	)
+	ok := wal.Exist(walPath)
+	if !ok {
+		if err = os.MkdirAll(walPath, 0o750); err != nil {
+			return nil, false, stacktrace.Propagate(err, "failed to create directory for wal storage at: %s", walPath)
+		}
+		w, err = wal.Create(logger, walPath, nil)
+		if err != nil {
+			return nil, false, stacktrace.Propagate(err, "failed to create wal at: %s", walPath)
+		}
+		if err = w.Close(); err != nil {
+			return nil, false, stacktrace.Propagate(err, "failed to close wal at: %s", walPath)
+		}
+		snapshot = &raftpb.Snapshot{}
+	} else {
+		snapshot, err = loadSnapshot(logger, walPath, snapper)
+		if err != nil {
+			return nil, false, stacktrace.Propagate(err, "failed to load snapshot")
+		}
+	}
+
+	// open the wal at the given snapshot and get all subsequent entries
+	w, err = wal.Open(logger, walPath, walpb.Snapshot{
+		Index:     snapshot.Metadata.Index,
+		Term:      snapshot.Metadata.Term,
+		ConfState: &snapshot.Metadata.ConfState,
+	})
+	if err != nil {
+		return nil, false, stacktrace.Propagate(err, "failed to open wal with index %d, term %d and confstate %v at: %s", snapshot.Metadata.Index, snapshot.Metadata.Term, snapshot.Metadata.ConfState, walPath)
+	}
+
+	_, state, entries, err := w.ReadAll()
+	if err != nil {
+		return nil, false, stacktrace.Propagate(err, "failed to read wal state and entries")
+	}
+
+	// initialize the raft memory storage with the loaded snapshot and wal entries
+	raftMemoryStorage := raft.NewMemoryStorage()
+	if !raft.IsEmptySnap(*snapshot) {
+		err = raftMemoryStorage.ApplySnapshot(*snapshot)
+		if err != nil {
+			return nil, false, stacktrace.Propagate(err, "failed to apply snapshot to raft memory storage")
+		}
+	}
+
+	err = raftMemoryStorage.SetHardState(state)
+	if err != nil {
+		return nil, false, stacktrace.Propagate(err, "failed to set hard state to raft memory storage")
+	}
+
+	err = raftMemoryStorage.Append(entries)
+	if err != nil {
+		return nil, false, stacktrace.Propagate(err, "failed to append entries to raft memory storage")
+	}
+
+	logger.Info("Loaded previous hardstate and entries to Raft memory storage", zap.Any("hard-state", state), zap.Int("entries-number", len(entries)))
+
+	return &storage{
+		logger: logger,
+
+		MemoryStorage: raftMemoryStorage,
+
+		wal: w,
+
+		snapper:                snapper,
+		providers:              make(map[string]snapshotProvider),
+		snapshotCatchUpEntries: snapshotCatchUpEntries,
+	}, ok, nil
+}
+
+func loadSnapshot(logger *zap.Logger, walPath string, snapshotter *snap.Snapshotter) (*raftpb.Snapshot, error) {
+	entries, err := wal.ValidSnapshotEntries(logger, walPath)
+	if err != nil {
+		return nil, stacktrace.Propagate(err, "failed to get valid snapshot entries")
+	}
+
+	snapshot, err := snapshotter.LoadNewestAvailable(entries)
+	if err != nil {
+		if errors.Is(err, snap.ErrNoSnapshot) {
+			return &raftpb.Snapshot{}, nil
+		}
+
+		return nil, stacktrace.Propagate(err, "failed to load newest available snapshot")
+	}
+
+	logger.Info("Loaded snapshot", zap.Uint64("index", snapshot.Metadata.Index), zap.Uint64("term", snapshot.Metadata.Term))
+	return snapshot, nil
+}
+
+// save saves the given snapshot to the snapshotter and the wal.
+func (s *storage) save(snapshot raftpb.Snapshot) error {
+	err := s.snapper.SaveSnap(snapshot)
+	if err != nil {
+		return stacktrace.Propagate(err, "failed to save snapshot")
+	}
+
+	err = s.wal.SaveSnapshot(walpb.Snapshot{
+		Index:     snapshot.Metadata.Index,
+		Term:      snapshot.Metadata.Term,
+		ConfState: &snapshot.Metadata.ConfState,
+	})
+	if err != nil {
+		return stacktrace.Propagate(err, "failed to save snapshot to wal")
+	}
+
+	// ReleaseLockTo releases the os-level flock locks held on WAL segment files
+	// for entries covered by the snapshot and that are no longer needed.
+	// We only call this on success. If saving the snapshot failed above, the files
+	// are still needed and must remain locked.
+	//
+	// In etcd's implementation this call marks files eligible for deletion
+	// by a background cleanup goroutine which we do not have yet.
+	//
+	// TODO: add a separate goroutine to cleanup old files
+	return s.wal.ReleaseLockTo(snapshot.Metadata.Index)
+}
+
+func (s *storage) registerSnapshotProvider(name string, provider func() ([]byte, error)) {
+	s.providers[name] = provider
+}
+
+// getSnapshot calls all registered snapshot providers and combines their data into a single snapshot.
+func (s *storage) getSnapshot() ([]byte, error) {
+	parts := make(map[string][]byte)
+	for name, provider := range s.providers {
+		data, err := provider()
+		if err != nil {
+			return nil, stacktrace.Propagate(err, "failed to get snapshot data from %q", name)
+		}
+
+		parts[name] = data
+	}
+
+	return json.Marshal(parts)
+}
+
+// snapshotter returns the snapshotter used by the storage.
+func (s *storage) snapshotter() *snap.Snapshotter {
+	return s.snapper
+}
+
+func (s *storage) triggerSnapshot(appliedIndex uint64, confState *raftpb.ConfState) error {
+	s.logger.Info("triggering snapshot", zap.Uint64("appliedIndex", appliedIndex))
+	data, err := s.getSnapshot()
+	if err != nil {
+		return stacktrace.Propagate(err, "failed to get snapshot data")
+	}
+
+	snap, err := s.CreateSnapshot(appliedIndex, confState, data)
+	if err != nil {
+		return stacktrace.Propagate(err, "failed to create snapshot from raft memory storage")
+	}
+
+	err = s.save(snap)
+	if err != nil {
+		return stacktrace.Propagate(err, "failed to save snapshot")
+	}
+
+	compactIndex := uint64(1)
+	if appliedIndex > s.snapshotCatchUpEntries {
+		compactIndex = appliedIndex - s.snapshotCatchUpEntries
+	}
+
+	return s.compactLog(compactIndex)
+}
+
+func (s *storage) compactLog(compactIndex uint64) error {
+	err := s.Compact(compactIndex)
+	if errors.Is(err, raft.ErrCompacted) {
+		s.logger.Warn("log already compacted", zap.Uint64("compactIndex", compactIndex))
+		return nil
+	}
+	if err != nil {
+		return stacktrace.Propagate(err, "failed to compact raft memory storage at index %d", compactIndex)
+	}
+
+	s.logger.Info("compacted log", zap.Uint64("compactIndex", compactIndex))
+	return nil
+}
+
+func (s *storage) handleReceivedState(snapshot raftpb.Snapshot, hardState raftpb.HardState, entries []raftpb.Entry) error {
+	if !raft.IsEmptySnap(snapshot) {
+		err := s.save(snapshot)
+		if err != nil {
+			return stacktrace.Propagate(err, "failed to save snapshot")
+		}
+	}
+
+	err := s.wal.Save(hardState, entries)
+	if err != nil {
+		return stacktrace.Propagate(err, "failed to save WAL entries")
+	}
+
+	if !raft.IsEmptySnap(snapshot) {
+		err := s.ApplySnapshot(snapshot)
+		if err != nil {
+			return stacktrace.Propagate(err, "failed to apply snapshot")
+		}
+	}
+
+	if err := s.Append(entries); err != nil {
+		return stacktrace.Propagate(err, "failed to append entries to raft storage")
+	}
+
+	return nil
 }