diff --git a/index/store/boltdb/writer.go b/index/store/boltdb/writer.go index 1cd29a34..4ba1d082 100644 --- a/index/store/boltdb/writer.go +++ b/index/store/boltdb/writer.go @@ -23,6 +23,10 @@ func (w *Writer) NewBatch() store.KVBatch { return store.NewEmulatedBatch(w.store.mo) } +func (w *Writer) NewBatchEx(options store.KVBatchOptions) ([]byte, store.KVBatch, error) { + return make([]byte, options.TotalBytes), w.NewBatch(), nil +} + func (w *Writer) ExecuteBatch(batch store.KVBatch) error { emulatedBatch, ok := batch.(*store.EmulatedBatch) diff --git a/index/store/goleveldb/writer.go b/index/store/goleveldb/writer.go index c5229005..defbea13 100644 --- a/index/store/goleveldb/writer.go +++ b/index/store/goleveldb/writer.go @@ -29,6 +29,10 @@ func (w *Writer) NewBatch() store.KVBatch { return &rv } +func (w *Writer) NewBatchEx(options store.KVBatchOptions) ([]byte, store.KVBatch, error) { + return make([]byte, options.TotalBytes), w.NewBatch(), nil +} + func (w *Writer) ExecuteBatch(b store.KVBatch) error { batch, ok := b.(*Batch) if !ok { diff --git a/index/store/gtreap/writer.go b/index/store/gtreap/writer.go index 4490b158..add8535c 100644 --- a/index/store/gtreap/writer.go +++ b/index/store/gtreap/writer.go @@ -29,6 +29,10 @@ func (w *Writer) NewBatch() store.KVBatch { return store.NewEmulatedBatch(w.s.mo) } +func (w *Writer) NewBatchEx(options store.KVBatchOptions) ([]byte, store.KVBatch, error) { + return make([]byte, options.TotalBytes), w.NewBatch(), nil +} + func (w *Writer) ExecuteBatch(batch store.KVBatch) error { emulatedBatch, ok := batch.(*store.EmulatedBatch) diff --git a/index/store/kvstore.go b/index/store/kvstore.go index 3bdd0c5c..55e3fe63 100644 --- a/index/store/kvstore.go +++ b/index/store/kvstore.go @@ -94,6 +94,14 @@ type KVWriter interface { // NewBatch returns a KVBatch for performaing batch operations on this kvstore NewBatch() KVBatch + // NewBatchEx returns a KVBatch and an associated byte array + // that's pre-sized based on the KVBatchOptions. The caller can + // use the returned byte array for keys and values associated with + // the batch. Once the batch is either executed or closed, the + // associated byte array should no longer be accessed by the + // caller. + NewBatchEx(KVBatchOptions) ([]byte, KVBatch, error) + // ExecuteBatch will execute the KVBatch, the provided KVBatch **MUST** have // been created by the same KVStore (though not necessarily the same KVWriter) // Batch execution is atomic, either all the operations or none will be performed @@ -103,6 +111,27 @@ type KVWriter interface { Close() error } +// KVBatchOptions provides the KVWriter.NewBatchEx() method with batch +// preparation and preallocation information. +type KVBatchOptions struct { + // TotalBytes is the sum of key and value bytes needed by the + // caller for the entire batch. It affects the size of the + // returned byte array of KVWrite.NewBatchEx(). + TotalBytes int + + // NumSets is the number of Set() calls the caller will invoke on + // the KVBatch. + NumSets int + + // NumMerges is the number of Merge() calls the caller will invoke + // on the KVBatch. + NumDeletes int + + // NumMerges is the number of Merge() calls the caller will invoke + // on the KVBatch. + NumMerges int +} + // KVBatch is an abstraction for making multiple KV mutations at once type KVBatch interface { diff --git a/index/store/metrics/writer.go b/index/store/metrics/writer.go index 0f23bfdf..420bf3ec 100644 --- a/index/store/metrics/writer.go +++ b/index/store/metrics/writer.go @@ -23,6 +23,10 @@ func (w *Writer) NewBatch() store.KVBatch { return &Batch{s: w.s, o: w.o.NewBatch()} } +func (w *Writer) NewBatchEx(options store.KVBatchOptions) ([]byte, store.KVBatch, error) { + return make([]byte, options.TotalBytes), w.NewBatch(), nil +} + func (w *Writer) ExecuteBatch(b store.KVBatch) (err error) { batch, ok := b.(*Batch) if !ok { diff --git a/index/store/null/null.go b/index/store/null/null.go index 10730664..90adc06a 100644 --- a/index/store/null/null.go +++ b/index/store/null/null.go @@ -96,6 +96,10 @@ func (w *writer) NewBatch() store.KVBatch { return &batch{} } +func (w *writer) NewBatchEx(options store.KVBatchOptions) ([]byte, store.KVBatch, error) { + return make([]byte, options.TotalBytes), w.NewBatch(), nil +} + func (w *writer) ExecuteBatch(store.KVBatch) error { return nil }