Skip to content

Commit e0521d0

Browse files
tamirmsclaude
andcommitted
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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BiCL19kUrno3Z4SAYgpAfh
1 parent 17f2607 commit e0521d0

4 files changed

Lines changed: 354 additions & 6 deletions

File tree

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
package rocksdb
2+
3+
import (
4+
"errors"
5+
"fmt"
6+
"iter"
7+
"os"
8+
"path/filepath"
9+
10+
"github.com/linxGnu/grocksdb"
11+
)
12+
13+
// loadDirName is where LoadSorted builds a file before RocksDB takes it: a
14+
// subdirectory of the store's own directory, which RocksDB does not scan.
15+
const loadDirName = "loading"
16+
17+
// LoadSorted writes entries, which must ascend by key, as one file and moves
18+
// it into cf. Readers see all of the entries or none, and they survive a
19+
// crash once LoadSorted returns. An empty sequence is an error.
20+
//
21+
// The entries bypass the WAL and the memtable, but other writes to the store
22+
// wait while RocksDB installs the file. The entries' key range must not
23+
// overlap keys cf already holds: the file then lands below every other file
24+
// of cf, and cf never needs compacting. An overlap fails the load.
25+
//
26+
// Each entry is written before the next is pulled, so entries may reuse
27+
// their buffers. The store's lifecycle lock is held throughout, so Close
28+
// waits for the load and entries must not call back into the store.
29+
func (s *Store) LoadSorted(cf string, entries iter.Seq2[[]byte, []byte]) error {
30+
s.mu.RLock()
31+
defer s.mu.RUnlock()
32+
if err := s.checkOpen(); err != nil {
33+
return err
34+
}
35+
if s.cfg.ReadOnly {
36+
return errors.New("rocksdb: LoadSorted on a read-only store")
37+
}
38+
if cf == "" {
39+
cf = defaultCFName
40+
}
41+
cfh, err := s.resolveCF(cf)
42+
if err != nil {
43+
return err
44+
}
45+
if err := s.loadSorted(cf, cfh, entries); err != nil {
46+
return fmt.Errorf("rocksdb: load sorted entries into %q: %w", cf, err)
47+
}
48+
return nil
49+
}
50+
51+
func (s *Store) loadSorted(cf string, cfh *grocksdb.ColumnFamilyHandle, entries iter.Seq2[[]byte, []byte]) error {
52+
dir := filepath.Join(s.cfg.Path, loadDirName)
53+
if err := os.MkdirAll(dir, dirPerm); err != nil {
54+
return err
55+
}
56+
file, err := os.CreateTemp(dir, "*.sst")
57+
if err != nil {
58+
return err
59+
}
60+
path := file.Name()
61+
// Nothing is left to remove once RocksDB has moved the file.
62+
defer func() { _ = os.Remove(path) }()
63+
if err := file.Close(); err != nil {
64+
return err
65+
}
66+
67+
env := grocksdb.NewDefaultEnvOptions()
68+
defer env.Destroy()
69+
// The writer copies cf's options, so the file carries cf's block size,
70+
// filter and compression.
71+
w := grocksdb.NewSSTFileWriter(env, s.cfOpts[cf])
72+
defer w.Destroy()
73+
if err := w.Open(path); err != nil {
74+
return err
75+
}
76+
for key, value := range entries {
77+
if err := w.Put(key, value); err != nil {
78+
return err
79+
}
80+
}
81+
if err := w.Finish(); err != nil {
82+
return err
83+
}
84+
85+
opts := grocksdb.NewDefaultIngestExternalFileOptions()
86+
defer opts.Destroy()
87+
opts.SetMoveFiles(true)
88+
opts.SetFailIfNotBottommostLevel(true)
89+
return s.db.IngestExternalFileCF(cfh, []string{path}, opts)
90+
}
Lines changed: 242 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,242 @@
1+
package rocksdb
2+
3+
import (
4+
"bytes"
5+
"fmt"
6+
"iter"
7+
"os"
8+
"path/filepath"
9+
"strconv"
10+
"testing"
11+
"time"
12+
13+
"github.com/linxGnu/grocksdb"
14+
"github.com/stretchr/testify/require"
15+
)
16+
17+
// pairs yields key/value pairs given as alternating strings.
18+
func pairs(kv ...string) iter.Seq2[[]byte, []byte] {
19+
return func(yield func(key, value []byte) bool) {
20+
for i := 0; i < len(kv); i += 2 {
21+
if !yield([]byte(kv[i]), []byte(kv[i+1])) {
22+
return
23+
}
24+
}
25+
}
26+
}
27+
28+
// numbered yields n ascending keys, each with value.
29+
func numbered(n int, value []byte) iter.Seq2[[]byte, []byte] {
30+
return func(yield func(key, value []byte) bool) {
31+
for i := range n {
32+
if !yield(fmt.Appendf(nil, "key-%08d", i), value) {
33+
return
34+
}
35+
}
36+
}
37+
}
38+
39+
func requireValue(t *testing.T, s *Store, cf, key, want string) {
40+
t.Helper()
41+
got, found, err := s.Get(cf, []byte(key))
42+
require.NoError(t, err)
43+
require.True(t, found, "key %q", key)
44+
require.Equal(t, want, string(got))
45+
}
46+
47+
func loadFiles(t *testing.T, s *Store) []string {
48+
t.Helper()
49+
files, err := filepath.Glob(filepath.Join(s.cfg.Path, loadDirName, "*"))
50+
require.NoError(t, err)
51+
return files
52+
}
53+
54+
func uintProperty(t *testing.T, s *Store, cf, name string) uint64 {
55+
t.Helper()
56+
n, err := strconv.ParseUint(s.db.GetPropertyCF(name, s.cfHandles[cf]), 10, 64)
57+
require.NoError(t, err)
58+
return n
59+
}
60+
61+
func TestStore_LoadSorted_RoundTrip(t *testing.T) {
62+
const loaded, other = "loaded", "other"
63+
cfg := Config{Path: t.TempDir(), ColumnFamilies: []string{loaded, other}, Logger: silentLogger()}
64+
s, err := New(cfg)
65+
require.NoError(t, err)
66+
67+
require.NoError(t, s.LoadSorted(loaded, pairs("a1", "one", "a2", "two")))
68+
// A batch to another CF sits in the memtable while a second, disjoint
69+
// file is loaded.
70+
require.NoError(t, s.Batch(func(b *BatchWriter) error {
71+
b.Put(other, []byte("k"), []byte("batched"))
72+
return nil
73+
}))
74+
require.NoError(t, s.LoadSorted(loaded, pairs("b1", "three")))
75+
require.Error(t, s.LoadSorted(loaded, pairs("a2", "x")), "a key a loaded file holds")
76+
require.Empty(t, loadFiles(t, s))
77+
78+
check := func(s *Store) {
79+
requireValue(t, s, loaded, "a1", "one")
80+
requireValue(t, s, loaded, "a2", "two")
81+
requireValue(t, s, loaded, "b1", "three")
82+
requireValue(t, s, other, "k", "batched")
83+
_, found, err := s.Get(other, []byte("a1"))
84+
require.NoError(t, err)
85+
require.False(t, found, "a loaded key stays in its column family")
86+
last, ok, err := s.LastKey(loaded)
87+
require.NoError(t, err)
88+
require.True(t, ok)
89+
require.Equal(t, "b1", string(last))
90+
}
91+
check(s)
92+
93+
require.NoError(t, s.Close())
94+
reopened, err := New(cfg)
95+
require.NoError(t, err)
96+
t.Cleanup(func() { _ = reopened.Close() })
97+
check(reopened)
98+
}
99+
100+
func TestStore_LoadSorted_FailedLoadLeavesNothing(t *testing.T) {
101+
s := openTestStore(t, nil)
102+
require.NoError(t, s.Put("", []byte("kept"), []byte("v")))
103+
104+
require.Error(t, s.LoadSorted("", pairs("b", "1", "a", "2")), "descending keys")
105+
require.Error(t, s.LoadSorted("", pairs()), "no entries")
106+
require.Error(t, s.LoadSorted("", pairs("kept", "x")), "a key the column family holds")
107+
require.ErrorIs(t, s.LoadSorted("not-configured", pairs("a", "1")), ErrCFNotFound)
108+
109+
require.Empty(t, loadFiles(t, s))
110+
for _, key := range []string{"a", "b"} {
111+
_, found, err := s.Get("", []byte(key))
112+
require.NoError(t, err)
113+
require.False(t, found, "key %q", key)
114+
}
115+
requireValue(t, s, "", "kept", "v")
116+
}
117+
118+
func TestStore_LoadSorted_ClosedAndReadOnlyStores(t *testing.T) {
119+
path := t.TempDir()
120+
s, err := New(Config{Path: path, Logger: silentLogger()})
121+
require.NoError(t, err)
122+
require.NoError(t, s.Close())
123+
require.ErrorIs(t, s.LoadSorted("", pairs("a", "1")), ErrStoreClosed)
124+
125+
readOnly, err := New(Config{Path: path, Logger: silentLogger(), ReadOnly: true})
126+
require.NoError(t, err)
127+
t.Cleanup(func() { _ = readOnly.Close() })
128+
require.Error(t, readOnly.LoadSorted("", pairs("a", "1")))
129+
require.NoDirExists(t, filepath.Join(path, loadDirName))
130+
}
131+
132+
// Close waits for a load in flight, so the load is durable or never started.
133+
func TestStore_LoadSorted_CloseWaitsForLoad(t *testing.T) {
134+
path := t.TempDir()
135+
s, err := New(Config{Path: path, Logger: silentLogger()})
136+
require.NoError(t, err)
137+
138+
loading, release := make(chan struct{}), make(chan struct{})
139+
loadErr := make(chan error, 1)
140+
go func() {
141+
loadErr <- s.LoadSorted("", func(yield func(key, value []byte) bool) {
142+
close(loading)
143+
<-release
144+
yield([]byte("a"), []byte("1"))
145+
})
146+
}()
147+
<-loading
148+
closed := make(chan struct{})
149+
go func() {
150+
_ = s.Close()
151+
close(closed)
152+
}()
153+
select {
154+
case <-closed:
155+
t.Fatal("Close returned while a load was in flight")
156+
case <-time.After(50 * time.Millisecond):
157+
}
158+
close(release)
159+
require.NoError(t, <-loadErr)
160+
<-closed
161+
162+
reopened, err := New(Config{Path: path, Logger: silentLogger(), MustExist: true})
163+
require.NoError(t, err)
164+
t.Cleanup(func() { _ = reopened.Close() })
165+
requireValue(t, reopened, "", "a", "1")
166+
}
167+
168+
// The file is written with the column family's own options: a compressed
169+
// column family gets a compressed file.
170+
func TestStore_LoadSorted_UsesTheColumnFamilyOptions(t *testing.T) {
171+
const raw, zstd = "raw", "zstd"
172+
s, err := New(Config{
173+
Path: t.TempDir(),
174+
ColumnFamilies: []string{raw, zstd},
175+
Logger: silentLogger(),
176+
PerCFOptions: map[string]CFOptions{zstd: {Compression: grocksdb.ZSTDCompression}},
177+
})
178+
require.NoError(t, err)
179+
t.Cleanup(func() { _ = s.Close() })
180+
181+
value := bytes.Repeat([]byte("compressible "), 100)
182+
for _, cf := range []string{raw, zstd} {
183+
require.NoError(t, s.LoadSorted(cf, numbered(1000, value)))
184+
}
185+
const size = "rocksdb.total-sst-files-size"
186+
require.Less(t, 4*uintProperty(t, s, zstd, size), uintProperty(t, s, raw, size))
187+
}
188+
189+
// With CacheIndexAndFilterBlocks, an open file's index and filter blocks
190+
// live in the block cache, not on the heap.
191+
func TestStore_CacheIndexAndFilterBlocks_KeepsThemOffTheHeap(t *testing.T) {
192+
const cached, plain = "cached", "plain"
193+
s, err := New(Config{
194+
Path: t.TempDir(),
195+
ColumnFamilies: []string{cached, plain},
196+
Logger: silentLogger(),
197+
Tuning: Tuning{BlockCacheMB: 4},
198+
PerCFOptions: map[string]CFOptions{
199+
cached: {BloomFilterBitsPerKey: 10, CacheIndexAndFilterBlocks: true},
200+
plain: {BloomFilterBitsPerKey: 10},
201+
},
202+
})
203+
require.NoError(t, err)
204+
t.Cleanup(func() { _ = s.Close() })
205+
206+
for _, cf := range []string{cached, plain} {
207+
require.NoError(t, s.LoadSorted(cf, numbered(20_000, []byte("v"))))
208+
requireValue(t, s, cf, "key-00000007", "v")
209+
}
210+
const heap = "rocksdb.estimate-table-readers-mem"
211+
require.Less(t, 4*uintProperty(t, s, cached, heap), uintProperty(t, s, plain, heap))
212+
}
213+
214+
// Files a stopped process was still building are removed by the next open
215+
// that takes the DB lock, and by no other open.
216+
func TestNew_RemovesUnfinishedLoadFiles(t *testing.T) {
217+
path := t.TempDir()
218+
cfg := Config{Path: path, Logger: silentLogger()}
219+
s, err := New(cfg)
220+
require.NoError(t, err)
221+
222+
// An open that fails on the lock leaves a live store's load alone.
223+
require.NoError(t, s.LoadSorted("", func(yield func(key, value []byte) bool) {
224+
_, err := New(cfg)
225+
require.Error(t, err)
226+
yield([]byte("a"), []byte("1"))
227+
}))
228+
require.NoError(t, s.Close())
229+
230+
orphan := filepath.Join(path, loadDirName, "orphan.sst")
231+
require.NoError(t, os.WriteFile(orphan, []byte("half a file"), 0o600))
232+
readOnly, err := New(Config{Path: path, Logger: silentLogger(), ReadOnly: true})
233+
require.NoError(t, err)
234+
require.NoError(t, readOnly.Close())
235+
require.FileExists(t, orphan)
236+
237+
reopened, err := New(Config{Path: path, Logger: silentLogger(), MustExist: true})
238+
require.NoError(t, err)
239+
t.Cleanup(func() { _ = reopened.Close() })
240+
require.NoFileExists(t, orphan)
241+
requireValue(t, reopened, "", "a", "1")
242+
}

