From 02c9953b087844bac17df29a5551cf7611c6cccc Mon Sep 17 00:00:00 2001
From: roseduan <roseduan520@gmail.com>
Date: Thu, 3 Aug 2023 21:23:50 +0800
Subject: [PATCH] Optimize: use sync.Pool to avoid create batch every time
 (#245)

---
 batch.go | 31 +++++++++++++++++++++++++++++++
 db.go    | 37 +++++++++++++++++++++++--------------
 2 files changed, 54 insertions(+), 14 deletions(-)

diff --git a/batch.go b/batch.go
index fc8af186..12423e8e 100644
--- a/batch.go
+++ b/batch.go
@@ -50,6 +50,37 @@ func (db *DB) NewBatch(options BatchOptions) *Batch {
 	return batch
 }
 
+func makeBatch() interface{} {
+	node, err := snowflake.NewNode(1)
+	if err != nil {
+		panic(fmt.Sprintf("snowflake.NewNode(1) failed: %v", err))
+	}
+	return &Batch{
+		options: DefaultBatchOptions,
+		batchId: node,
+	}
+}
+
+func (b *Batch) init(rdonly, sync bool, db *DB) *Batch {
+	b.options.ReadOnly = rdonly
+	b.options.Sync = sync
+	b.db = db
+	b.lock()
+	return b
+}
+
+func (b *Batch) withPendingWrites() *Batch {
+	b.pendingWrites = make(map[string]*LogRecord)
+	return b
+}
+
+func (b *Batch) reset() {
+	b.db = nil
+	b.pendingWrites = nil
+	b.committed = false
+	b.rollbacked = false
+}
+
 func (b *Batch) lock() {
 	if b.options.ReadOnly {
 		b.db.mu.RLock()
diff --git a/db.go b/db.go
index 3d5902d3..de797fef 100644
--- a/db.go
+++ b/db.go
@@ -44,6 +44,7 @@ type DB struct {
 	mu           sync.RWMutex
 	closed       bool
 	mergeRunning uint32 // indicate if the database is merging
+	batchPool    sync.Pool
 }
 
 // Stat represents the statistics of the database.
@@ -109,6 +110,7 @@ func Open(options Options) (*DB, error) {
 		index:     index.NewIndexer(),
 		options:   options,
 		fileLock:  fileLock,
+		batchPool: sync.Pool{New: makeBatch},
 	}
 
 	// load index frm hint file
@@ -178,12 +180,15 @@ func (db *DB) Stat() *Stat {
 // Actually, it will open a new batch and commit it.
 // You can think the batch has only one Put operation.
 func (db *DB) Put(key []byte, value []byte) error {
-	options := DefaultBatchOptions
+	batch := db.batchPool.Get().(*Batch)
+	defer func() {
+		batch.reset()
+		db.batchPool.Put(batch)
+	}()
 	// This is a single delete operation, we can set Sync to false.
 	// Because the data will be written to the WAL,
 	// and the WAL file will be synced to disk according to the DB options.
-	options.Sync = false
-	batch := db.NewBatch(options)
+	batch.init(false, false, db).withPendingWrites()
 	if err := batch.Put(key, value); err != nil {
 		return err
 	}
@@ -194,12 +199,12 @@ func (db *DB) Put(key []byte, value []byte) error {
 // Actually, it will open a new batch and commit it.
 // You can think the batch has only one Get operation.
 func (db *DB) Get(key []byte) ([]byte, error) {
-	options := DefaultBatchOptions
-	// Read-only operation
-	options.ReadOnly = true
-	batch := db.NewBatch(options)
+	batch := db.batchPool.Get().(*Batch)
+	batch.init(true, false, db)
 	defer func() {
 		_ = batch.Commit()
+		batch.reset()
+		db.batchPool.Put(batch)
 	}()
 	return batch.Get(key)
 }
@@ -208,13 +213,17 @@ func (db *DB) Get(key []byte) ([]byte, error) {
 // Actually, it will open a new batch and commit it.
 // You can think the batch has only one Delete operation.
 func (db *DB) Delete(key []byte) error {
-	options := DefaultBatchOptions
+	batch := db.batchPool.Get().(*Batch)
+	defer func() {
+		batch.reset()
+		db.batchPool.Put(batch)
+	}()
 	// This is a single delete operation, we can set Sync to false.
 	// Because the data will be written to the WAL,
 	// and the WAL file will be synced to disk according to the DB options.
-	options.Sync = false
-	batch := db.NewBatch(options)
+	batch.init(false, false, db).withPendingWrites()
 	if err := batch.Delete(key); err != nil {
+		_ = batch.Rollback()
 		return err
 	}
 	return batch.Commit()
@@ -224,12 +233,12 @@ func (db *DB) Delete(key []byte) error {
 // Actually, it will open a new batch and commit it.
 // You can think the batch has only one Exist operation.
 func (db *DB) Exist(key []byte) (bool, error) {
-	options := DefaultBatchOptions
-	// Read-only operation
-	options.ReadOnly = true
-	batch := db.NewBatch(options)
+	batch := db.batchPool.Get().(*Batch)
+	batch.init(true, false, db)
 	defer func() {
 		_ = batch.Commit()
+		batch.reset()
+		db.batchPool.Put(batch)
 	}()
 	return batch.Exist(key)
 }