Skip to content

Commit 6c53bc4

Browse files
julienrbrtjulienrbrt
committed
Optimize metadata writes with batching
1 parent af7e326 commit 6c53bc4

6 files changed

Lines changed: 150 additions & 9 deletions

File tree

‎pkg/store/cached_store.go‎

Lines changed: 26 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -112,16 +112,33 @@ func (cs *CachedStore) startWriteLoop() {
112112
go func() {
113113
defer close(cs.done)
114114
for op := range cs.writeCh {
115-
var err error
116-
if op.isDelete {
117-
err = cs.Store.DeleteMetadata(context.Background(), op.key)
118-
} else {
119-
err = cs.Store.SetMetadata(context.Background(), op.key, op.value)
115+
ops := []asyncWriteOp{op}
116+
drain:
117+
for {
118+
select {
119+
case op := <-cs.writeCh:
120+
ops = append(ops, op)
121+
default:
122+
break drain
123+
}
120124
}
121-
if err != nil {
122-
cs.logger.Error().Err(err).Str("key", op.key).
123-
Bool("delete", op.isDelete).
124-
Msg("async metadata write failed")
125+
126+
var puts []MetadataKV
127+
var deletes []string
128+
for _, o := range ops {
129+
if o.isDelete {
130+
deletes = append(deletes, o.key)
131+
} else {
132+
puts = append(puts, MetadataKV{Key: o.key, Value: o.value})
133+
}
134+
}
135+
136+
if err := cs.Store.BatchMetadata(context.Background(), puts, deletes); err != nil {
137+
for _, o := range ops {
138+
cs.logger.Error().Err(err).Str("key", o.key).
139+
Bool("delete", o.isDelete).
140+
Msg("async metadata batch write failed")
141+
}
125142
}
126143
}
127144
}()

‎pkg/store/store.go‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -190,6 +190,30 @@ func (s *DefaultStore) SetMetadata(ctx context.Context, key string, value []byte
190190
return nil
191191
}
192192

193+
func (s *DefaultStore) BatchMetadata(ctx context.Context, puts []MetadataKV, deletes []string) error {
194+
if len(puts) == 0 && len(deletes) == 0 {
195+
return nil
196+
}
197+
batch, err := s.db.Batch(ctx)
198+
if err != nil {
199+
return fmt.Errorf("failed to create metadata batch: %w", err)
200+
}
201+
for _, kv := range puts {
202+
if err := batch.Put(ctx, ds.NewKey(GetMetaKey(kv.Key)), kv.Value); err != nil {
203+
return fmt.Errorf("failed to batch-put metadata key '%s': %w", kv.Key, err)
204+
}
205+
}
206+
for _, key := range deletes {
207+
if err := batch.Delete(ctx, ds.NewKey(GetMetaKey(key))); err != nil {
208+
return fmt.Errorf("failed to batch-delete metadata key '%s': %w", key, err)
209+
}
210+
}
211+
if err := batch.Commit(ctx); err != nil {
212+
return fmt.Errorf("failed to commit metadata batch: %w", err)
213+
}
214+
return nil
215+
}
216+
193217
// GetMetadata returns values stored for given key with SetMetadata.
194218
func (s *DefaultStore) GetMetadata(ctx context.Context, key string) ([]byte, error) {
195219
data, err := s.db.Get(ctx, ds.NewKey(GetMetaKey(key)))

‎pkg/store/tracing.go‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -211,6 +211,25 @@ func (t *tracedStore) DeleteMetadata(ctx context.Context, key string) error {
211211
return nil
212212
}
213213

214+
func (t *tracedStore) BatchMetadata(ctx context.Context, puts []MetadataKV, deletes []string) error {
215+
ctx, span := t.tracer.Start(ctx, "Store.BatchMetadata",
216+
trace.WithAttributes(
217+
attribute.Int("puts", len(puts)),
218+
attribute.Int("deletes", len(deletes)),
219+
),
220+
)
221+
defer span.End()
222+
223+
err := t.inner.BatchMetadata(ctx, puts, deletes)
224+
if err != nil {
225+
span.RecordError(err)
226+
span.SetStatus(codes.Error, err.Error())
227+
return err
228+
}
229+
230+
return nil
231+
}
232+
214233
func (t *tracedStore) DeleteStateAtHeight(ctx context.Context, height uint64) error {
215234
ctx, span := t.tracer.Start(ctx, "Store.DeleteStateAtHeight",
216235
trace.WithAttributes(attribute.Int64("height", int64(height))),

‎pkg/store/tracing_test.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ type tracingMockStore struct {
2828
getMetadataFn func(ctx context.Context, key string) ([]byte, error)
2929

3030
setMetadataFn func(ctx context.Context, key string, value []byte) error
31+
batchMetadataFn func(ctx context.Context, puts []MetadataKV, deletes []string) error
3132
deleteMetadataFn func(ctx context.Context, key string) error
3233
rollbackFn func(ctx context.Context, height uint64, aggregator bool) error
3334
pruneBlocksFn func(ctx context.Context, height uint64) error
@@ -105,6 +106,13 @@ func (m *tracingMockStore) SetMetadata(ctx context.Context, key string, value []
105106
return nil
106107
}
107108

109+
func (m *tracingMockStore) BatchMetadata(ctx context.Context, puts []MetadataKV, deletes []string) error {
110+
if m.batchMetadataFn != nil {
111+
return m.batchMetadataFn(ctx, puts, deletes)
112+
}
113+
return nil
114+
}
115+
108116
func (m *tracingMockStore) DeleteMetadata(ctx context.Context, key string) error {
109117
if m.deleteMetadataFn != nil {
110118
return m.deleteMetadataFn(ctx, key)

‎pkg/store/types.go‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,12 @@ type Store interface {
5454
NewBatch(ctx context.Context) (Batch, error)
5555
}
5656

57+
// MetadataKV is a key-value pair for batched metadata operations.
58+
type MetadataKV struct {
59+
Key string
60+
Value []byte
61+
}
62+
5763
type Metadata interface {
5864
// SetMetadata saves arbitrary value in the store.
5965
//
@@ -62,6 +68,10 @@ type Metadata interface {
6268

6369
// DeleteMetadata removes a metadata key from the store.
6470
DeleteMetadata(ctx context.Context, key string) error
71+
72+
// BatchMetadata writes and deletes metadata keys in a single Badger
73+
// WriteBatch transaction, reducing contention on the write pipeline.
74+
BatchMetadata(ctx context.Context, puts []MetadataKV, deletes []string) error
6575
}
6676

6777
type Reader interface {

‎test/mocks/store.go‎

Lines changed: 63 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)