‎cmd/stellar-rpc/internal/rpcv2/rocksdb/rocksdb.go‎

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -125,7 +125,7 @@ type Store struct {
125125

126126
db *grocksdb.DB
127127
opts *grocksdb.Options
128-
cfOpts []*grocksdb.Options
128+
cfOpts map[string]*grocksdb.Options
129129
cfHandles map[string]*grocksdb.ColumnFamilyHandle
130130
ro *grocksdb.ReadOptions
131131
wo *grocksdb.WriteOptions
@@ -748,15 +748,22 @@ func (s *Store) constructAndOpen() error {
748748
return fmt.Errorf("rocksdb: open %s: %w", abs, err)
749749
}
750750

751-
cfMap := make(map[string]*grocksdb.ColumnFamilyHandle, len(cfHandles))
752-
for i, name := range cfNames {
753-
cfMap[name] = cfHandles[i]
751+
if !s.cfg.ReadOnly {
752+
// Files LoadSorted was still building when the last process stopped.
753+
// Only the holder of the DB lock may remove them.
754+
if err := os.RemoveAll(filepath.Join(abs, loadDirName)); err != nil {
755+
s.cfg.Logger.WithError(err).Warnf("rocksdb: removing unfinished load files in %s", abs)
756+
}
754757
}
755758

756759
s.db = db
757760
s.opts = opts
758-
s.cfOpts = cfOpts
759-
s.cfHandles = cfMap
761+
s.cfOpts = make(map[string]*grocksdb.Options, len(cfNames))
762+
s.cfHandles = make(map[string]*grocksdb.ColumnFamilyHandle, len(cfNames))
763+
for i, name := range cfNames {
764+
s.cfOpts[name] = cfOpts[i]
765+
s.cfHandles[name] = cfHandles[i]
766+
}
760767
s.ro = grocksdb.NewDefaultReadOptions()
761768
s.wo = grocksdb.NewDefaultWriteOptions()
762769

@@ -876,6 +883,9 @@ func (s *Store) applySharedTableOptions(cfNames []string, cfOpts []*grocksdb.Opt
876883
if override.BlockSize > 0 {
877884
bbto.SetBlockSize(override.BlockSize)
878885
}
886+
if override.CacheIndexAndFilterBlocks {
887+
bbto.SetCacheIndexAndFilterBlocks(true)
888+
}
879889
o.SetBlockBasedTableFactory(bbto)
880890
s.bbtos = append(s.bbtos, bbto)
881891
}

‎cmd/stellar-rpc/internal/rpcv2/rocksdb/tuning.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,12 @@ type CFOptions struct {
5353
// BloomFilterBitsPerKey installs this CF's bloom filter; 0 = none —
5454
// right for a CF never probed for keys it may not hold.
5555
BloomFilterBitsPerKey int
56+
57+
// CacheIndexAndFilterBlocks keeps each SST's index and filter blocks in
58+
// the block cache instead of on the heap for as long as the file is
59+
// open. For a CF that accumulates many files, it bounds their memory by
60+
// the cache size.
61+
CacheIndexAndFilterBlocks bool
5662
}
5763

5864
// Tuning — DB-wide RocksDB knobs shared across every CF of one store. Zero

0 commit comments

Comments
 (0)