Merge branch 'main' into feature/s3-doc

Author: bootjp

adapter/dynamodb.go

@@ -62,10 +62,10 @@ const (
 	batchWriteItemMaxItems      = 25
 	dynamoMaxRequestBodyBytes   = 1 << 20
 
-	dynamoTableMetaPrefix       = "!ddb|meta|table|"
-	dynamoTableGenerationPrefix = "!ddb|meta|gen|"
-	dynamoItemPrefix            = "!ddb|item|"
-	dynamoGSIPrefix             = "!ddb|gsi|"
+	dynamoTableMetaPrefix       = kv.DynamoTableMetaPrefix
+	dynamoTableGenerationPrefix = kv.DynamoTableGenerationPrefix
+	dynamoItemPrefix            = kv.DynamoItemPrefix
+	dynamoGSIPrefix             = kv.DynamoGSIPrefix
 	dynamoScanPageLimit         = 1024
 	dynamoKeyEscapeByte         = byte(0x00)
 	dynamoKeyTerminatorByte     = byte(0x01)

kv/shard_key.go

@@ -11,24 +11,45 @@ const redisInternalRoutePrefix = "!redis|"
 
 var redisInternalRoutePrefixBytes = []byte(redisInternalRoutePrefix)
 
+const (
+	dynamoRoutePrefix = "!ddb|route|table|"
+
+	// DynamoTableMetaPrefix prefixes DynamoDB table metadata keys.
+	DynamoTableMetaPrefix = "!ddb|meta|table|"
+	// DynamoTableGenerationPrefix prefixes DynamoDB table generation keys.
+	DynamoTableGenerationPrefix = "!ddb|meta|gen|"
+	// DynamoItemPrefix prefixes DynamoDB item storage keys.
+	DynamoItemPrefix = "!ddb|item|"
+	// DynamoGSIPrefix prefixes DynamoDB GSI storage keys.
+	DynamoGSIPrefix = "!ddb|gsi|"
+)
+
+var (
+	dynamoRoutePrefixBytes           = []byte(dynamoRoutePrefix)
+	dynamoTableMetaPrefixBytes       = []byte(DynamoTableMetaPrefix)
+	dynamoTableGenerationPrefixBytes = []byte(DynamoTableGenerationPrefix)
+	dynamoItemPrefixBytes            = []byte(DynamoItemPrefix)
+	dynamoGSIPrefixBytes             = []byte(DynamoGSIPrefix)
+)
+
 // routeKey normalizes internal keys (e.g., list metadata/items) to the logical
 // user key used for shard routing.
 func routeKey(key []byte) []byte {
 	if key == nil {
 		return nil
 	}
+	if embedded, ok := txnRouteKey(key); ok {
+		return normalizeRouteKey(embedded)
+	}
+	return normalizeRouteKey(key)
+}
+
+func normalizeRouteKey(key []byte) []byte {
 	if user := redisRouteKey(key); user != nil {
 		return user
 	}
-	if embedded, ok := txnRouteKey(key); ok {
-		if user := s3keys.ExtractRouteKey(embedded); user != nil {
-			return user
-		}
-		// Transaction internal keys embed the logical key after the prefix.
-		if user := store.ExtractListUserKey(embedded); user != nil {
-			return user
-		}
-		return embedded
+	if table := dynamoRouteKey(key); table != nil {
+		return table
 	}
 	if user := s3keys.ExtractRouteKey(key); user != nil {
 		return user
@@ -50,3 +71,36 @@ func redisRouteKey(key []byte) []byte {
 	}
 	return rest[sep+1:]
 }
+
+func dynamoRouteKey(key []byte) []byte {
+	switch {
+	case bytes.HasPrefix(key, dynamoTableMetaPrefixBytes):
+		return dynamoRouteTableKey(key[len(dynamoTableMetaPrefixBytes):])
+	case bytes.HasPrefix(key, dynamoTableGenerationPrefixBytes):
+		return dynamoRouteTableKey(key[len(dynamoTableGenerationPrefixBytes):])
+	case bytes.HasPrefix(key, dynamoItemPrefixBytes):
+		return dynamoRouteFromTablePrefixedKey(key[len(dynamoItemPrefixBytes):])
+	case bytes.HasPrefix(key, dynamoGSIPrefixBytes):
+		return dynamoRouteFromTablePrefixedKey(key[len(dynamoGSIPrefixBytes):])
+	default:
+		return nil
+	}
+}
+
+func dynamoRouteFromTablePrefixedKey(rest []byte) []byte {
+	sep := bytes.IndexByte(rest, '|')
+	if sep <= 0 {
+		return nil
+	}
+	return dynamoRouteTableKey(rest[:sep])
+}
+
+func dynamoRouteTableKey(tableSegment []byte) []byte {
+	if len(tableSegment) == 0 {
+		return nil
+	}
+	out := make([]byte, 0, len(dynamoRoutePrefixBytes)+len(tableSegment))
+	out = append(out, dynamoRoutePrefixBytes...)
+	out = append(out, tableSegment...)
+	return out
+}

kv/shard_key_test.go

@@ -1,6 +1,7 @@
 package kv
 
 import (
+	"encoding/base64"
 	"testing"
 
 	"github.com/bootjp/elastickv/internal/s3keys"
@@ -27,3 +28,22 @@ func TestRouteKey_NormalizesTxnWrappedS3Key(t *testing.T) {
 	embedded := s3keys.UploadPartKey("bucket-a", 7, "path/to/object", "upload-1", 3)
 	require.Equal(t, s3keys.RouteKey("bucket-a", 7, "path/to/object"), routeKey(txnLockKey(embedded)))
 }
+
+func TestRouteKey_NormalizesDynamoKeysToTable(t *testing.T) {
+	t.Parallel()
+
+	tableSegment := []byte(base64.RawURLEncoding.EncodeToString([]byte("users")))
+	indexSegment := base64.RawURLEncoding.EncodeToString([]byte("status-index"))
+	want := dynamoRouteTableKey(tableSegment)
+
+	metaKey := append([]byte(DynamoTableMetaPrefix), tableSegment...)
+	generationKey := append([]byte(DynamoTableGenerationPrefix), tableSegment...)
+	itemKey := append([]byte(DynamoItemPrefix+string(tableSegment)+"|7|"), []byte("pk\x00\x01")...)
+	gsiKey := append([]byte(DynamoGSIPrefix+string(tableSegment)+"|7|"+indexSegment+"|"), []byte("idx\x00\x01pk\x00\x01")...)
+
+	require.Equal(t, want, routeKey(metaKey))
+	require.Equal(t, want, routeKey(generationKey))
+	require.Equal(t, want, routeKey(itemKey))
+	require.Equal(t, want, routeKey(gsiKey))
+	require.Equal(t, want, routeKey(txnLockKey(itemKey)))
+}

kv/sharded_coordinator.go

@@ -105,15 +105,13 @@ func (c *ShardedCoordinator) dispatchTxn(startTS uint64, commitTS uint64, elems
 		return nil, errors.WithStack(ErrTxnPrimaryKeyRequired)
 	}
 
-	if commitTS == 0 {
-		commitTS = c.nextTxnTSAfter(startTS)
-	} else {
-		// Observe caller-provided commitTS to keep the HLC monotonic; without
-		// this the clock could later issue timestamps smaller than commitTS.
-		c.clock.Observe(commitTS)
+	commitTS, err = c.resolveTxnCommitTS(startTS, commitTS)
+	if err != nil {
+		return nil, err
 	}
-	if commitTS == 0 || commitTS <= startTS {
-		return nil, errors.WithStack(ErrTxnCommitTSRequired)
+
+	if len(gids) == 1 {
+		return c.dispatchSingleShardTxn(startTS, commitTS, primaryKey, gids[0], elems)
 	}
 
 	prepared, err := c.prewriteTxn(startTS, commitTS, primaryKey, grouped, gids)
@@ -131,6 +129,37 @@ func (c *ShardedCoordinator) dispatchTxn(startTS uint64, commitTS uint64, elems
 	return &CoordinateResponse{CommitIndex: maxIndex}, nil
 }
 
+func (c *ShardedCoordinator) resolveTxnCommitTS(startTS, commitTS uint64) (uint64, error) {
+	if commitTS == 0 {
+		commitTS = c.nextTxnTSAfter(startTS)
+	} else {
+		// Observe caller-provided commitTS to keep the HLC monotonic; without
+		// this the clock could later issue timestamps smaller than commitTS.
+		c.clock.Observe(commitTS)
+	}
+	if commitTS == 0 || commitTS <= startTS {
+		return 0, errors.WithStack(ErrTxnCommitTSRequired)
+	}
+	return commitTS, nil
+}
+
+func (c *ShardedCoordinator) dispatchSingleShardTxn(startTS, commitTS uint64, primaryKey []byte, gid uint64, elems []*Elem[OP]) (*CoordinateResponse, error) {
+	g, err := c.txnGroupForID(gid)
+	if err != nil {
+		return nil, err
+	}
+	resp, err := g.Txn.Commit([]*pb.Request{
+		onePhaseTxnRequest(startTS, commitTS, primaryKey, elems),
+	})
+	if err != nil {
+		return nil, errors.WithStack(err)
+	}
+	if resp == nil {
+		return &CoordinateResponse{}, nil
+	}
+	return &CoordinateResponse{CommitIndex: resp.CommitIndex}, nil
+}
+
 type preparedGroup struct {
 	gid  uint64
 	keys []*pb.Mutation
@@ -487,16 +516,9 @@ func (c *ShardedCoordinator) txnLogs(reqs *OperationGroup[OP]) ([]*pb.Request, e
 	if len(gids) != 1 {
 		return nil, errors.WithStack(ErrInvalidRequest)
 	}
-	commitTS := reqs.CommitTS
-	if commitTS == 0 {
-		commitTS = c.nextTxnTSAfter(reqs.StartTS)
-	} else {
-		// Observe caller-provided commitTS to keep the HLC monotonic; without
-		// this the clock could later issue timestamps smaller than commitTS.
-		c.clock.Observe(commitTS)
-	}
-	if commitTS == 0 || commitTS <= reqs.StartTS {
-		return nil, errors.WithStack(ErrTxnCommitTSRequired)
+	commitTS, err := c.resolveTxnCommitTS(reqs.StartTS, reqs.CommitTS)
+	if err != nil {
+		return nil, err
 	}
 	return buildTxnLogs(reqs.StartTS, commitTS, grouped, gids)
 }

kv/sharded_coordinator_txn_test.go

@@ -166,6 +166,49 @@ func TestShardedCoordinatorDispatchTxn_CrossShardPhasesAndCommitIndex(t *testing
 	require.Equal(t, commitMeta1.CommitTS, commitMeta2.CommitTS)
 }
 
+func TestShardedCoordinatorDispatchTxn_SingleShardUsesOnePhase(t *testing.T) {
+	t.Parallel()
+
+	engine := distribution.NewEngine()
+	engine.UpdateRoute([]byte(""), nil, 1)
+
+	g1Txn := &recordingTransactional{
+		responses: []*TransactionResponse{
+			{CommitIndex: 17},
+		},
+	}
+
+	coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{
+		1: {Txn: g1Txn},
+	}, 1, NewHLC(), nil)
+
+	startTS := uint64(10)
+	resp, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{
+		IsTxn:   true,
+		StartTS: startTS,
+		Elems: []*Elem[OP]{
+			{Op: Put, Key: []byte("b"), Value: []byte("v1")},
+			{Op: Put, Key: []byte("c"), Value: []byte("v2")},
+		},
+	})
+	require.NoError(t, err)
+	require.NotNil(t, resp)
+	require.Equal(t, uint64(17), resp.CommitIndex)
+
+	require.Len(t, g1Txn.requests, 1)
+	req := g1Txn.requests[0]
+	require.Equal(t, pb.Phase_NONE, req.Phase)
+	require.Equal(t, startTS, req.Ts)
+	require.Len(t, req.Mutations, 3)
+	require.Equal(t, []byte("b"), req.Mutations[1].Key)
+	require.Equal(t, []byte("c"), req.Mutations[2].Key)
+
+	meta := requestTxnMeta(t, req)
+	require.Equal(t, []byte("b"), meta.PrimaryKey)
+	require.Zero(t, meta.LockTTLms)
+	require.Greater(t, meta.CommitTS, startTS)
+}
+
 func TestShardedCoordinatorDispatchTxn_UsesProvidedCommitTS(t *testing.T) {
 	t.Parallel()