From e1db3b3ebab9165c9ada1310bd11bd742a851bec Mon Sep 17 00:00:00 2001 From: tamirms Date: Tue, 6 Oct 2026 00:45:31 +0100 Subject: [PATCH] rocksdb: load a file of sorted entries into a column family Store.LoadSorted writes an ascending sequence of entries as one SST file under the store's own loading directory and moves it into the column family with IngestExternalFile. Readers see all of the entries or none, the entries survive a crash once the call returns, and they bypass the WAL and the memtable. The keys must not overlap keys the column family already holds: the file then lands below every other file and never needs compacting, and an overlap fails the load. The store's lifecycle lock is held for the whole load, so Close waits for it, and a read-write open removes the files a stopped process left behind. The per-column-family options are now kept by name, which the file writer needs to give a file the column family's own block size, filter and compression. CFOptions gains CacheIndexAndFilterBlocks, which keeps an SST's index and filter blocks in the block cache instead of on the heap, so a column family made of many loaded files cannot grow the heap with them. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01BiCL19kUrno3Z4SAYgpAfh --- .../internal/rpcv2/rocksdb/load.go | 90 +++++++ .../internal/rpcv2/rocksdb/load_test.go | 242 ++++++++++++++++++ .../internal/rpcv2/rocksdb/rocksdb.go | 22 +- .../internal/rpcv2/rocksdb/tuning.go | 6 + 4 files changed, 354 insertions(+), 6 deletions(-) create mode 100644 cmd/stellar-rpc/internal/rpcv2/rocksdb/load.go create mode 100644 cmd/stellar-rpc/internal/rpcv2/rocksdb/load_test.go diff --git a/cmd/stellar-rpc/internal/rpcv2/rocksdb/load.go b/cmd/stellar-rpc/internal/rpcv2/rocksdb/load.go new file mode 100644 index 000000000..16a85820b --- /dev/null +++ b/cmd/stellar-rpc/internal/rpcv2/rocksdb/load.go @@ -0,0 +1,90 @@ +package rocksdb + +import ( + "errors" + "fmt" + "iter" + "os" + "path/filepath" + + "github.com/linxGnu/grocksdb" +) + +// loadDirName is where LoadSorted builds a file before RocksDB takes it: a +// subdirectory of the store's own directory, which RocksDB does not scan. +const loadDirName = "loading" + +// LoadSorted writes entries, which must ascend by key, as one file and moves +// it into cf. Readers see all of the entries or none, and they survive a +// crash once LoadSorted returns. An empty sequence is an error. +// +// The entries bypass the WAL and the memtable, but other writes to the store +// wait while RocksDB installs the file. The entries' key range must not +// overlap keys cf already holds: the file then lands below every other file +// of cf, and cf never needs compacting. An overlap fails the load. +// +// Each entry is written before the next is pulled, so entries may reuse +// their buffers. The store's lifecycle lock is held throughout, so Close +// waits for the load and entries must not call back into the store. +func (s *Store) LoadSorted(cf string, entries iter.Seq2[[]byte, []byte]) error { + s.mu.RLock() + defer s.mu.RUnlock() + if err := s.checkOpen(); err != nil { + return err + } + if s.cfg.ReadOnly { + return errors.New("rocksdb: LoadSorted on a read-only store") + } + if cf == "" { + cf = defaultCFName + } + cfh, err := s.resolveCF(cf) + if err != nil { + return err + } + if err := s.loadSorted(cf, cfh, entries); err != nil { + return fmt.Errorf("rocksdb: load sorted entries into %q: %w", cf, err) + } + return nil +} + +func (s *Store) loadSorted(cf string, cfh *grocksdb.ColumnFamilyHandle, entries iter.Seq2[[]byte, []byte]) error { + dir := filepath.Join(s.cfg.Path, loadDirName) + if err := os.MkdirAll(dir, dirPerm); err != nil { + return err + } + file, err := os.CreateTemp(dir, "*.sst") + if err != nil { + return err + } + path := file.Name() + // Nothing is left to remove once RocksDB has moved the file. + defer func() { _ = os.Remove(path) }() + if err := file.Close(); err != nil { + return err + } + + env := grocksdb.NewDefaultEnvOptions() + defer env.Destroy() + // The writer copies cf's options, so the file carries cf's block size, + // filter and compression. + w := grocksdb.NewSSTFileWriter(env, s.cfOpts[cf]) + defer w.Destroy() + if err := w.Open(path); err != nil { + return err + } + for key, value := range entries { + if err := w.Put(key, value); err != nil { + return err + } + } + if err := w.Finish(); err != nil { + return err + } + + opts := grocksdb.NewDefaultIngestExternalFileOptions() + defer opts.Destroy() + opts.SetMoveFiles(true) + opts.SetFailIfNotBottommostLevel(true) + return s.db.IngestExternalFileCF(cfh, []string{path}, opts) +} diff --git a/cmd/stellar-rpc/internal/rpcv2/rocksdb/load_test.go b/cmd/stellar-rpc/internal/rpcv2/rocksdb/load_test.go new file mode 100644 index 000000000..62e507607 --- /dev/null +++ b/cmd/stellar-rpc/internal/rpcv2/rocksdb/load_test.go @@ -0,0 +1,242 @@ +package rocksdb + +import ( + "bytes" + "fmt" + "iter" + "os" + "path/filepath" + "strconv" + "testing" + "time" + + "github.com/linxGnu/grocksdb" + "github.com/stretchr/testify/require" +) + +// pairs yields key/value pairs given as alternating strings. +func pairs(kv ...string) iter.Seq2[[]byte, []byte] { + return func(yield func(key, value []byte) bool) { + for i := 0; i < len(kv); i += 2 { + if !yield([]byte(kv[i]), []byte(kv[i+1])) { + return + } + } + } +} + +// numbered yields n ascending keys, each with value. +func numbered(n int, value []byte) iter.Seq2[[]byte, []byte] { + return func(yield func(key, value []byte) bool) { + for i := range n { + if !yield(fmt.Appendf(nil, "key-%08d", i), value) { + return + } + } + } +} + +func requireValue(t *testing.T, s *Store, cf, key, want string) { + t.Helper() + got, found, err := s.Get(cf, []byte(key)) + require.NoError(t, err) + require.True(t, found, "key %q", key) + require.Equal(t, want, string(got)) +} + +func loadFiles(t *testing.T, s *Store) []string { + t.Helper() + files, err := filepath.Glob(filepath.Join(s.cfg.Path, loadDirName, "*")) + require.NoError(t, err) + return files +} + +func uintProperty(t *testing.T, s *Store, cf, name string) uint64 { + t.Helper() + n, err := strconv.ParseUint(s.db.GetPropertyCF(name, s.cfHandles[cf]), 10, 64) + require.NoError(t, err) + return n +} + +func TestStore_LoadSorted_RoundTrip(t *testing.T) { + const loaded, other = "loaded", "other" + cfg := Config{Path: t.TempDir(), ColumnFamilies: []string{loaded, other}, Logger: silentLogger()} + s, err := New(cfg) + require.NoError(t, err) + + require.NoError(t, s.LoadSorted(loaded, pairs("a1", "one", "a2", "two"))) + // A batch to another CF sits in the memtable while a second, disjoint + // file is loaded. + require.NoError(t, s.Batch(func(b *BatchWriter) error { + b.Put(other, []byte("k"), []byte("batched")) + return nil + })) + require.NoError(t, s.LoadSorted(loaded, pairs("b1", "three"))) + require.Error(t, s.LoadSorted(loaded, pairs("a2", "x")), "a key a loaded file holds") + require.Empty(t, loadFiles(t, s)) + + check := func(s *Store) { + requireValue(t, s, loaded, "a1", "one") + requireValue(t, s, loaded, "a2", "two") + requireValue(t, s, loaded, "b1", "three") + requireValue(t, s, other, "k", "batched") + _, found, err := s.Get(other, []byte("a1")) + require.NoError(t, err) + require.False(t, found, "a loaded key stays in its column family") + last, ok, err := s.LastKey(loaded) + require.NoError(t, err) + require.True(t, ok) + require.Equal(t, "b1", string(last)) + } + check(s) + + require.NoError(t, s.Close()) + reopened, err := New(cfg) + require.NoError(t, err) + t.Cleanup(func() { _ = reopened.Close() }) + check(reopened) +} + +func TestStore_LoadSorted_FailedLoadLeavesNothing(t *testing.T) { + s := openTestStore(t, nil) + require.NoError(t, s.Put("", []byte("kept"), []byte("v"))) + + require.Error(t, s.LoadSorted("", pairs("b", "1", "a", "2")), "descending keys") + require.Error(t, s.LoadSorted("", pairs()), "no entries") + require.Error(t, s.LoadSorted("", pairs("kept", "x")), "a key the column family holds") + require.ErrorIs(t, s.LoadSorted("not-configured", pairs("a", "1")), ErrCFNotFound) + + require.Empty(t, loadFiles(t, s)) + for _, key := range []string{"a", "b"} { + _, found, err := s.Get("", []byte(key)) + require.NoError(t, err) + require.False(t, found, "key %q", key) + } + requireValue(t, s, "", "kept", "v") +} + +func TestStore_LoadSorted_ClosedAndReadOnlyStores(t *testing.T) { + path := t.TempDir() + s, err := New(Config{Path: path, Logger: silentLogger()}) + require.NoError(t, err) + require.NoError(t, s.Close()) + require.ErrorIs(t, s.LoadSorted("", pairs("a", "1")), ErrStoreClosed) + + readOnly, err := New(Config{Path: path, Logger: silentLogger(), ReadOnly: true}) + require.NoError(t, err) + t.Cleanup(func() { _ = readOnly.Close() }) + require.Error(t, readOnly.LoadSorted("", pairs("a", "1"))) + require.NoDirExists(t, filepath.Join(path, loadDirName)) +} + +// Close waits for a load in flight, so the load is durable or never started. +func TestStore_LoadSorted_CloseWaitsForLoad(t *testing.T) { + path := t.TempDir() + s, err := New(Config{Path: path, Logger: silentLogger()}) + require.NoError(t, err) + + loading, release := make(chan struct{}), make(chan struct{}) + loadErr := make(chan error, 1) + go func() { + loadErr <- s.LoadSorted("", func(yield func(key, value []byte) bool) { + close(loading) + <-release + yield([]byte("a"), []byte("1")) + }) + }() + <-loading + closed := make(chan struct{}) + go func() { + _ = s.Close() + close(closed) + }() + select { + case <-closed: + t.Fatal("Close returned while a load was in flight") + case <-time.After(50 * time.Millisecond): + } + close(release) + require.NoError(t, <-loadErr) + <-closed + + reopened, err := New(Config{Path: path, Logger: silentLogger(), MustExist: true}) + require.NoError(t, err) + t.Cleanup(func() { _ = reopened.Close() }) + requireValue(t, reopened, "", "a", "1") +} + +// The file is written with the column family's own options: a compressed +// column family gets a compressed file. +func TestStore_LoadSorted_UsesTheColumnFamilyOptions(t *testing.T) { + const raw, zstd = "raw", "zstd" + s, err := New(Config{ + Path: t.TempDir(), + ColumnFamilies: []string{raw, zstd}, + Logger: silentLogger(), + PerCFOptions: map[string]CFOptions{zstd: {Compression: grocksdb.ZSTDCompression}}, + }) + require.NoError(t, err) + t.Cleanup(func() { _ = s.Close() }) + + value := bytes.Repeat([]byte("compressible "), 100) + for _, cf := range []string{raw, zstd} { + require.NoError(t, s.LoadSorted(cf, numbered(1000, value))) + } + const size = "rocksdb.total-sst-files-size" + require.Less(t, 4*uintProperty(t, s, zstd, size), uintProperty(t, s, raw, size)) +} + +// With CacheIndexAndFilterBlocks, an open file's index and filter blocks +// live in the block cache, not on the heap. +func TestStore_CacheIndexAndFilterBlocks_KeepsThemOffTheHeap(t *testing.T) { + const cached, plain = "cached", "plain" + s, err := New(Config{ + Path: t.TempDir(), + ColumnFamilies: []string{cached, plain}, + Logger: silentLogger(), + Tuning: Tuning{BlockCacheMB: 4}, + PerCFOptions: map[string]CFOptions{ + cached: {BloomFilterBitsPerKey: 10, CacheIndexAndFilterBlocks: true}, + plain: {BloomFilterBitsPerKey: 10}, + }, + }) + require.NoError(t, err) + t.Cleanup(func() { _ = s.Close() }) + + for _, cf := range []string{cached, plain} { + require.NoError(t, s.LoadSorted(cf, numbered(20_000, []byte("v")))) + requireValue(t, s, cf, "key-00000007", "v") + } + const heap = "rocksdb.estimate-table-readers-mem" + require.Less(t, 4*uintProperty(t, s, cached, heap), uintProperty(t, s, plain, heap)) +} + +// Files a stopped process was still building are removed by the next open +// that takes the DB lock, and by no other open. +func TestNew_RemovesUnfinishedLoadFiles(t *testing.T) { + path := t.TempDir() + cfg := Config{Path: path, Logger: silentLogger()} + s, err := New(cfg) + require.NoError(t, err) + + // An open that fails on the lock leaves a live store's load alone. + require.NoError(t, s.LoadSorted("", func(yield func(key, value []byte) bool) { + _, err := New(cfg) + require.Error(t, err) + yield([]byte("a"), []byte("1")) + })) + require.NoError(t, s.Close()) + + orphan := filepath.Join(path, loadDirName, "orphan.sst") + require.NoError(t, os.WriteFile(orphan, []byte("half a file"), 0o600)) + readOnly, err := New(Config{Path: path, Logger: silentLogger(), ReadOnly: true}) + require.NoError(t, err) + require.NoError(t, readOnly.Close()) + require.FileExists(t, orphan) + + reopened, err := New(Config{Path: path, Logger: silentLogger(), MustExist: true}) + require.NoError(t, err) + t.Cleanup(func() { _ = reopened.Close() }) + require.NoFileExists(t, orphan) + requireValue(t, reopened, "", "a", "1") +} diff --git a/cmd/stellar-rpc/internal/rpcv2/rocksdb/rocksdb.go b/cmd/stellar-rpc/internal/rpcv2/rocksdb/rocksdb.go index edb3dfde4..87eee7d33 100644 --- a/cmd/stellar-rpc/internal/rpcv2/rocksdb/rocksdb.go +++ b/cmd/stellar-rpc/internal/rpcv2/rocksdb/rocksdb.go @@ -125,7 +125,7 @@ type Store struct { db *grocksdb.DB opts *grocksdb.Options - cfOpts []*grocksdb.Options + cfOpts map[string]*grocksdb.Options cfHandles map[string]*grocksdb.ColumnFamilyHandle ro *grocksdb.ReadOptions wo *grocksdb.WriteOptions @@ -748,15 +748,22 @@ func (s *Store) constructAndOpen() error { return fmt.Errorf("rocksdb: open %s: %w", abs, err) } - cfMap := make(map[string]*grocksdb.ColumnFamilyHandle, len(cfHandles)) - for i, name := range cfNames { - cfMap[name] = cfHandles[i] + if !s.cfg.ReadOnly { + // Files LoadSorted was still building when the last process stopped. + // Only the holder of the DB lock may remove them. + if err := os.RemoveAll(filepath.Join(abs, loadDirName)); err != nil { + s.cfg.Logger.WithError(err).Warnf("rocksdb: removing unfinished load files in %s", abs) + } } s.db = db s.opts = opts - s.cfOpts = cfOpts - s.cfHandles = cfMap + s.cfOpts = make(map[string]*grocksdb.Options, len(cfNames)) + s.cfHandles = make(map[string]*grocksdb.ColumnFamilyHandle, len(cfNames)) + for i, name := range cfNames { + s.cfOpts[name] = cfOpts[i] + s.cfHandles[name] = cfHandles[i] + } s.ro = grocksdb.NewDefaultReadOptions() s.wo = grocksdb.NewDefaultWriteOptions() @@ -876,6 +883,9 @@ func (s *Store) applySharedTableOptions(cfNames []string, cfOpts []*grocksdb.Opt if override.BlockSize > 0 { bbto.SetBlockSize(override.BlockSize) } + if override.CacheIndexAndFilterBlocks { + bbto.SetCacheIndexAndFilterBlocks(true) + } o.SetBlockBasedTableFactory(bbto) s.bbtos = append(s.bbtos, bbto) } diff --git a/cmd/stellar-rpc/internal/rpcv2/rocksdb/tuning.go b/cmd/stellar-rpc/internal/rpcv2/rocksdb/tuning.go index 111bae826..0f720d886 100644 --- a/cmd/stellar-rpc/internal/rpcv2/rocksdb/tuning.go +++ b/cmd/stellar-rpc/internal/rpcv2/rocksdb/tuning.go @@ -53,6 +53,12 @@ type CFOptions struct { // BloomFilterBitsPerKey installs this CF's bloom filter; 0 = none — // right for a CF never probed for keys it may not hold. BloomFilterBitsPerKey int + + // CacheIndexAndFilterBlocks keeps each SST's index and filter blocks in + // the block cache instead of on the heap for as long as the file is + // open. For a CF that accumulates many files, it bounds their memory by + // the cache size. + CacheIndexAndFilterBlocks bool } // Tuning — DB-wide RocksDB knobs shared across every CF of one store. Zero