Fix cursor automatically moving to next position when delete and put are in the same TXN

cursor.go

@@ -151,6 +151,20 @@ func (c *Cursor) Delete() error {
 	}
 	c.node().del(key)
 
+	pos := c.curPos()
+	// If the deleted item is the first item in the node, we should keep
+	// the cursor's position (decrease by 1). Otherwise, the cursor will
+	// automatically move to the next item, and when clients call `Next`
+	// method afterward, it will skip one item in the node.
+	//
+	// Note if there isn't any write on current page in current transaction,
+	// then the cursor only iterates items against the page instead of the
+	// in-memory node. Since the page is immutable, so we don't need to change
+	// pos.index in such case.
+	if pos != nil && pos.index == 0 && pos.node != nil {
+		pos.index = -1
+	}
+
 	return nil
 }
 
@@ -377,6 +391,14 @@ func (c *Cursor) keyValue() ([]byte, []byte, uint32) {
 	return elem.Key(), elem.Value(), elem.Flags()
 }
 
+// curPos returns current position of the cursor.
+func (c *Cursor) curPos() *elemRef {
+	if len(c.stack) == 0 {
+		return nil
+	}
+	return &c.stack[len(c.stack)-1]
+}
+
 // node returns the node that the cursor is currently positioned on.
 func (c *Cursor) node() *node {
 	common.Assert(len(c.stack) > 0, "accessing a node with a zero-length cursor stack")

cursor_test.go

@@ -11,6 +11,9 @@ import (
 	"testing"
 	"testing/quick"
 
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
 	bolt "go.etcd.io/bbolt"
 	"go.etcd.io/bbolt/errors"
 	"go.etcd.io/bbolt/internal/btesting"
@@ -742,6 +745,95 @@ func TestCursor_QuickCheck_BucketsOnly_Reverse(t *testing.T) {
 	}
 }
 
+func TestCursor_DeleteAndPut_InTransactions_Forward(t *testing.T) {
+	testCursorWithBothDeleteAndPut(t,
+		func(c *bolt.Cursor) ([]byte, []byte) {
+			return c.First()
+		}, func(c *bolt.Cursor) ([]byte, []byte) {
+			return c.Next()
+		})
+}
+
+func TestCursor_DeleteAndPut_InTransactions_Reverse(t *testing.T) {
+	testCursorWithBothDeleteAndPut(t,
+		func(c *bolt.Cursor) ([]byte, []byte) {
+			return c.Last()
+		}, func(c *bolt.Cursor) ([]byte, []byte) {
+			return c.Prev()
+		})
+}
+
+func testCursorWithBothDeleteAndPut(t *testing.T, initPos func(*bolt.Cursor) ([]byte, []byte), nextPos func(*bolt.Cursor) ([]byte, []byte)) {
+	testCases := []struct {
+		name        string
+		separateTxn bool
+	}{
+		{
+			name:        "put and delete operations in one transaction",
+			separateTxn: false,
+		},
+		{
+			name:        "put and delete operations in separate transactions",
+			separateTxn: true,
+		},
+	}
+
+	const count = 10
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			db := btesting.MustCreateDB(t)
+
+			err := db.Update(func(tx *bolt.Tx) error {
+				b, _ := tx.CreateBucket([]byte("widgets"))
+				for i := 0; i < count; i++ {
+					k := make([]byte, 8)
+					binary.BigEndian.PutUint64(k, uint64(i))
+					if perr := b.Put(k, make([]byte, 100)); perr != nil {
+						return perr
+					}
+				}
+
+				if !tc.separateTxn {
+					c := b.Cursor()
+					for key, _ := initPos(c); key != nil; key, _ = nextPos(c) {
+						if derr := c.Delete(); derr != nil {
+							return derr
+						}
+					}
+				}
+
+				return nil
+			})
+			require.NoError(t, err)
+
+			if tc.separateTxn {
+				err := db.Update(func(tx *bolt.Tx) error {
+					b := tx.Bucket([]byte("widgets"))
+
+					c := b.Cursor()
+					for key, _ := initPos(c); key != nil; key, _ = nextPos(c) {
+						if err := c.Delete(); err != nil {
+							return err
+						}
+					}
+
+					return nil
+				})
+				require.NoError(t, err)
+			}
+
+			err = db.View(func(tx *bolt.Tx) error {
+				b := tx.Bucket([]byte("widgets"))
+				assert.Equal(t, 0, b.Stats().KeyN)
+				return nil
+			})
+			require.NoError(t, err)
+		})
+	}
+}
+
 func ExampleCursor() {
 	// Open the database.
 	db, err := bolt.Open(tempfile(), 0600, nil)