From 3ed248c4eba0b80c26d96916ae4af4d3e52673f7 Mon Sep 17 00:00:00 2001
From: Benjamin Wang <wachao@vmware.com>
Date: Tue, 27 Jun 2023 16:48:51 +0100
Subject: [PATCH] Improve CreateBucketIfNotExists to avoid double searching the
 same key

Benchmark with this change:
BenchmarkBucket_CreateBucketIfNotExists-10    	     123	   9573035 ns/op	   17930 B/op	      37 allocs/op

Benchmark with old implementnation:
BenchmarkBucket_CreateBucketIfNotExists-10    	     121	  10474415 ns/op	   18147 B/op	      46 allocs/op

Signed-off-by: Benjamin Wang <wachao@vmware.com>
---
 bucket.go      | 54 ++++++++++++++++++++++++++++++++++++++++++++------
 bucket_test.go | 28 ++++++++++++++++++++++++++
 2 files changed, 76 insertions(+), 6 deletions(-)

diff --git a/bucket.go b/bucket.go
index 6a3c48ae8..f8741dcb6 100644
--- a/bucket.go
+++ b/bucket.go
@@ -190,13 +190,55 @@ func (b *Bucket) CreateBucket(key []byte) (*Bucket, error) {
 // Returns an error if the bucket name is blank, or if the bucket name is too long.
 // The bucket instance is only valid for the lifetime of the transaction.
 func (b *Bucket) CreateBucketIfNotExists(key []byte) (*Bucket, error) {
-	child, err := b.CreateBucket(key)
-	if err == errors.ErrBucketExists {
-		return b.Bucket(key), nil
-	} else if err != nil {
-		return nil, err
+	if b.tx.db == nil {
+		return nil, errors.ErrTxClosed
+	} else if !b.tx.writable {
+		return nil, errors.ErrTxNotWritable
+	} else if len(key) == 0 {
+		return nil, errors.ErrBucketNameRequired
+	}
+
+	if b.buckets != nil {
+		if child := b.buckets[string(key)]; child != nil {
+			return child, nil
+		}
+	}
+
+	// Move cursor to correct position.
+	c := b.Cursor()
+	k, v, flags := c.seek(key)
+
+	// Return an error if there is an existing non-bucket key.
+	if bytes.Equal(key, k) {
+		if (flags & common.BucketLeafFlag) != 0 {
+			var child = b.openBucket(v)
+			if b.buckets != nil {
+				b.buckets[string(key)] = child
+			}
+
+			return child, nil
+		}
+		return nil, errors.ErrIncompatibleValue
+	}
+
+	// Create empty, inline bucket.
+	var bucket = Bucket{
+		InBucket:    &common.InBucket{},
+		rootNode:    &node{isLeaf: true},
+		FillPercent: DefaultFillPercent,
 	}
-	return child, nil
+	var value = bucket.write()
+
+	// Insert into node.
+	key = cloneBytes(key)
+	c.node().put(key, key, value, 0, common.BucketLeafFlag)
+
+	// Since subbuckets are not allowed on inline buckets, we need to
+	// dereference the inline page, if it exists. This will cause the bucket
+	// to be treated as a regular, non-inline bucket for the rest of the tx.
+	b.page = nil
+
+	return b.Bucket(key), nil
 }
 
 // DeleteBucket deletes a bucket at the given key.
diff --git a/bucket_test.go b/bucket_test.go
index 21becb330..b60a1b912 100644
--- a/bucket_test.go
+++ b/bucket_test.go
@@ -1883,6 +1883,34 @@ func TestBucket_Delete_Quick(t *testing.T) {
 	}
 }
 
+func BenchmarkBucket_CreateBucketIfNotExists(b *testing.B) {
+	db := btesting.MustCreateDB(b)
+	defer db.MustClose()
+
+	const bucketCount = 1_000_000
+
+	err := db.Update(func(tx *bolt.Tx) error {
+		for i := 0; i < bucketCount; i++ {
+			bucketName := fmt.Sprintf("bucket_%d", i)
+			_, berr := tx.CreateBucket([]byte(bucketName))
+			require.NoError(b, berr)
+		}
+		return nil
+	})
+	require.NoError(b, err)
+
+	b.ResetTimer()
+	b.ReportAllocs()
+
+	for i := 0; i < b.N; i++ {
+		err := db.Update(func(tx *bolt.Tx) error {
+			_, berr := tx.CreateBucketIfNotExists([]byte("bucket_100"))
+			return berr
+		})
+		require.NoError(b, err)
+	}
+}
+
 func ExampleBucket_Put() {
 	// Open the database.
 	db, err := bolt.Open(tempfile(), 0600, nil)