From b5cb99cf148727c27ab35524a4d3bf1f90f1da52 Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Thu, 17 Oct 2024 17:27:06 -0700
Subject: [PATCH] add batch

---
 storage/operation/badgerimpl/dbstore.go |  4 ++++
 storage/operation/badgerimpl/writer.go  |  1 +
 storage/operation/pebbleimpl/dbstore.go |  4 ++++
 storage/operation/pebbleimpl/writer.go  |  1 +
 storage/operations.go                   | 13 +++++++++++++
 5 files changed, 23 insertions(+)

diff --git a/storage/operation/badgerimpl/dbstore.go b/storage/operation/badgerimpl/dbstore.go
index b8460165e32..bbe5e274d4f 100644
--- a/storage/operation/badgerimpl/dbstore.go
+++ b/storage/operation/badgerimpl/dbstore.go
@@ -21,3 +21,7 @@ func (b *dbStore) Reader() storage.Reader {
 func (b *dbStore) WithReaderBatchWriter(fn func(storage.ReaderBatchWriter) error) error {
 	return WithReaderBatchWriter(b.db, fn)
 }
+
+func (b *dbStore) NewBatch() storage.Batch {
+	return NewReaderBatchWriter(b.db)
+}
diff --git a/storage/operation/badgerimpl/writer.go b/storage/operation/badgerimpl/writer.go
index 3837be3917f..8fbee66d17f 100644
--- a/storage/operation/badgerimpl/writer.go
+++ b/storage/operation/badgerimpl/writer.go
@@ -18,6 +18,7 @@ type ReaderBatchWriter struct {
 }
 
 var _ storage.ReaderBatchWriter = (*ReaderBatchWriter)(nil)
+var _ storage.Batch = (*ReaderBatchWriter)(nil)
 
 func (b *ReaderBatchWriter) GlobalReader() storage.Reader {
 	return b.globalReader
diff --git a/storage/operation/pebbleimpl/dbstore.go b/storage/operation/pebbleimpl/dbstore.go
index fcc5b14a06a..c7362681f24 100644
--- a/storage/operation/pebbleimpl/dbstore.go
+++ b/storage/operation/pebbleimpl/dbstore.go
@@ -21,3 +21,7 @@ func (b *dbStore) Reader() storage.Reader {
 func (b *dbStore) WithReaderBatchWriter(fn func(storage.ReaderBatchWriter) error) error {
 	return WithReaderBatchWriter(b.db, fn)
 }
+
+func (b *dbStore) NewBatch() storage.Batch {
+	return NewReaderBatchWriter(b.db)
+}
diff --git a/storage/operation/pebbleimpl/writer.go b/storage/operation/pebbleimpl/writer.go
index ad639223209..a206d05d6b7 100644
--- a/storage/operation/pebbleimpl/writer.go
+++ b/storage/operation/pebbleimpl/writer.go
@@ -15,6 +15,7 @@ type ReaderBatchWriter struct {
 }
 
 var _ storage.ReaderBatchWriter = (*ReaderBatchWriter)(nil)
+var _ storage.Batch = (*ReaderBatchWriter)(nil)
 
 func (b *ReaderBatchWriter) GlobalReader() storage.Reader {
 	return b.globalReader
diff --git a/storage/operations.go b/storage/operations.go
index 95e25fcfa9d..b36d230103e 100644
--- a/storage/operations.go
+++ b/storage/operations.go
@@ -104,6 +104,19 @@ type DB interface {
 	// atomic batch updates to the database.
 	// Any error returned are considered fatal and the batch is not committed.
 	WithReaderBatchWriter(func(ReaderBatchWriter) error) error
+
+	// NewBatch create a new batch for writing.
+	NewBatch() Batch
+}
+
+// Batch is an interface for a batch of writes to a storage backend.
+// The batch is pending until it is committed.
+// Useful for dynamically adding writes to the batch
+type Batch interface {
+	ReaderBatchWriter
+
+	// Commit applies the batched updates to the database.
+	Commit() error
 }
 
 // OnlyWriter is an adapter to convert a function that takes a Writer