Refactor to use MVCCStore with explicit timestamps

Author: bootjp

adapter/dynamodb.go

@@ -26,13 +26,13 @@ const updateSplitCount = 2
 
 type DynamoDBServer struct {
 	listen           net.Listener
-	store            store.ScanStore
+	store            store.MVCCStore
 	coordinator      kv.Coordinator
 	dynamoTranscoder *dynamodbTranscoder
 	httpServer       *http.Server
 }
 
-func NewDynamoDBServer(listen net.Listener, st store.ScanStore, coordinate *kv.Coordinate) *DynamoDBServer {
+func NewDynamoDBServer(listen net.Listener, st store.MVCCStore, coordinate *kv.Coordinate) *DynamoDBServer {
 	d := &DynamoDBServer{
 		listen:           listen,
 		store:            st,
@@ -114,7 +114,8 @@ func (d *DynamoDBServer) getItem(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, "missing key", http.StatusBadRequest)
 		return
 	}
-	v, err := d.store.Get(r.Context(), []byte(keyAttr.S))
+	readTS := snapshotTS(d.coordinator.Clock())
+	v, err := d.store.GetAt(r.Context(), []byte(keyAttr.S), readTS)
 	if err != nil {
 		if errors.Is(err, store.ErrKeyNotFound) {
 			w.Header().Set("Content-Type", "application/x-amz-json-1.0")
@@ -233,7 +234,8 @@ func (d *DynamoDBServer) validateCondition(ctx context.Context, expr string, nam
 	if expr == "" {
 		return nil
 	}
-	exists, err := d.store.Exists(ctx, key)
+	readTS := snapshotTS(d.coordinator.Clock())
+	exists, err := d.store.ExistsAt(ctx, key, readTS)
 	if err != nil {
 		return errors.WithStack(err)
 	}

adapter/dynamodb_test.go

@@ -12,6 +12,7 @@ import (
 	"github.com/aws/aws-sdk-go-v2/service/dynamodb"
 	"github.com/aws/aws-sdk-go-v2/service/dynamodb/types"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestDynamoDB_PutItem_GetItem(t *testing.T) {
@@ -248,7 +249,7 @@ func TestDynamoDB_TransactWriteItems_Concurrent(t *testing.T) {
 			})
 			assert.NoError(t, err, "Get failed for key1 in goroutine %d", i)
 			value1Attr, ok := out1.Item["value"].(*types.AttributeValueMemberS)
-			assert.True(t, ok, "Type assertion failed for key1 in goroutine %d", i)
+			require.True(t, ok, "Type assertion failed for key1 in goroutine %d", i)
 			assert.Equal(t, value1, value1Attr.Value, "Value mismatch for key1 in goroutine %d", i)
 
 			out2, err := client.GetItem(context.Background(), &dynamodb.GetItemInput{
@@ -259,7 +260,7 @@ func TestDynamoDB_TransactWriteItems_Concurrent(t *testing.T) {
 			})
 			assert.NoError(t, err, "Get failed for key2 in goroutine %d", i)
 			value2Attr, ok := out2.Item["value"].(*types.AttributeValueMemberS)
-			assert.True(t, ok, "Type assertion failed for key2 in goroutine %d", i)
+			require.True(t, ok, "Type assertion failed for key2 in goroutine %d", i)
 			assert.Equal(t, value2, value2Attr.Value, "Value mismatch for key2 in goroutine %d", i)
 		}(i)
 	}

adapter/grpc.go

@@ -22,13 +22,13 @@ type GRPCServer struct {
 	log            *slog.Logger
 	grpcTranscoder *grpcTranscoder
 	coordinator    kv.Coordinator
-	store          store.ScanStore
+	store          store.MVCCStore
 
 	pb.UnimplementedRawKVServer
 	pb.UnimplementedTransactionalKVServer
 }
 
-func NewGRPCServer(store store.ScanStore, coordinate *kv.Coordinate) *GRPCServer {
+func NewGRPCServer(store store.MVCCStore, coordinate *kv.Coordinate) *GRPCServer {
 	return &GRPCServer{
 		log: slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{
 			Level: slog.LevelWarn,
@@ -40,8 +40,13 @@ func NewGRPCServer(store store.ScanStore, coordinate *kv.Coordinate) *GRPCServer
 }
 
 func (r GRPCServer) RawGet(ctx context.Context, req *pb.RawGetRequest) (*pb.RawGetResponse, error) {
+	readTS := req.GetTs()
+	if readTS == 0 {
+		readTS = snapshotTS(r.coordinator.Clock())
+	}
+
 	if r.coordinator.IsLeader() {
-		v, err := r.store.Get(ctx, req.Key)
+		v, err := r.store.GetAt(ctx, req.Key, readTS)
 		if err != nil {
 			switch {
 			case errors.Is(err, store.ErrKeyNotFound):
@@ -93,7 +98,8 @@ func (r GRPCServer) tryLeaderGet(key []byte) ([]byte, error) {
 	defer conn.Close()
 
 	cli := pb.NewRawKVClient(conn)
-	resp, err := cli.RawGet(context.Background(), &pb.RawGetRequest{Key: key})
+	ts := snapshotTS(r.coordinator.Clock())
+	resp, err := cli.RawGet(context.Background(), &pb.RawGetRequest{Key: key, Ts: ts})
 	if err != nil {
 		return nil, errors.WithStack(err)
 	}
@@ -180,7 +186,8 @@ func (r GRPCServer) Get(ctx context.Context, req *pb.GetRequest) (*pb.GetRespons
 		return nil, errors.WithStack(err)
 	}
 
-	v, err := r.store.Get(ctx, req.Key)
+	readTS := snapshotTS(r.coordinator.Clock())
+	v, err := r.store.GetAt(ctx, req.Key, readTS)
 	if err != nil {
 		switch {
 		case errors.Is(err, store.ErrKeyNotFound):
@@ -227,7 +234,8 @@ func (r GRPCServer) Scan(ctx context.Context, req *pb.ScanRequest) (*pb.ScanResp
 			Kv: nil,
 		}, errors.WithStack(err)
 	}
-	res, err := r.store.Scan(ctx, req.StartKey, req.EndKey, limit)
+	readTS := snapshotTS(r.coordinator.Clock())
+	res, err := r.store.ScanAt(ctx, req.StartKey, req.EndKey, limit, readTS)
 	if err != nil {
 		return &pb.ScanResponse{
 			Kv: nil,

adapter/internal.go

@@ -35,18 +35,7 @@ func (i *Internal) Forward(_ context.Context, req *pb.ForwardRequest) (*pb.Forwa
 		return nil, errors.WithStack(ErrNotLeader)
 	}
 
-	// Ensure leader issues start_ts when followers forward txn groups without it.
-	if req.IsTxn {
-		var startTs uint64
-		for _, r := range req.Requests {
-			if r.Ts == 0 {
-				if startTs == 0 {
-					startTs = i.clock.Next()
-				}
-				r.Ts = startTs
-			}
-		}
-	}
+	i.stampTimestamps(req)
 
 	r, err := i.transactionManager.Commit(req.Requests)
 	if err != nil {
@@ -61,3 +50,27 @@ func (i *Internal) Forward(_ context.Context, req *pb.ForwardRequest) (*pb.Forwa
 		CommitIndex: r.CommitIndex,
 	}, nil
 }
+
+func (i *Internal) stampTimestamps(req *pb.ForwardRequest) {
+	if req == nil {
+		return
+	}
+	if req.IsTxn {
+		var startTs uint64
+		for _, r := range req.Requests {
+			if r.Ts == 0 {
+				if startTs == 0 {
+					startTs = i.clock.Next()
+				}
+				r.Ts = startTs
+			}
+		}
+		return
+	}
+
+	for _, r := range req.Requests {
+		if r.Ts == 0 {
+			r.Ts = i.clock.Next()
+		}
+	}
+}

adapter/redis.go

@@ -37,10 +37,9 @@ var argsLen = map[string]int{
 
 type RedisServer struct {
 	listen          net.Listener
-	store           store.ScanStore
+	store           store.MVCCStore
 	coordinator     kv.Coordinator
 	redisTranscoder *redisTranscoder
-	listStore       *store.ListStore
 	// TODO manage membership from raft log
 	leaderRedis map[raft.ServerAddress]string
 
@@ -72,15 +71,12 @@ type redisResult struct {
 	err     error
 }
 
-func store2list(st store.ScanStore) *store.ListStore { return store.NewListStore(st) }
-
-func NewRedisServer(listen net.Listener, store store.ScanStore, coordinate *kv.Coordinate, leaderRedis map[raft.ServerAddress]string) *RedisServer {
+func NewRedisServer(listen net.Listener, store store.MVCCStore, coordinate *kv.Coordinate, leaderRedis map[raft.ServerAddress]string) *RedisServer {
 	r := &RedisServer{
 		listen:          listen,
 		store:           store,
 		coordinator:     coordinate,
 		redisTranscoder: newRedisTranscoder(),
-		listStore:       store2list(store),
 		leaderRedis:     leaderRedis,
 	}
 
@@ -112,6 +108,10 @@ func getConnState(conn redcon.Conn) *connState {
 	return st
 }
 
+func (r *RedisServer) readTS() uint64 {
+	return snapshotTS(r.coordinator.Clock())
+}
+
 func (r *RedisServer) Run() error {
 	err := redcon.Serve(r.listen,
 		func(conn redcon.Conn, cmd redcon.Command) {
@@ -211,7 +211,8 @@ func (r *RedisServer) get(conn redcon.Conn, cmd redcon.Command) {
 	}
 
 	if r.coordinator.IsLeader() {
-		v, err := r.store.Get(context.Background(), cmd.Args[1])
+		readTS := r.readTS()
+		v, err := r.store.GetAt(context.Background(), cmd.Args[1], readTS)
 		if err != nil {
 			switch {
 			case errors.Is(err, store.ErrKeyNotFound):
@@ -276,7 +277,8 @@ func (r *RedisServer) exists(conn redcon.Conn, cmd redcon.Command) {
 		return
 	}
 
-	ok, err := r.store.Exists(context.Background(), cmd.Args[1])
+	readTS := r.readTS()
+	ok, err := r.store.ExistsAt(context.Background(), cmd.Args[1], readTS)
 	if err != nil {
 		conn.WriteError(err.Error())
 		return
@@ -325,7 +327,8 @@ func (r *RedisServer) localKeys(pattern []byte) ([][]byte, error) {
 }
 
 func (r *RedisServer) localKeysExact(pattern []byte) ([][]byte, error) {
-	res, err := r.store.Exists(context.Background(), pattern)
+	readTS := r.readTS()
+	res, err := r.store.ExistsAt(context.Background(), pattern, readTS)
 	if err != nil {
 		return nil, errors.WithStack(err)
 	}
@@ -338,7 +341,8 @@ func (r *RedisServer) localKeysExact(pattern []byte) ([][]byte, error) {
 func (r *RedisServer) localKeysPattern(pattern []byte) ([][]byte, error) {
 	start := r.patternStart(pattern)
 
-	keys, err := r.store.Scan(context.Background(), start, nil, math.MaxInt)
+	readTS := r.readTS()
+	keys, err := r.store.ScanAt(context.Background(), start, nil, math.MaxInt, readTS)
 	if err != nil {
 		return nil, errors.WithStack(err)
 	}
@@ -828,13 +832,24 @@ func clampRange(start, end, length int) (int, int) {
 }
 
 func (r *RedisServer) loadListMeta(ctx context.Context, key []byte) (store.ListMeta, bool, error) {
-	meta, exists, err := r.listStore.LoadMeta(ctx, key)
-	return meta, exists, errors.WithStack(err)
+	readTS := r.readTS()
+	val, err := r.store.GetAt(ctx, store.ListMetaKey(key), readTS)
+	if err != nil {
+		if errors.Is(err, store.ErrKeyNotFound) {
+			return store.ListMeta{}, false, nil
+		}
+		return store.ListMeta{}, false, errors.WithStack(err)
+	}
+	meta, err := store.UnmarshalListMeta(val)
+	if err != nil {
+		return store.ListMeta{}, false, errors.WithStack(err)
+	}
+	return meta, true, nil
 }
 
 func (r *RedisServer) isListKey(ctx context.Context, key []byte) (bool, error) {
-	isList, err := r.listStore.IsList(ctx, key)
-	return isList, errors.WithStack(err)
+	_, exists, err := r.loadListMeta(ctx, key)
+	return exists, err
 }
 
 func (r *RedisServer) buildRPushOps(meta store.ListMeta, key []byte, values [][]byte) ([]*kv.Elem[kv.OP], store.ListMeta, error) {
@@ -895,7 +910,8 @@ func (r *RedisServer) deleteList(ctx context.Context, key []byte) error {
 	start := listItemKey(key, math.MinInt64)
 	end := listItemKey(key, math.MaxInt64)
 
-	kvs, err := r.store.Scan(ctx, start, end, math.MaxInt)
+	readTS := r.readTS()
+	kvs, err := r.store.ScanAt(ctx, start, end, math.MaxInt, readTS)
 	if err != nil {
 		return errors.WithStack(err)
 	}
@@ -926,7 +942,8 @@ func (r *RedisServer) fetchListRange(ctx context.Context, key []byte, meta store
 	startKey := listItemKey(key, startSeq)
 	endKey := listItemKey(key, endSeq+1) // exclusive
 
-	kvs, err := r.store.Scan(ctx, startKey, endKey, int(endIdx-startIdx+1))
+	readTS := r.readTS()
+	kvs, err := r.store.ScanAt(ctx, startKey, endKey, int(endIdx-startIdx+1), readTS)
 	if err != nil {
 		return nil, errors.WithStack(err)
 	}
@@ -1039,7 +1056,8 @@ func (r *RedisServer) tryLeaderGet(key []byte) ([]byte, error) {
 	defer conn.Close()
 
 	cli := pb.NewRawKVClient(conn)
-	resp, err := cli.RawGet(context.Background(), &pb.RawGetRequest{Key: key})
+	ts := r.readTS()
+	resp, err := cli.RawGet(context.Background(), &pb.RawGetRequest{Key: key, Ts: ts})
 	if err != nil {
 		return nil, errors.WithStack(err)
 	}
@@ -1048,8 +1066,9 @@ func (r *RedisServer) tryLeaderGet(key []byte) ([]byte, error) {
 }
 
 func (r *RedisServer) getValue(key []byte) ([]byte, error) {
+	readTS := r.readTS()
 	if r.coordinator.IsLeader() {
-		v, err := r.store.Get(context.Background(), key)
+		v, err := r.store.GetAt(context.Background(), key, readTS)
 		return v, errors.WithStack(err)
 	}
 	return r.tryLeaderGet(key)

adapter/ts.go

@@ -0,0 +1,17 @@
+package adapter
+
+import "github.com/bootjp/elastickv/kv"
+
+// snapshotTS returns a timestamp suitable for snapshot reads without
+// unnecessarily advancing the logical clock. It relies solely on the shared
+// HLC; if none has been issued yet, fall back to MaxUint64 to see latest
+// committed versions irrespective of local clock lag.
+func snapshotTS(clock *kv.HLC) uint64 {
+	if clock == nil {
+		return ^uint64(0)
+	}
+	if cur := clock.Current(); cur != 0 {
+		return cur
+	}
+	return ^uint64(0)
+}

kv/coordinator.go

@@ -34,6 +34,7 @@ type Coordinator interface {
 	Dispatch(reqs *OperationGroup[OP]) (*CoordinateResponse, error)
 	IsLeader() bool
 	RaftLeader() raft.ServerAddress
+	Clock() *HLC
 }
 
 func (c *Coordinate) Dispatch(reqs *OperationGroup[OP]) (*CoordinateResponse, error) {