store/kv: optimize compactKeepIndex with binary search; add secondary shard idempotency

Agent-Logs-Url: https://github.com/bootjp/elastickv/sessions/a06bfa98-ebb4-486c-905f-68dd201385b1

Co-authored-by: bootjp <1306365+bootjp@users.noreply.github.com>

Author: Copilot

kv/fsm.go

@@ -365,6 +365,26 @@ func (f *kvFSM) handleCommitRequest(ctx context.Context, r *pb.Request) error {
 	if err != nil {
 		return err
 	}
+
+	// Secondary-shard fallback: txnCommitKey lives only on the primary shard.
+	// If this shard doesn't hold the primary key, detect an already-committed
+	// state by checking whether any target data key already has a committed
+	// version at or beyond commitTS. If so, use commitTS as the conflict-check
+	// baseline so that idempotent re-application doesn't trip on the
+	// already-written version.
+	if applyStartTS == startTS {
+		for _, mut := range uniq {
+			latestTS, exists, err := f.store.LatestCommitTS(ctx, mut.Key)
+			if err != nil {
+				return err
+			}
+			if exists && latestTS >= commitTS {
+				applyStartTS = commitTS
+				break
+			}
+		}
+	}
+
 	storeMuts, err := f.buildCommitStoreMutations(ctx, uniq, meta, startTS, commitTS)
 	if err != nil {
 		return err

kv/fsm_txn_test.go

@@ -118,6 +118,89 @@ func TestCommitIsIdempotentAfterCommitRecordExists(t *testing.T) {
 	require.Equal(t, commitTS, gotCommitTS)
 }
 
+func TestCommitIsIdempotentOnSecondaryShardWhenKeyAlreadyCommitted(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	st := store.NewMVCCStore()
+	fsm, ok := NewKvFSM(st).(*kvFSM)
+	require.True(t, ok)
+
+	startTS := uint64(15)
+	commitTS := uint64(25)
+	primaryKey := []byte("p") // lives on another shard; no txnCommitKey here
+	userKey := []byte("k")
+
+	// PREPARE: write lock and intent for userKey.
+	prepare := &pb.Request{
+		IsTxn: true,
+		Phase: pb.Phase_PREPARE,
+		Ts:    startTS,
+		Mutations: []*pb.Mutation{
+			{Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: primaryKey, LockTTLms: defaultTxnLockTTLms})},
+			{Op: pb.Op_PUT, Key: userKey, Value: []byte("v")},
+		},
+	}
+	require.NoError(t, applyFSMRequest(t, fsm, prepare))
+
+	// Simulate a partial commit: the data key is written at commitTS but the
+	// txn lock/intent are still present (inconsistent state that the
+	// secondary-shard idempotency check must handle).
+	require.NoError(t, st.PutAt(ctx, userKey, []byte("v"), commitTS, 0))
+
+	// COMMIT on this secondary shard (no txnCommitKey for primaryKey here).
+	// Without the secondary-shard LatestCommitTS check this would fail with a
+	// write-conflict error because userKey@commitTS > startTS.
+	commit := &pb.Request{
+		IsTxn: true,
+		Phase: pb.Phase_COMMIT,
+		Ts:    startTS,
+		Mutations: []*pb.Mutation{
+			{Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: primaryKey, CommitTS: commitTS})},
+			{Op: pb.Op_PUT, Key: userKey},
+		},
+	}
+	require.NoError(t, applyFSMRequest(t, fsm, commit))
+
+	// The committed value should still be readable.
+	v, err := st.GetAt(ctx, userKey, ^uint64(0))
+	require.NoError(t, err)
+	require.Equal(t, []byte("v"), v)
+
+	// The lock and intent should be cleaned up.
+	_, err = st.GetAt(ctx, txnLockKey(userKey), ^uint64(0))
+	require.ErrorIs(t, err, store.ErrKeyNotFound)
+	_, err = st.GetAt(ctx, txnIntentKey(userKey), ^uint64(0))
+	require.ErrorIs(t, err, store.ErrKeyNotFound)
+}
+
+func TestCommitPropagatesSecondaryShardLatestCommitTSError(t *testing.T) {
+	t.Parallel()
+
+	underlying := store.NewMVCCStore()
+	userKey := []byte("k")
+	errorStore := erroringLatestCommitStore{MVCCStore: underlying, key: userKey}
+	fsm, ok := NewKvFSM(errorStore).(*kvFSM)
+	require.True(t, ok)
+
+	startTS := uint64(16)
+	commitTS := uint64(26)
+	primaryKey := []byte("p") // no txnCommitKey in store → secondary-shard path
+
+	commit := &pb.Request{
+		IsTxn: true,
+		Phase: pb.Phase_COMMIT,
+		Ts:    startTS,
+		Mutations: []*pb.Mutation{
+			{Op: pb.Op_PUT, Key: []byte(txnMetaPrefix), Value: EncodeTxnMeta(TxnMeta{PrimaryKey: primaryKey, CommitTS: commitTS})},
+			{Op: pb.Op_PUT, Key: userKey},
+		},
+	}
+	err := applyFSMRequest(t, fsm, commit)
+	require.Error(t, err)
+	require.ErrorIs(t, err, ErrTestLatestCommitTS)
+}
+
 func TestCommitRejectsMissingPrimaryKey(t *testing.T) {
 	t.Parallel()
 

store/compaction_unit_test.go

@@ -55,8 +55,7 @@ func TestCompactKeepIndex_MultipleVersions(t *testing.T) {
 		{TS: 300, Value: []byte("v300")},
 		{TS: 400, Value: []byte("v400")},
 	}
-	// minTS=250: iterates from end: i=3(400>250), i=2(300>250), i=1(200<=250)
-	// keepIdx=1, > 0 → compact versions before idx 1
+	// minTS=250: first index with TS>250 is idx 2 (TS=300); keepIdx = 2-1 = 1.
 	assert.Equal(t, 1, compactKeepIndex(versions, 250))
 }
 
@@ -67,7 +66,7 @@ func TestCompactKeepIndex_AllBelowMinTS(t *testing.T) {
 		{TS: 20, Value: []byte("v20")},
 		{TS: 30, Value: []byte("v30")},
 	}
-	// minTS=100: iterates from end: i=2(30<=100) → keepIdx=2, > 0 → return 2
+	// minTS=100: first index with TS>100 is len(versions)=3; keepIdx = 3-1 = 2.
 	// compact versions before idx 2 (i.e. TS=10 and TS=20)
 	assert.Equal(t, 2, compactKeepIndex(versions, 100))
 }

store/mvcc_store.go

@@ -901,12 +901,13 @@ func compactKeepIndex(versions []VersionedValue, minTS uint64) int {
 	if len(versions) == 0 {
 		return -1
 	}
-	keepIdx := -1
-	for i, version := range versions {
-		if version.TS <= minTS {
-			keepIdx = i
-		}
-	}
+	// versions are sorted ascending by TS; use binary search for O(log n).
+	// Find the first index where TS > minTS; everything before it has TS <= minTS.
+	n := sort.Search(len(versions), func(i int) bool {
+		return versions[i].TS > minTS
+	})
+	// n-1 is the last index with TS <= minTS.
+	keepIdx := n - 1
 	if keepIdx <= 0 {
 		return -1
 	}