Shared index.

Author: ncruces

litestream/api.go

@@ -19,15 +19,10 @@ func init() {
 
 var (
 	liteMtx sync.RWMutex
-	// +checklocks:memoryMtx
-	liteDBs = map[string]liteDB{}
+	// +checklocks:liteMtx
+	liteDBs = map[string]*liteDB{}
 )
 
-type liteDB struct {
-	client litestream.ReplicaClient
-	opts   *ReplicaOptions
-}
-
 // ReplicaOptions represents options for [NewReplica].
 type ReplicaOptions struct {
 	// Where to log error messages. May be nil.
@@ -53,7 +48,7 @@ func NewReplica(name string, client litestream.ReplicaClient, options ReplicaOpt
 
 	liteMtx.Lock()
 	defer liteMtx.Unlock()
-	liteDBs[name] = liteDB{
+	liteDBs[name] = &liteDB{
 		client: client,
 		opts:   &options,
 	}

litestream/go.mod

@@ -5,11 +5,31 @@ go 1.24.4
 require (
 	github.com/benbjohnson/litestream v0.5.0
 	github.com/ncruces/go-sqlite3 v0.29.1
+	github.com/ncruces/wbt v0.2.0
 	github.com/superfly/ltx v0.5.0
 )
 
 require (
 	filippo.io/age v1.1.1 // indirect
+	github.com/aws/aws-sdk-go-v2 v1.37.1 // indirect
+	github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.0 // indirect
+	github.com/aws/aws-sdk-go-v2/config v1.30.2 // indirect
+	github.com/aws/aws-sdk-go-v2/credentials v1.18.2 // indirect
+	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.1 // indirect
+	github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.18.2 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.1 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.1 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/ini v1.8.3 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.1 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.0 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.8.1 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.1 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.1 // indirect
+	github.com/aws/aws-sdk-go-v2/service/s3 v1.85.1 // indirect
+	github.com/aws/aws-sdk-go-v2/service/sso v1.26.1 // indirect
+	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.31.1 // indirect
+	github.com/aws/aws-sdk-go-v2/service/sts v1.35.1 // indirect
+	github.com/aws/smithy-go v1.22.5 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/dustin/go-humanize v1.0.1 // indirect

litestream/go.sum

@@ -118,6 +118,8 @@ github.com/ncruces/go-strftime v0.1.9 h1:bY0MQC28UADQmHmaF5dgpLmImcShSi2kHU9XLdh
 github.com/ncruces/go-strftime v0.1.9/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
 github.com/ncruces/julianday v1.0.0 h1:fH0OKwa7NWvniGQtxdJRxAgkBMolni2BjDHaWTxqt7M=
 github.com/ncruces/julianday v1.0.0/go.mod h1:Dusn2KvZrrovOMJuOt0TNXL6tB7U2E8kvza5fFc9G7g=
+github.com/ncruces/wbt v0.2.0 h1:Q9zlKOBSZc7Yy/R2cGa35g6RKUUE3BjNIW3tfGC4F04=
+github.com/ncruces/wbt v0.2.0/go.mod h1:DtF92amvMxH69EmBFUSFWRDAlo6hOEfoNQnClxj9C/c=
 github.com/pierrec/lz4/v4 v4.1.22 h1:cKFw6uJDK+/gfw5BcDL0JL5aBsAFdsIT18eRtLj7VIU=
 github.com/pierrec/lz4/v4 v4.1.22/go.mod h1:gZWDp/Ze/IJXGXf23ltt2EXimqmTUXEy0GFuRQyBid4=
 github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ=

litestream/vfs.go

@@ -6,12 +6,14 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"sync"
 	"time"
 
 	"github.com/benbjohnson/litestream"
 	"github.com/ncruces/go-sqlite3"
 	"github.com/ncruces/go-sqlite3/util/vfsutil"
 	"github.com/ncruces/go-sqlite3/vfs"
+	"github.com/ncruces/wbt"
 	"github.com/superfly/ltx"
 )
 
@@ -30,21 +32,14 @@ func (liteVFS) Open(name string, flags vfs.OpenFlag) (vfs.File, vfs.OpenFlag, er
 	liteMtx.RLock()
 	defer liteMtx.RUnlock()
 	if db, ok := liteDBs[name]; ok {
-		f := liteFile{
-			liteDB: db,
-			txids:  new(levelTXIDs),
-			pages:  map[uint32]ltx.PageIndexElem{},
-		}
-
 		// Build the page index so we can lookup individual pages.
-		if err := f.buildIndex(context.Background()); err != nil {
-			f.opts.Logger.Error("build index", "error", err)
+		if err := db.buildIndex(context.Background()); err != nil {
+			db.opts.Logger.Error("build index", "error", err)
 			return nil, 0, err
 		}
-		return &f, flags | vfs.OPEN_READONLY, nil
+		return &liteFile{db: db}, flags | vfs.OPEN_READONLY, nil
 	}
 	return nil, flags, sqlite3.CANTOPEN
-
 }
 
 func (liteVFS) Delete(name string, dirSync bool) error {
@@ -62,22 +57,21 @@ func (liteVFS) FullPathname(name string) (string, error) {
 }
 
 type liteFile struct {
-	liteDB
-
-	conn  *sqlite3.Conn
-	txids *levelTXIDs
-	pages map[uint32]ltx.PageIndexElem
-
-	lastPoll  time.Time
-	pageSize  uint32
-	pageCount uint32
-	lock      vfs.LockLevel
+	db       *liteDB
+	conn     *sqlite3.Conn
+	pages    *pageIndex
+	txid     ltx.TXID
+	pageSize uint32
 }
 
 func (f *liteFile) Close() error { return nil }
 
 func (f *liteFile) ReadAt(p []byte, off int64) (n int, err error) {
-	err = f.pollReplica()
+	ctx := f.context()
+	pages, txid := f.pages, f.txid
+	if pages == nil {
+		pages, txid, err = f.db.pollReplica(ctx)
+	}
 	if err != nil {
 		return 0, err
 	}
@@ -87,19 +81,14 @@ func (f *liteFile) ReadAt(p []byte, off int64) (n int, err error) {
 		pgno += uint32(off / int64(f.pageSize))
 	}
 
-	elem, ok := f.pages[pgno]
+	elem, ok := pages.Get(pgno)
 	if !ok {
 		return 0, io.EOF
 	}
 
-	ctx := context.Background()
-	if f.conn != nil {
-		ctx = f.conn.GetInterrupt()
-	}
-
-	_, data, err := litestream.FetchPage(ctx, f.client, elem.Level, elem.MinTXID, elem.MaxTXID, elem.Offset, elem.Size)
+	_, data, err := litestream.FetchPage(ctx, f.db.client, elem.Level, elem.MinTXID, elem.MaxTXID, elem.Offset, elem.Size)
 	if err != nil {
-		f.opts.Logger.Error("fetch page", "error", err)
+		f.db.opts.Logger.Error("fetch page", "error", err)
 		return 0, err
 	}
 
@@ -109,7 +98,7 @@ func (f *liteFile) ReadAt(p []byte, off int64) (n int, err error) {
 		data[18] >= 1 && data[19] >= 1 &&
 		data[18] <= 3 && data[19] <= 3 {
 		data[18], data[19] = 0x01, 0x01
-		binary.BigEndian.PutUint32(data[24:28], uint32(f.txids[0]))
+		binary.BigEndian.PutUint32(data[24:28], uint32(txid))
 		f.pageSize = uint32(256 * binary.LittleEndian.Uint16(data[16:18]))
 	}
 
@@ -133,23 +122,22 @@ func (f *liteFile) Sync(flag vfs.SyncFlag) error {
 }
 
 func (f *liteFile) Size() (size int64, err error) {
-	size = int64(f.pageCount) * int64(f.pageSize)
+	if max := f.pages.Max(); max != nil {
+		size = int64(max.Key()) * int64(f.pageSize)
+	}
 	return
 }
 
-func (f *liteFile) Lock(lock vfs.LockLevel) error {
+func (f *liteFile) Lock(lock vfs.LockLevel) (err error) {
 	if lock >= vfs.LOCK_RESERVED {
 		return sqlite3.IOERR_LOCK
 	}
-	if err := f.pollReplica(); err != nil {
-		return err
-	}
-	f.lock = max(f.lock, lock)
-	return nil
+	f.pages, f.txid, err = f.db.pollReplica(f.context())
+	return err
 }
 
 func (f *liteFile) Unlock(lock vfs.LockLevel) error {
-	f.lock = min(f.lock, lock)
+	f.pages, f.txid = nil, 0
 	return nil
 }
 
@@ -172,7 +160,32 @@ func (f *liteFile) SetDB(conn any) {
 	f.conn = conn.(*sqlite3.Conn)
 }
 
-func (f *liteFile) buildIndex(ctx context.Context) error {
+func (f *liteFile) context() context.Context {
+	if f.conn != nil {
+		return f.conn.GetInterrupt()
+	}
+	return context.Background()
+}
+
+type liteDB struct {
+	client   litestream.ReplicaClient
+	opts     *ReplicaOptions
+	pages    *pageIndex // +checklocks:mtx
+	lastPoll time.Time  // +checklocks:mtx
+	txids    levelTXIDs // +checklocks:mtx
+	mtx      sync.Mutex
+}
+
+func (f *liteDB) buildIndex(ctx context.Context) error {
+	f.mtx.Lock()
+	defer f.mtx.Unlock()
+
+	// Skip if we already have an index.
+	if f.pages != nil {
+		return nil
+	}
+
+	// Build the index from scratch from a Litestream restore plan.
 	infos, err := litestream.CalcRestorePlan(ctx, f.client, 0, time.Time{}, f.opts.Logger)
 	if err != nil {
 		if !errors.Is(err, litestream.ErrTxNotAvailable) {
@@ -182,7 +195,7 @@ func (f *liteFile) buildIndex(ctx context.Context) error {
 	}
 
 	for _, info := range infos {
-		err := f.updateIndex(ctx, info)
+		err := f.updateInfo(ctx, info)
 		if err != nil {
 			return err
 		}
@@ -192,81 +205,86 @@ func (f *liteFile) buildIndex(ctx context.Context) error {
 	return nil
 }
 
-func (f *liteFile) pollReplica() error {
-	// Can't poll in a transaction.
-	if f.lock > vfs.LOCK_NONE {
-		return nil
-	}
+func (f *liteDB) pollReplica(ctx context.Context) (*pageIndex, ltx.TXID, error) {
+	f.mtx.Lock()
+	defer f.mtx.Unlock()
 
 	// Limit polling interval.
 	if time.Since(f.lastPoll) < f.opts.PollInterval {
-		return nil
-	}
-
-	ctx := context.Background()
-	if f.conn != nil {
-		ctx = f.conn.GetInterrupt()
+		return f.pages, f.txids[0], nil
 	}
 
+	// Updating from MinLevel to SnapshotLevel is non-racy,
+	// since LTX files are compacted into higher levels
+	// before the lower level LTX files are deleted.
 	for level := f.opts.MinLevel; level <= litestream.SnapshotLevel; level++ {
-		err := func() error {
-			var nextTXID ltx.TXID
-			// Snapshots must start from scratch,
-			// other levels can start from where they were left.
-			if level != litestream.SnapshotLevel {
-				nextTXID = f.txids[level] + 1
-			}
-
-			// Start reading from the next LTX file after the current position.
-			itr, err := f.client.LTXFiles(ctx, level, nextTXID)
-			if err != nil {
-				return fmt.Errorf("ltx files: %w", err)
-			}
-			defer itr.Close()
-
-			// Build an update across all new LTX files.
-			for itr.Next() {
-				info := itr.Item()
-
-				// Skip LTX files already in the index.
-				if info.MaxTXID <= f.txids[level] {
-					continue
-				}
-
-				err := f.updateIndex(ctx, info)
-				if err != nil {
-					return err
-				}
-			}
-			if err := itr.Err(); err != nil {
-				return err
-			}
-			return itr.Close()
-		}()
-		if err != nil {
+		if err := f.updateLevel(ctx, level); err != nil {
 			f.opts.Logger.Error("cannot poll replica", "error", err)
-			return err
+			return nil, 0, err
 		}
 	}
 
 	f.lastPoll = time.Now()
-	return nil
+	return f.pages, f.txids[0], nil
 }
 
-func (f *liteFile) updateIndex(ctx context.Context, info *ltx.FileInfo) error {
-	// Read page index.
+// +checklocks:f.mtx
+func (f *liteDB) updateLevel(ctx context.Context, level int) error {
+	var nextTXID ltx.TXID
+	// Snapshots must start from scratch,
+	// other levels can start from where they were left.
+	if level != litestream.SnapshotLevel {
+		nextTXID = f.txids[level] + 1
+	}
+
+	// Start reading from the next LTX file after the current position.
+	itr, err := f.client.LTXFiles(ctx, level, nextTXID)
+	if err != nil {
+		return fmt.Errorf("ltx files: %w", err)
+	}
+	defer itr.Close()
+
+	// Build an update across all new LTX files.
+	for itr.Next() {
+		info := itr.Item()
+
+		// Skip LTX files already fully loaded into the index.
+		if info.MaxTXID <= f.txids[level] {
+			continue
+		}
+
+		err := f.updateInfo(ctx, info)
+		if err != nil {
+			return err
+		}
+	}
+	if err := itr.Err(); err != nil {
+		return err
+	}
+	return itr.Close()
+}
+
+// +checklocks:f.mtx
+func (f *liteDB) updateInfo(ctx context.Context, info *ltx.FileInfo) error {
 	idx, err := litestream.FetchPageIndex(ctx, f.client, info)
 	if err != nil {
 		return fmt.Errorf("fetch page index: %w", err)
 	}
 
-	// Replace pages in overall index with new pages.
+	// Replace pages in the index with new pages.
 	for k, v := range idx {
-		f.pageCount = max(f.pageCount, k)
-		f.pages[k] = v
+		// Patch avoids mutating the index for an unmodified page.
+		f.pages = f.pages.Patch(k, func(node *pageIndex) (ltx.PageIndexElem, bool) {
+			return v, node == nil || v != node.Value()
+		})
 	}
-	f.txids[info.Level] = max(f.txids[info.Level], info.MaxTXID)
+
+	// Track the MaxTXID for each level.
+	maxTXID := &f.txids[info.Level]
+	*maxTXID = max(*maxTXID, info.MaxTXID)
 	return nil
 }
 
+// Type aliases; these are a mouthful.
+type pageIndex = wbt.Tree[uint32, ltx.PageIndexElem]
 type levelTXIDs = [litestream.SnapshotLevel + 1]ltx.TXID