Merge pull request #499 from bootjp/feat/read-skew-fix

store: add read-set validation to ApplyMutations for SSI

Author: bootjp

adapter/distribution_server_test.go

@@ -664,7 +664,7 @@ func (s *distributionCoordinatorStub) applyDispatch(
 	startTS uint64,
 	commitTS uint64,
 ) error {
-	if err := s.store.ApplyMutations(ctx, mutations, startTS, commitTS); err != nil {
+	if err := s.store.ApplyMutations(ctx, mutations, nil, startTS, commitTS); err != nil {
 		return err
 	}
 	if s.afterDispatch != nil {

adapter/redis.go

@@ -1923,7 +1923,11 @@ func (t *txnContext) commit() error {
 		return nil
 	}
 
-	group := &kv.OperationGroup[kv.OP]{IsTxn: true, Elems: elems, StartTS: t.startTS}
+	readKeys := make([][]byte, 0, len(t.readKeys))
+	for _, k := range t.readKeys {
+		readKeys = append(readKeys, k)
+	}
+	group := &kv.OperationGroup[kv.OP]{IsTxn: true, Elems: elems, StartTS: t.startTS, ReadKeys: readKeys}
 	ctx, cancel := context.WithTimeout(context.Background(), redisDispatchTimeout)
 	defer cancel()
 	if _, err := t.server.coordinator.Dispatch(ctx, group); err != nil {

adapter/s3_test.go

@@ -1916,7 +1916,7 @@ func TestS3Server_BackwardCompatibility_NoBucketAclFieldIsPrivate(t *testing.T)
 	commitTS := coord.Clock().Next()
 	err = st.ApplyMutations(context.Background(), []*store.KVPairMutation{
 		{Op: store.OpTypePut, Key: s3keys.BucketMetaKey("legacy-bucket"), Value: legacyJSON},
-	}, commitTS-1, commitTS)
+	}, nil, commitTS-1, commitTS)
 	require.NoError(t, err)
 
 	// Create a server WITH credentials; the legacy bucket has no acl field.

distribution/catalog.go

@@ -627,7 +627,7 @@ func (s *CatalogStore) applySaveMutations(ctx context.Context, plan savePlan, mu
 	if err != nil {
 		return err
 	}
-	if err := s.store.ApplyMutations(ctx, mutations, plan.readTS, commitTS); err != nil {
+	if err := s.store.ApplyMutations(ctx, mutations, nil, plan.readTS, commitTS); err != nil {
 		if errors.Is(err, store.ErrWriteConflict) {
 			return errors.WithStack(ErrCatalogVersionMismatch)
 		}

kv/coordinator.go

@@ -143,8 +143,13 @@ func (c *Coordinate) dispatchTxn(reqs []*Elem[OP], startTS uint64, commitTS uint
 		return nil, errors.WithStack(ErrTxnCommitTSRequired)
 	}
 
+	// Read-set validation for single-shard transactions is performed by the
+	// adapter BEFORE Raft submission (validateReadSet). Passing readKeys
+	// into the Raft log would cause the FSM to reject transactions after
+	// they are already committed in the log, forcing retries at a later
+	// timestamp and breaking realtime ordering of appends.
 	r, err := c.transactionManager.Commit([]*pb.Request{
-		onePhaseTxnRequest(startTS, commitTS, primary, reqs),
+		onePhaseTxnRequest(startTS, commitTS, primary, reqs, nil),
 	})
 	if err != nil {
 		return nil, errors.WithStack(err)
@@ -258,7 +263,7 @@ func (c *Coordinate) redirect(ctx context.Context, reqs *OperationGroup[OP]) (*C
 			commitTS = 0
 		}
 		requests = []*pb.Request{
-			onePhaseTxnRequest(reqs.StartTS, commitTS, primary, reqs.Elems),
+			onePhaseTxnRequest(reqs.StartTS, commitTS, primary, reqs.Elems, nil),
 		}
 	} else {
 		for _, req := range reqs.Elems {
@@ -318,7 +323,7 @@ func elemToMutation(req *Elem[OP]) *pb.Mutation {
 	panic("unreachable")
 }
 
-func onePhaseTxnRequest(startTS, commitTS uint64, primaryKey []byte, reqs []*Elem[OP]) *pb.Request {
+func onePhaseTxnRequest(startTS, commitTS uint64, primaryKey []byte, reqs []*Elem[OP], readKeys [][]byte) *pb.Request {
 	muts := make([]*pb.Mutation, 0, len(reqs)+1)
 	muts = append(muts, txnMetaMutation(primaryKey, 0, commitTS))
 	for _, req := range reqs {
@@ -329,6 +334,7 @@ func onePhaseTxnRequest(startTS, commitTS uint64, primaryKey []byte, reqs []*Ele
 		Phase:     pb.Phase_NONE,
 		Ts:        startTS,
 		Mutations: muts,
+		ReadKeys:  readKeys,
 	}
 }
 

kv/coordinator_txn_test.go

@@ -116,3 +116,6 @@ func TestCoordinateDispatchTxn_UsesProvidedCommitTS(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, commitTS, meta.CommitTS)
 }
+
+// ReadKeys omission from single-shard Raft entries is tested in
+// TestShardedCoordinatorDispatchTxn_SingleShardOmitsReadKeysFromRaftEntry.

kv/fsm.go

@@ -184,7 +184,7 @@ func (f *kvFSM) handleRawRequest(ctx context.Context, r *pb.Request, commitTS ui
 	}
 	// Raw requests always commit against the latest state; use commitTS as both
 	// the validation snapshot and the commit timestamp.
-	return errors.WithStack(f.store.ApplyMutations(ctx, muts, commitTS, commitTS))
+	return errors.WithStack(f.store.ApplyMutations(ctx, muts, nil, commitTS, commitTS))
 }
 
 // extractDelPrefix checks if the mutations contain a DEL_PREFIX operation.
@@ -313,18 +313,16 @@ func (f *kvFSM) handlePrepareRequest(ctx context.Context, r *pb.Request) error {
 		return err
 	}
 
-	if err := f.store.ApplyMutations(ctx, storeMuts, startTS, startTS); err != nil {
+	if err := f.store.ApplyMutations(ctx, storeMuts, r.ReadKeys, startTS, startTS); err != nil {
 		return errors.WithStack(err)
 	}
 	return nil
 }
 
 // handleOnePhaseTxnRequest applies a single-shard transaction atomically.
-// The isolation level is Snapshot Isolation (SI): only write-write conflicts
-// are detected via ApplyMutations. Read-write conflicts (write skew) are NOT
-// prevented because the read-set is not tracked. Callers requiring
-// Serializable Snapshot Isolation (SSI) must implement read-set validation
-// at a higher layer.
+// Both write-write and read-write conflicts are checked: the read set carried
+// in r.ReadKeys is validated alongside the mutation keys inside
+// ApplyMutations under the store's apply lock.
 func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, commitTS uint64) error {
 	meta, muts, err := extractTxnMeta(r.Mutations)
 	if err != nil {
@@ -350,7 +348,7 @@ func (f *kvFSM) handleOnePhaseTxnRequest(ctx context.Context, r *pb.Request, com
 	if err != nil {
 		return err
 	}
-	return errors.WithStack(f.store.ApplyMutations(ctx, storeMuts, startTS, commitTS))
+	return errors.WithStack(f.store.ApplyMutations(ctx, storeMuts, r.ReadKeys, startTS, commitTS))
 }
 
 func (f *kvFSM) handleCommitRequest(ctx context.Context, r *pb.Request) error {
@@ -419,7 +417,7 @@ func (f *kvFSM) commitApplyStartTS(ctx context.Context, primaryKey []byte, start
 // The secondary-shard LatestCommitTS scan is intentionally deferred to the
 // write-conflict path so the hot (first-time) commit path pays no extra cost.
 func (f *kvFSM) applyCommitWithIdempotencyFallback(ctx context.Context, storeMuts []*store.KVPairMutation, uniq []*pb.Mutation, applyStartTS, commitTS uint64) error {
-	err := f.store.ApplyMutations(ctx, storeMuts, applyStartTS, commitTS)
+	err := f.store.ApplyMutations(ctx, storeMuts, nil, applyStartTS, commitTS)
 	if err == nil {
 		return nil
 	}
@@ -436,7 +434,7 @@ func (f *kvFSM) applyCommitWithIdempotencyFallback(ctx context.Context, storeMut
 			return errors.WithStack(lErr)
 		}
 		if exists && latestTS >= commitTS {
-			return errors.WithStack(f.store.ApplyMutations(ctx, storeMuts, commitTS, commitTS))
+			return errors.WithStack(f.store.ApplyMutations(ctx, storeMuts, nil, commitTS, commitTS))
 		}
 	}
 	return errors.WithStack(err)
@@ -475,7 +473,7 @@ func (f *kvFSM) handleAbortRequest(ctx context.Context, r *pb.Request, abortTS u
 	if len(storeMuts) == 0 {
 		return nil
 	}
-	return errors.WithStack(f.store.ApplyMutations(ctx, storeMuts, startTS, abortTS))
+	return errors.WithStack(f.store.ApplyMutations(ctx, storeMuts, nil, startTS, abortTS))
 }
 
 func (f *kvFSM) buildPrepareStoreMutations(ctx context.Context, muts []*pb.Mutation, primaryKey []byte, startTS, expireAt uint64) ([]*store.KVPairMutation, error) {

kv/leader_routed_store.go

@@ -232,11 +232,11 @@ func (s *LeaderRoutedStore) LatestCommitTS(ctx context.Context, key []byte) (uin
 	return s.proxyRawLatestCommitTS(ctx, key)
 }
 
-func (s *LeaderRoutedStore) ApplyMutations(ctx context.Context, mutations []*store.KVPairMutation, startTS, commitTS uint64) error {
+func (s *LeaderRoutedStore) ApplyMutations(ctx context.Context, mutations []*store.KVPairMutation, readKeys [][]byte, startTS, commitTS uint64) error {
 	if s == nil || s.local == nil {
 		return errors.WithStack(store.ErrNotSupported)
 	}
-	return errors.WithStack(s.local.ApplyMutations(ctx, mutations, startTS, commitTS))
+	return errors.WithStack(s.local.ApplyMutations(ctx, mutations, readKeys, startTS, commitTS))
 }
 
 func (s *LeaderRoutedStore) DeletePrefixAt(ctx context.Context, prefix []byte, excludePrefix []byte, commitTS uint64) error {

kv/shard_store.go

@@ -1116,7 +1116,7 @@ func cleanupTSWithNow(startTS, now uint64) uint64 {
 //
 // All mutations must belong to the same shard. Cross-shard mutation batches are
 // not supported.
-func (s *ShardStore) ApplyMutations(ctx context.Context, mutations []*store.KVPairMutation, startTS, commitTS uint64) error {
+func (s *ShardStore) ApplyMutations(ctx context.Context, mutations []*store.KVPairMutation, readKeys [][]byte, startTS, commitTS uint64) error {
 	if len(mutations) == 0 {
 		return nil
 	}
@@ -1135,7 +1135,7 @@ func (s *ShardStore) ApplyMutations(ctx context.Context, mutations []*store.KVPa
 			return errors.WithStack(ErrCrossShardMutationBatchNotSupported)
 		}
 	}
-	return errors.WithStack(firstGroup.Store.ApplyMutations(ctx, mutations, startTS, commitTS))
+	return errors.WithStack(firstGroup.Store.ApplyMutations(ctx, mutations, readKeys, startTS, commitTS))
 }
 
 // DeletePrefixAt applies a prefix delete to every shard in the store.

kv/sharded_coordinator.go

@@ -90,7 +90,7 @@ func (c *ShardedCoordinator) Dispatch(ctx context.Context, reqs *OperationGroup[
 	}
 
 	if reqs.IsTxn {
-		return c.dispatchTxn(reqs.StartTS, reqs.CommitTS, reqs.Elems)
+		return c.dispatchTxn(ctx, reqs.StartTS, reqs.CommitTS, reqs.Elems, reqs.ReadKeys)
 	}
 
 	logs, err := c.requestLogs(reqs)
@@ -193,7 +193,7 @@ func (c *ShardedCoordinator) broadcastToAllGroups(requests []*pb.Request) (*Coor
 	return &CoordinateResponse{CommitIndex: maxIndex.Load()}, nil
 }
 
-func (c *ShardedCoordinator) dispatchTxn(startTS uint64, commitTS uint64, elems []*Elem[OP]) (*CoordinateResponse, error) {
+func (c *ShardedCoordinator) dispatchTxn(ctx context.Context, startTS uint64, commitTS uint64, elems []*Elem[OP], readKeys [][]byte) (*CoordinateResponse, error) {
 	grouped, gids, err := c.groupMutations(elems)
 	if err != nil {
 		return nil, err
@@ -212,7 +212,7 @@ func (c *ShardedCoordinator) dispatchTxn(startTS uint64, commitTS uint64, elems
 		return c.dispatchSingleShardTxn(startTS, commitTS, primaryKey, gids[0], elems)
 	}
 
-	prepared, err := c.prewriteTxn(startTS, commitTS, primaryKey, grouped, gids)
+	prepared, err := c.prewriteTxn(ctx, startTS, commitTS, primaryKey, grouped, gids, readKeys)
 	if err != nil {
 		return nil, err
 	}
@@ -246,8 +246,9 @@ func (c *ShardedCoordinator) dispatchSingleShardTxn(startTS, commitTS uint64, pr
 	if err != nil {
 		return nil, err
 	}
+	// Single-shard: read-set validated pre-Raft by the adapter.
 	resp, err := g.Txn.Commit([]*pb.Request{
-		onePhaseTxnRequest(startTS, commitTS, primaryKey, elems),
+		onePhaseTxnRequest(startTS, commitTS, primaryKey, elems, nil),
 	})
 	if err != nil {
 		return nil, errors.WithStack(err)
@@ -263,10 +264,12 @@ type preparedGroup struct {
 	keys []*pb.Mutation
 }
 
-func (c *ShardedCoordinator) prewriteTxn(startTS, commitTS uint64, primaryKey []byte, grouped map[uint64][]*pb.Mutation, gids []uint64) ([]preparedGroup, error) {
+func (c *ShardedCoordinator) prewriteTxn(ctx context.Context, startTS, commitTS uint64, primaryKey []byte, grouped map[uint64][]*pb.Mutation, gids []uint64, readKeys [][]byte) ([]preparedGroup, error) {
 	prepareMeta := txnMetaMutation(primaryKey, defaultTxnLockTTLms, 0)
 	prepared := make([]preparedGroup, 0, len(gids))
 
+	groupedReadKeys := c.groupReadKeysByShardID(readKeys)
+
 	for _, gid := range gids {
 		g, err := c.txnGroupForID(gid)
 		if err != nil {
@@ -277,6 +280,7 @@ func (c *ShardedCoordinator) prewriteTxn(startTS, commitTS uint64, primaryKey []
 			Phase:     pb.Phase_PREPARE,
 			Ts:        startTS,
 			Mutations: append([]*pb.Mutation{prepareMeta}, grouped[gid]...),
+			ReadKeys:  groupedReadKeys[gid],
 		}
 		if _, err := g.Txn.Commit([]*pb.Request{req}); err != nil {
 			c.abortPreparedTxn(startTS, primaryKey, prepared, abortTSFrom(startTS, commitTS))
@@ -285,6 +289,14 @@ func (c *ShardedCoordinator) prewriteTxn(startTS, commitTS uint64, primaryKey []
 		prepared = append(prepared, preparedGroup{gid: gid, keys: keyMutations(grouped[gid])})
 	}
 
+	// Validate read keys on read-only shards (shards that have read keys
+	// but no mutations in this transaction). Without this, a concurrent
+	// write to a read-only shard would go undetected.
+	if err := c.validateReadOnlyShards(ctx, groupedReadKeys, gids, startTS); err != nil {
+		c.abortPreparedTxn(startTS, primaryKey, prepared, abortTSFrom(startTS, commitTS))
+		return nil, err
+	}
+
 	return prepared, nil
 }
 
@@ -586,6 +598,81 @@ func (c *ShardedCoordinator) engineGroupIDForKey(key []byte) uint64 {
 	return route.GroupID
 }
 
+func (c *ShardedCoordinator) groupReadKeysByShardID(readKeys [][]byte) map[uint64][][]byte {
+	if len(readKeys) == 0 {
+		return nil
+	}
+	grouped := make(map[uint64][][]byte)
+	for _, key := range readKeys {
+		gid := c.engineGroupIDForKey(key)
+		if gid == 0 {
+			continue
+		}
+		grouped[gid] = append(grouped[gid], key)
+	}
+	return grouped
+}
+
+// validateReadOnlyShards checks read-write conflicts on shards that have
+// read keys but no mutations in this transaction. writeGIDs is the set of
+// shards that already received a PREPARE with their readKeys attached.
+//
+// Because these shards have no mutations, we cannot send a PREPARE request
+// (the FSM rejects empty mutation lists). Instead we issue a linearizable
+// read barrier on each read-only shard's Raft group (ensuring the local
+// FSM has applied all committed log entries) and then check LatestCommitTS
+// against the local store.
+//
+// NOTE: This check is performed outside the FSM's applyMu lock, so there
+// is a small TOCTOU window between the linearizable read barrier and the
+// LatestCommitTS check. A concurrent write that commits in this window may
+// go undetected. Full SSI for read-only shards in multi-shard transactions
+// would require a dedicated "read-validate" FSM request phase. For
+// single-shard transactions and write-shard read keys, validation is fully
+// atomic under applyMu.
+func (c *ShardedCoordinator) validateReadOnlyShards(ctx context.Context, groupedReadKeys map[uint64][][]byte, writeGIDs []uint64, startTS uint64) error {
+	if len(groupedReadKeys) == 0 {
+		return nil
+	}
+	writeSet := make(map[uint64]struct{}, len(writeGIDs))
+	for _, gid := range writeGIDs {
+		writeSet[gid] = struct{}{}
+	}
+	for gid, keys := range groupedReadKeys {
+		if _, isWrite := writeSet[gid]; isWrite {
+			continue
+		}
+		if err := c.validateReadKeysOnShard(ctx, gid, keys, startTS); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (c *ShardedCoordinator) validateReadKeysOnShard(ctx context.Context, gid uint64, keys [][]byte, startTS uint64) error {
+	g, ok := c.groups[gid]
+	if !ok {
+		return nil
+	}
+	// Linearizable read barrier: wait until the shard's FSM has applied
+	// all Raft-committed entries so LatestCommitTS reflects the latest
+	// committed state. Without this, a concurrent write that is committed
+	// in Raft but not yet applied locally would be invisible.
+	if _, err := linearizableReadEngineCtx(ctx, engineForGroup(g)); err != nil {
+		return errors.WithStack(err)
+	}
+	for _, key := range keys {
+		ts, exists, err := g.Store.LatestCommitTS(ctx, key)
+		if err != nil {
+			return errors.WithStack(err)
+		}
+		if exists && ts > startTS {
+			return errors.WithStack(store.NewWriteConflictError(key))
+		}
+	}
+	return nil
+}
+
 var _ Coordinator = (*ShardedCoordinator)(nil)
 
 func validateOperationGroup(reqs *OperationGroup[OP]) error {