fix key adapter test

br/pkg/lightning/backend/local/key_adapter.go

@@ -25,11 +25,9 @@ import (
 type KeyAdapter interface {
 	// Encode encodes the key with its corresponding rowID. It appends the encoded key to dst and returns the
 	// resulting slice. The encoded key is guaranteed to be in ascending order for comparison.
-	// To reuse key's storage for the encoded output, use key[:0] as dst. Otherwise, dst must not overlap key.
 	Encode(dst []byte, key []byte, rowID int64) []byte
 
 	// Decode decodes the original key to dst. It appends the encoded key to dst and returns the resulting slice.
-	// To reuse key's storage for the encoded output, use data[:0] as dst. Otherwise, dst must not overlap data.
 	Decode(dst []byte, data []byte) ([]byte, error)
 
 	// EncodedLen returns the encoded key length.
@@ -81,9 +79,13 @@ func (dupDetectKeyAdapter) Decode(dst []byte, data []byte) ([]byte, error) {
 	if err != nil {
 		return nil, err
 	}
-	if len(dst) == 0 || len(dst)+len(key) <= cap(dst) {
+	if len(dst) == 0 {
 		return key, nil
 	}
+	if len(dst)+len(key) <= cap(dst) {
+		dst = dst[:len(dst)+len(key)]
+		return dst, nil
+	}
 	// New slice is allocated, append key to dst manually.
 	return append(dst, key...), nil
 }

br/pkg/lightning/backend/local/key_adapter_test.go

@@ -20,6 +20,7 @@ import (
 	"math"
 	"sort"
 	"testing"
+	"unsafe"
 
 	"github.com/stretchr/testify/require"
 )
@@ -72,6 +73,7 @@ func TestDupDetectKeyAdapter(t *testing.T) {
 	keyAdapter := dupDetectKeyAdapter{}
 	for _, input := range inputs {
 		result := keyAdapter.Encode(nil, input.key, input.rowID)
+		require.Equal(t, keyAdapter.EncodedLen(input.key), len(result))
 
 		// Decode the result.
 		key, err := keyAdapter.Decode(nil, result)
@@ -111,42 +113,62 @@ func TestDupDetectEncodeDupKey(t *testing.T) {
 	require.NotEqual(t, result1, result2)
 }
 
+func startWithSameMemory(x []byte, y []byte) bool {
+	return cap(x) > 0 && cap(y) > 0 && uintptr(unsafe.Pointer(&x[:cap(x)][0])) == uintptr(unsafe.Pointer(&y[:cap(y)][0]))
+}
+
 func TestEncodeKeyToPreAllocatedBuf(t *testing.T) {
 	t.Parallel()
 
 	keyAdapters := []KeyAdapter{noopKeyAdapter{}, dupDetectKeyAdapter{}}
 	for _, keyAdapter := range keyAdapters {
 		key := randBytes(32)
 		buf := make([]byte, 256)
-		buf2 := keyAdapter.Encode(buf[:0], key, 1)
+		buf2 := keyAdapter.Encode(buf[:4], key, 1)
+		require.True(t, startWithSameMemory(buf, buf2))
 		// Verify the encoded result first.
-		key2, err := keyAdapter.Decode(nil, buf2)
+		key2, err := keyAdapter.Decode(nil, buf2[4:])
 		require.NoError(t, err)
 		require.EqualValues(t, key, key2)
-		// There should be no new slice allocated.
-		// If we change a byte in `buf`, `buf2` can read the new byte.
-		require.EqualValues(t, buf2, buf[:len(buf2)])
-		buf[0]++
-		require.Equal(t, buf2[0], buf[0])
 	}
 }
 
 func TestDecodeKeyToPreAllocatedBuf(t *testing.T) {
 	t.Parallel()
 
+	data := []byte{
+		0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8, 0xff, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0xf7,
+		0x0, 0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8, 0x9, 0xa, 0xb, 0xc, 0xd, 0xe, 0xf,
+	}
+	keyAdapters := []KeyAdapter{noopKeyAdapter{}, dupDetectKeyAdapter{}}
+	for _, keyAdapter := range keyAdapters {
+		key, err := keyAdapter.Decode(nil, data)
+		require.NoError(t, err)
+		buf := make([]byte, 4+len(data))
+		buf2, err := keyAdapter.Decode(buf[:4], data)
+		require.NoError(t, err)
+		require.True(t, startWithSameMemory(buf, buf2))
+		require.Equal(t, key, buf2[4:])
+	}
+}
+
+func TestDecodeKeyDstIsInsufficient(t *testing.T) {
+	t.Parallel()
+
+	data := []byte{
+		0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8, 0xff, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0xf7,
+		0x0, 0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8, 0x9, 0xa, 0xb, 0xc, 0xd, 0xe, 0xf,
+	}
 	keyAdapters := []KeyAdapter{noopKeyAdapter{}, dupDetectKeyAdapter{}}
 	for _, keyAdapter := range keyAdapters {
-		data := []byte{
-			0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8, 0xff, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0xf7,
-			0x0, 0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8, 0x9, 0xa, 0xb, 0xc, 0xd, 0xe, 0xf,
-		}
-		buf := make([]byte, len(data))
-		key, err := keyAdapter.Decode(buf[:0], data)
+		key, err := keyAdapter.Decode(nil, data)
+		require.NoError(t, err)
+		buf := make([]byte, 4, 6)
+		copy(buf, []byte{'a', 'b', 'c', 'd'})
+		buf2, err := keyAdapter.Decode(buf[:4], data)
 		require.NoError(t, err)
-		// There should be no new slice allocated.
-		// If we change a byte in `buf`, `buf2` can read the new byte.
-		require.EqualValues(t, key[:len(buf)], buf)
-		buf[0]++
-		require.Equal(t, key[0], buf[0])
+		require.False(t, startWithSameMemory(buf, buf2))
+		require.EqualValues(t, buf[:4], buf2[:4])
+		require.Equal(t, key, buf2[4:])
 	}
 }