Filter replicated databases

client.go

@@ -22,7 +22,8 @@ type Client interface {
 	Commit(ctx context.Context, primaryURL string, nodeID uint64, name string, lockID int64, r io.Reader) error
 
 	// Stream starts a long-running connection to stream changes from another node.
-	Stream(ctx context.Context, primaryURL string, nodeID uint64, posMap map[string]ltx.Pos) (Stream, error)
+	// If filter is specified, only those databases will be replicated.
+	Stream(ctx context.Context, primaryURL string, nodeID uint64, posMap map[string]ltx.Pos, filter []string) (Stream, error)
 }
 
 // Stream represents a stream of frames.

cmd/litefs/config.go

@@ -150,6 +150,9 @@ type LeaseConfig struct {
 	// become primary again.
 	DemoteDelay time.Duration `yaml:"demote-delay"`
 
+	// Specifies a subset of databases to replica.
+	Databases []string `yaml:"databases"`
+
 	// Consul lease settings.
 	Consul struct {
 		URL       string        `yaml:"url"`

cmd/litefs/mount_linux.go

@@ -46,6 +46,8 @@ type MountCommand struct {
 
 	// Used for generating the advertise URL for testing.
 	AdvertiseURLFn func() string
+
+	OnInitStore func()
 }
 
 // NewMountCommand returns a new instance of MountCommand.
@@ -158,6 +160,10 @@ func (c *MountCommand) Validate(ctx context.Context) (err error) {
 		return fmt.Errorf("invalid lease type, must be either 'consul' or 'static', got: '%v'", c.Config.Lease.Type)
 	}
 
+	if c.Config.Lease.Candidate && len(c.Config.Lease.Databases) > 0 {
+		return fmt.Errorf("cannot specify a database replication filter on candidate nodes")
+	}
+
 	return nil
 }
 
@@ -362,8 +368,13 @@ func (c *MountCommand) initStore(ctx context.Context) error {
 	c.Store.ReconnectDelay = c.Config.Lease.ReconnectDelay
 	c.Store.DemoteDelay = c.Config.Lease.DemoteDelay
 	c.Store.Client = http.NewClient()
+	c.Store.DatabaseFilter = c.Config.Lease.Databases
 	c.initEnvironment(ctx)
 
+	if c.OnInitStore != nil {
+		c.OnInitStore()
+	}
+
 	if err := c.initStoreBackupClient(ctx); err != nil {
 		return err
 	}

cmd/litefs/mount_test.go

@@ -1935,6 +1935,40 @@ func TestMultiNode_Autopromotion(t *testing.T) {
 	}
 }
 
+func TestMultiNode_DatabaseFilter(t *testing.T) {
+	cmd0 := runMountCommand(t, newMountCommand(t, t.TempDir(), nil))
+	waitForPrimary(t, cmd0)
+	cmd1 := newMountCommand(t, t.TempDir(), cmd0)
+	cmd1.OnInitStore = func() {
+		cmd1.Store.DatabaseFilter = []string{"x.db"}
+	}
+	runMountCommand(t, cmd1)
+
+	db0 := testingutil.OpenSQLDB(t, filepath.Join(cmd0.Config.FUSE.Dir, "x.db"))
+	if _, err := db0.Exec(`CREATE TABLE t (x)`); err != nil {
+		t.Fatal(err)
+	} else if err := db0.Close(); err != nil {
+		t.Fatal(err)
+	}
+
+	db1 := testingutil.OpenSQLDB(t, filepath.Join(cmd0.Config.FUSE.Dir, "y.db"))
+	if _, err := db1.Exec(`CREATE TABLE t (y)`); err != nil {
+		t.Fatal(err)
+	} else if err := db1.Close(); err != nil {
+		t.Fatal(err)
+	}
+
+	waitForSync(t, "x.db", cmd0, cmd1)
+
+	// Only the filtered database should exist.
+	if _, err := os.Stat(filepath.Join(cmd1.Config.FUSE.Dir, "x.db")); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := os.Stat(filepath.Join(cmd1.Config.FUSE.Dir, "y.db")); !os.IsNotExist(err) {
+		t.Fatal("expected second database to not exist on replica")
+	}
+}
+
 func TestMultiNode_StaticLeaser(t *testing.T) {
 	dir0, dir1 := t.TempDir(), t.TempDir()
 	cmd0 := newMountCommand(t, dir0, nil)

http/client.go

@@ -12,6 +12,7 @@ import (
 	"net/http"
 	"net/url"
 	"strconv"
+	"strings"
 
 	"github.com/superfly/litefs"
 	"github.com/superfly/litefs/internal/chunk"
@@ -354,7 +355,7 @@ func (c *Client) Commit(ctx context.Context, primaryURL string, nodeID uint64, n
 }
 
 // Stream returns a snapshot and continuous stream of WAL updates.
-func (c *Client) Stream(ctx context.Context, primaryURL string, nodeID uint64, posMap map[string]ltx.Pos) (litefs.Stream, error) {
+func (c *Client) Stream(ctx context.Context, primaryURL string, nodeID uint64, posMap map[string]ltx.Pos, filter []string) (litefs.Stream, error) {
 	u, err := url.Parse(primaryURL)
 	if err != nil {
 		return nil, fmt.Errorf("invalid client URL: %w", err)
@@ -364,11 +365,17 @@ func (c *Client) Stream(ctx context.Context, primaryURL string, nodeID uint64, p
 		return nil, fmt.Errorf("URL host required")
 	}
 
+	q := make(url.Values)
+	if len(filter) > 0 {
+		q.Set("filter", strings.Join(filter, ","))
+	}
+
 	// Strip off everything but the scheme & host.
 	*u = url.URL{
-		Scheme: u.Scheme,
-		Host:   u.Host,
-		Path:   "/stream",
+		Scheme:   u.Scheme,
+		Host:     u.Host,
+		Path:     "/stream",
+		RawQuery: q.Encode(),
 	}
 
 	var buf bytes.Buffer

http/server.go

@@ -492,6 +492,7 @@ func (s *Server) handlePostStream(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, "Upgrade to HTTP/2 required", http.StatusUpgradeRequired)
 		return
 	}
+	q := r.URL.Query()
 
 	// Prevent nodes from connecting to themselves.
 	id, _ := litefs.ParseNodeID(r.Header.Get(HeaderNodeID))
@@ -536,6 +537,14 @@ func (s *Server) handlePostStream(w http.ResponseWriter, r *http.Request) {
 		dirtySet[db.Name()] = struct{}{}
 	}
 
+	// Determine filtered set of databases, if any.
+	filterSet := make(map[string]struct{})
+	if filter := q.Get("filter"); filter != "" {
+		for _, name := range strings.Split(filter, ",") {
+			filterSet[name] = struct{}{}
+		}
+	}
+
 	// Flush header so client can resume control.
 	w.WriteHeader(http.StatusOK)
 	w.(http.Flusher).Flush()
@@ -554,6 +563,15 @@ func (s *Server) handlePostStream(w http.ResponseWriter, r *http.Request) {
 	var readySent bool
 	var handoffLeaseID string
 	for {
+		// Restrict dirty set to only databases in the filter set.
+		if len(filterSet) > 0 && len(dirtySet) > 0 {
+			for name := range dirtySet {
+				if _, ok := filterSet[name]; !ok {
+					delete(dirtySet, name)
+				}
+			}
+		}
+
 		// Send pending transactions for each database.
 		for name := range dirtySet {
 			if err := s.streamDB(r.Context(), w, name, posMap); err != nil {

mock/client.go

@@ -14,7 +14,7 @@ type Client struct {
 	AcquireHaltLockFunc func(ctx context.Context, primaryURL string, nodeID uint64, name string, lockID int64) (*litefs.HaltLock, error)
 	ReleaseHaltLockFunc func(ctx context.Context, primaryURL string, nodeID uint64, name string, lockID int64) error
 	CommitFunc          func(ctx context.Context, primaryURL string, nodeID uint64, name string, lockID int64, r io.Reader) error
-	StreamFunc          func(ctx context.Context, primaryURL string, nodeID uint64, posMap map[string]ltx.Pos) (litefs.Stream, error)
+	StreamFunc          func(ctx context.Context, primaryURL string, nodeID uint64, posMap map[string]ltx.Pos, filter []string) (litefs.Stream, error)
 }
 
 func (c *Client) AcquireHaltLock(ctx context.Context, primaryURL string, nodeID uint64, name string, lockID int64) (*litefs.HaltLock, error) {
@@ -29,8 +29,8 @@ func (c *Client) Commit(ctx context.Context, primaryURL string, nodeID uint64, n
 	return c.CommitFunc(ctx, primaryURL, nodeID, name, lockID, r)
 }
 
-func (c *Client) Stream(ctx context.Context, primaryURL string, nodeID uint64, posMap map[string]ltx.Pos) (litefs.Stream, error) {
-	return c.StreamFunc(ctx, primaryURL, nodeID, posMap)
+func (c *Client) Stream(ctx context.Context, primaryURL string, nodeID uint64, posMap map[string]ltx.Pos, filter []string) (litefs.Stream, error) {
+	return c.StreamFunc(ctx, primaryURL, nodeID, posMap, filter)
 }
 
 type Stream struct {

store.go

@@ -127,6 +127,9 @@ type Store struct {
 	// Interface to interact with the host environment.
 	Environment Environment
 
+	// Specifies a subset of databases to replicate from the primary.
+	DatabaseFilter []string
+
 	// If true, computes and verifies the checksum of the entire database
 	// after every transaction. Should only be used during testing.
 	StrictVerify bool
@@ -1357,7 +1360,7 @@ func (s *Store) monitorLeaseAsReplica(ctx context.Context, info PrimaryInfo) (ha
 	}()
 
 	posMap := s.PosMap()
-	st, err := s.Client.Stream(ctx, info.AdvertiseURL, s.id, posMap)
+	st, err := s.Client.Stream(ctx, info.AdvertiseURL, s.id, posMap, s.DatabaseFilter)
 	if err != nil {
 		return "", fmt.Errorf("connect to primary: %s ('%s')", err, info.AdvertiseURL)
 	}

store_test.go

@@ -220,7 +220,7 @@ func TestStore_PrimaryCtx(t *testing.T) {
 		}
 
 		client := mock.Client{
-			StreamFunc: func(ctx context.Context, rawurl string, nodeID uint64, posMap map[string]ltx.Pos) (litefs.Stream, error) {
+			StreamFunc: func(ctx context.Context, rawurl string, nodeID uint64, posMap map[string]ltx.Pos, filter []string) (litefs.Stream, error) {
 				return &mock.Stream{
 					ReadCloser:    io.NopCloser(&bytes.Buffer{}),
 					ClusterIDFunc: func() string { return "" },
@@ -253,7 +253,7 @@ func TestStore_PrimaryCtx(t *testing.T) {
 	t.Run("InitialReplica", func(t *testing.T) {
 		leaser := litefs.NewStaticLeaser(false, "localhost", "http://localhost:20202")
 		client := mock.Client{
-			StreamFunc: func(ctx context.Context, rawurl string, nodeID uint64, posMap map[string]ltx.Pos) (litefs.Stream, error) {
+			StreamFunc: func(ctx context.Context, rawurl string, nodeID uint64, posMap map[string]ltx.Pos, filter []string) (litefs.Stream, error) {
 				var buf bytes.Buffer
 				if err := litefs.WriteStreamFrame(&buf, &litefs.ReadyStreamFrame{}); err != nil {
 					return nil, err