Refactor list meta serialization to binary format

Author: bootjp

adapter/dynamodb_transcoder.go

@@ -101,7 +101,7 @@ func (t *dynamodbTranscoder) valueAttrToOps(key []byte, val attributeValue) (*kv
 			Tail: int64(len(val.L)),
 			Len:  int64(len(val.L)),
 		}
-		b, err := json.Marshal(meta)
+		b, err := store.MarshalListMeta(meta)
 		if err != nil {
 			return nil, errors.WithStack(err)
 		}

adapter/redis.go

@@ -3,7 +3,6 @@ package adapter
 import (
 	"bytes"
 	"context"
-	"encoding/json"
 	"math"
 	"net"
 	"sort"
@@ -742,7 +741,7 @@ func (t *txnContext) buildListElems() ([]*kv.Elem[kv.OP], error) {
 
 		st.meta.Len += int64(len(st.appends))
 		st.meta.Tail = st.meta.Head + st.meta.Len
-		metaBytes, err := json.Marshal(st.meta)
+		metaBytes, err := store.MarshalListMeta(st.meta)
 		if err != nil {
 			return nil, errors.WithStack(err)
 		}
@@ -854,7 +853,7 @@ func (r *RedisServer) buildRPushOps(meta store.ListMeta, key []byte, values [][]
 	meta.Len += int64(len(values))
 	meta.Tail = meta.Head + meta.Len
 
-	b, err := json.Marshal(meta)
+	b, err := store.MarshalListMeta(meta)
 	if err != nil {
 		return nil, meta, errors.WithStack(err)
 	}

store/bolt_store.go

@@ -175,6 +175,22 @@ func (t *boltStoreTxn) Exists(_ context.Context, key []byte) (bool, error) {
 	return t.bucket.Get(key) != nil, nil
 }
 
+func (t *boltStoreTxn) Scan(_ context.Context, start []byte, end []byte, limit int) ([]*KVPair, error) {
+	if limit <= 0 {
+		return nil, nil
+	}
+
+	var res []*KVPair
+	c := t.bucket.Cursor()
+	for k, v := c.Seek(start); k != nil && (end == nil || bytes.Compare(k, end) < 0); k, v = c.Next() {
+		res = append(res, &KVPair{Key: k, Value: v})
+		if len(res) >= limit {
+			break
+		}
+	}
+	return res, nil
+}
+
 func (s *boltStore) Txn(ctx context.Context, fn func(ctx context.Context, txn Txn) error) error {
 	btxn, err := s.bbolt.Begin(true)
 	if err != nil {

store/list_store.go

@@ -3,21 +3,23 @@ package store
 import (
 	"bytes"
 	"context"
-	"encoding/hex"
-	"encoding/json"
+	"encoding/binary"
 	"math"
 
 	"github.com/cockroachdb/errors"
 )
 
 // Wide-column style list storage using per-element keys.
-// Item keys: !lst|itm|<userKey>|%020d
-// Meta key : !lst|meta|<userKey> -> {"h":head,"t":tail,"l":len}
+// Item keys: !lst|itm|<userKey><seq(8-byte sortable binary)>
+// Meta key : !lst|meta|<userKey> -> [Head(8)][Tail(8)][Len(8)]
 
 const (
 	ListMetaPrefix = "!lst|meta|"
 	ListItemPrefix = "!lst|itm|"
-	ListSeqWidth   = 20
+	// limit per scan when deleting to avoid OOM.
+	deleteBatchSize    = 1024
+	listMetaBinarySize = 24
+	scanAdvanceByte    = byte(0x00)
 )
 
 type ListMeta struct {
@@ -43,24 +45,28 @@ func (s *ListStore) IsList(ctx context.Context, key []byte) (bool, error) {
 
 // PutList replaces the entire list.
 func (s *ListStore) PutList(ctx context.Context, key []byte, list []string) error {
-	// delete existing meta/items (best-effort)
-	if err := s.deleteList(ctx, key); err != nil {
-		return err
-	}
-
 	meta := ListMeta{Head: 0, Tail: int64(len(list)), Len: int64(len(list))}
-	metaBytes, err := json.Marshal(meta)
+	metaBytes, err := marshalListMeta(meta)
 	if err != nil {
 		return errors.WithStack(err)
 	}
 
 	return errors.WithStack(s.store.Txn(ctx, func(ctx context.Context, txn Txn) error {
+		scanTxn, ok := txn.(ScanTxn)
+		if !ok {
+			return errors.WithStack(ErrNotSupported)
+		}
+
+		if err := s.deleteListTxn(ctx, scanTxn, key); err != nil {
+			return err
+		}
+
 		for i, v := range list {
-			if err := txn.Put(ctx, ListItemKey(key, int64(i)), []byte(v)); err != nil {
+			if err := scanTxn.Put(ctx, ListItemKey(key, int64(i)), []byte(v)); err != nil {
 				return errors.WithStack(err)
 			}
 		}
-		if err := txn.Put(ctx, ListMetaKey(key), metaBytes); err != nil {
+		if err := scanTxn.Put(ctx, ListMetaKey(key), metaBytes); err != nil {
 			return errors.WithStack(err)
 		}
 		return nil
@@ -107,7 +113,7 @@ func (s *ListStore) RPush(ctx context.Context, key []byte, values ...string) (in
 		}
 		meta.Len += int64(len(values))
 		meta.Tail = meta.Head + meta.Len
-		metaBytes, err := json.Marshal(meta)
+		metaBytes, err := marshalListMeta(meta)
 		if err != nil {
 			return errors.WithStack(err)
 		}
@@ -168,11 +174,8 @@ func (s *ListStore) LoadMeta(ctx context.Context, key []byte) (ListMeta, bool, e
 	if len(val) == 0 {
 		return ListMeta{}, false, nil
 	}
-	var meta ListMeta
-	if err := json.Unmarshal(val, &meta); err != nil {
-		return ListMeta{}, false, errors.WithStack(err)
-	}
-	return meta, true, nil
+	meta, err := unmarshalListMeta(val)
+	return meta, err == nil, errors.WithStack(err)
 }
 
 func (s *ListStore) loadMetaTxn(ctx context.Context, txn Txn, key []byte) (ListMeta, bool, error) {
@@ -186,31 +189,44 @@ func (s *ListStore) loadMetaTxn(ctx context.Context, txn Txn, key []byte) (ListM
 	if len(val) == 0 {
 		return ListMeta{}, false, nil
 	}
-	var meta ListMeta
-	if err := json.Unmarshal(val, &meta); err != nil {
-		return ListMeta{}, false, errors.WithStack(err)
-	}
-	return meta, true, nil
+	meta, err := unmarshalListMeta(val)
+	return meta, err == nil, errors.WithStack(err)
 }
 
-func (s *ListStore) deleteList(ctx context.Context, key []byte) error {
-	start := ListItemKey(key, mathMinInt64) // use smallest seq
-	end := ListItemKey(key, mathMaxInt64)
+// deleteListTxn deletes list items and metadata within the provided transaction.
+func (s *ListStore) deleteListTxn(ctx context.Context, txn ScanTxn, key []byte) error {
+	start := ListItemKey(key, mathMinInt64) // inclusive
+	end := ListItemKey(key, mathMaxInt64)   // inclusive sentinel
 
-	items, err := s.store.Scan(ctx, start, end, math.MaxInt)
-	if err != nil && !errors.Is(err, ErrKeyNotFound) {
-		return errors.WithStack(err)
-	}
+	for {
+		kvs, err := txn.Scan(ctx, start, end, deleteBatchSize)
+		if err != nil && !errors.Is(err, ErrKeyNotFound) {
+			return errors.WithStack(err)
+		}
+		if len(kvs) == 0 {
+			break
+		}
 
-	return errors.WithStack(s.store.Txn(ctx, func(ctx context.Context, txn Txn) error {
-		for _, kvp := range items {
+		for _, kvp := range kvs {
 			if err := txn.Delete(ctx, kvp.Key); err != nil {
 				return errors.WithStack(err)
 			}
 		}
-		_ = txn.Delete(ctx, ListMetaKey(key))
-		return nil
-	}))
+
+		// advance start just after the last processed key to guarantee forward progress
+		lastKey := kvs[len(kvs)-1].Key
+		start = append(bytes.Clone(lastKey), scanAdvanceByte)
+
+		if len(kvs) < deleteBatchSize {
+			break
+		}
+	}
+
+	// delete meta last (ignore missing)
+	if err := txn.Delete(ctx, ListMetaKey(key)); err != nil && !errors.Is(err, ErrKeyNotFound) {
+		return errors.WithStack(err)
+	}
+	return nil
 }
 
 // ListMetaKey builds the metadata key for a user key.
@@ -220,20 +236,55 @@ func ListMetaKey(userKey []byte) []byte {
 
 // ListItemKey builds the item key for a user key and sequence number.
 func ListItemKey(userKey []byte, seq int64) []byte {
-	// Offset sign bit (seq ^ minInt64) to preserve order, then big-endian encode and hex.
+	// Offset sign bit (seq ^ minInt64) to preserve order, then big-endian encode (8 bytes).
 	var raw [8]byte
 	encodeSortableInt64(raw[:], seq)
-	hexSeq := make([]byte, hex.EncodedLen(len(raw)))
-	hex.Encode(hexSeq, raw[:])
 
-	buf := make([]byte, 0, len(ListItemPrefix)+len(userKey)+1+len(hexSeq))
+	buf := make([]byte, 0, len(ListItemPrefix)+len(userKey)+len(raw))
 	buf = append(buf, ListItemPrefix...)
 	buf = append(buf, userKey...)
-	buf = append(buf, '|')
-	buf = append(buf, hexSeq...)
+	buf = append(buf, raw[:]...)
 	return buf
 }
 
+// MarshalListMeta encodes ListMeta into a fixed 24-byte binary format.
+func MarshalListMeta(meta ListMeta) ([]byte, error) { return marshalListMeta(meta) }
+
+// UnmarshalListMeta decodes ListMeta from the fixed 24-byte binary format.
+func UnmarshalListMeta(b []byte) (ListMeta, error) { return unmarshalListMeta(b) }
+
+func marshalListMeta(meta ListMeta) ([]byte, error) {
+	if meta.Head < 0 || meta.Tail < 0 || meta.Len < 0 {
+		return nil, errors.WithStack(errors.Newf("list meta contains negative value: head=%d tail=%d len=%d", meta.Head, meta.Tail, meta.Len))
+	}
+
+	buf := make([]byte, listMetaBinarySize)
+	binary.BigEndian.PutUint64(buf[0:8], uint64(meta.Head))
+	binary.BigEndian.PutUint64(buf[8:16], uint64(meta.Tail))
+	binary.BigEndian.PutUint64(buf[16:24], uint64(meta.Len))
+	return buf, nil
+}
+
+func unmarshalListMeta(b []byte) (ListMeta, error) {
+	if len(b) != listMetaBinarySize {
+		return ListMeta{}, errors.Wrap(errors.Newf("invalid list meta length: %d", len(b)), "unmarshal list meta")
+	}
+
+	head := binary.BigEndian.Uint64(b[0:8])
+	tail := binary.BigEndian.Uint64(b[8:16])
+	length := binary.BigEndian.Uint64(b[16:24])
+
+	if head > math.MaxInt64 || tail > math.MaxInt64 || length > math.MaxInt64 {
+		return ListMeta{}, errors.New("list meta value overflows int64")
+	}
+
+	return ListMeta{
+		Head: int64(head),
+		Tail: int64(tail),
+		Len:  int64(length),
+	}, nil
+}
+
 // encodeSortableInt64 writes seq with sign bit flipped (seq ^ minInt64) in big-endian order.
 const sortableInt64Bytes = 8
 
@@ -286,11 +337,10 @@ func ExtractListUserKey(key []byte) []byte {
 		return bytes.TrimPrefix(key, []byte(ListMetaPrefix))
 	case IsListItemKey(key):
 		trimmed := bytes.TrimPrefix(key, []byte(ListItemPrefix))
-		idx := bytes.LastIndexByte(trimmed, '|')
-		if idx == -1 {
+		if len(trimmed) < sortableInt64Bytes {
 			return nil
 		}
-		return trimmed[:idx]
+		return trimmed[:len(trimmed)-sortableInt64Bytes]
 	default:
 		return nil
 	}

store/list_store_test.go

@@ -3,9 +3,11 @@ package store
 import (
 	"context"
 	"fmt"
+	"io"
 	"sync"
 	"testing"
 
+	"github.com/cockroachdb/errors"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -80,3 +82,71 @@ func TestListStore_RPushConcurrent(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Len(t, list, n)
 }
+
+// failingScanStore simulates a transaction commit failure to verify atomicity.
+type failingScanStore struct {
+	inner ScanStore
+	fail  bool
+}
+
+func newFailingScanStore(inner ScanStore, fail bool) *failingScanStore {
+	return &failingScanStore{inner: inner, fail: fail}
+}
+
+func (s *failingScanStore) Get(ctx context.Context, key []byte) ([]byte, error) {
+	return s.inner.Get(ctx, key)
+}
+
+func (s *failingScanStore) Put(ctx context.Context, key []byte, value []byte) error {
+	return s.inner.Put(ctx, key, value)
+}
+
+func (s *failingScanStore) Delete(ctx context.Context, key []byte) error {
+	return s.inner.Delete(ctx, key)
+}
+
+func (s *failingScanStore) Exists(ctx context.Context, key []byte) (bool, error) {
+	return s.inner.Exists(ctx, key)
+}
+
+func (s *failingScanStore) Snapshot() (io.ReadWriter, error) { return nil, ErrNotSupported }
+func (s *failingScanStore) Restore(io.Reader) error          { return ErrNotSupported }
+func (s *failingScanStore) Close() error                     { return nil }
+
+func (s *failingScanStore) Scan(ctx context.Context, start []byte, end []byte, limit int) ([]*KVPair, error) {
+	return s.inner.Scan(ctx, start, end, limit)
+}
+
+// Txn executes the function; if fail is set, it aborts commit and returns an error.
+func (s *failingScanStore) Txn(ctx context.Context, f func(ctx context.Context, txn Txn) error) error {
+	err := s.inner.Txn(ctx, func(ctx context.Context, txn Txn) error {
+		if s.fail {
+			return errors.New("injected commit failure")
+		}
+		return f(ctx, txn)
+	})
+
+	return err
+}
+
+func TestListStore_PutList_RollbackOnTxnFailure(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	rawBase := NewRbMemoryStore()
+	ls := NewListStore(rawBase)
+
+	initial := []string{"a", "b", "c"}
+	assert.NoError(t, ls.PutList(ctx, []byte("k"), initial))
+
+	failStore := newFailingScanStore(rawBase, true)
+	lsFail := NewListStore(failStore)
+
+	err := lsFail.PutList(ctx, []byte("k"), []string{"x", "y"})
+	assert.Error(t, err, "expected injected failure")
+
+	// Original list must remain intact because txn never committed.
+	out, err := ls.GetList(ctx, []byte("k"))
+	assert.NoError(t, err)
+	assert.Equal(t, initial, out)
+}