fix(lockservice): honor context cancellation in lock waits

pkg/cnservice/server.go

@@ -894,8 +894,13 @@ func (s *service) initLockService() {
 	cfg := s.getLockServiceConfig()
 	s.lockService = lockservice.NewLockService(
 		cfg,
-		lockservice.WithWait(func() {
-			<-s.hakeeperConnected
+		lockservice.WithWait(func(ctx context.Context) error {
+			select {
+			case <-s.hakeeperConnected:
+				return nil
+			case <-ctx.Done():
+				return ctx.Err()
+			}
 		}))
 	runtime.ServiceRuntime(s.cfg.UUID).SetGlobalVariables(runtime.LockService, s.lockService)
 	lockservice.SetLockServiceByServiceID(s.cfg.UUID, s.lockService)

pkg/lockservice/deadlock.go

@@ -38,7 +38,7 @@ var (
 type detector struct {
 	logger            *log.MOLogger
 	c                 chan deadlockTxn
-	waitTxnsFetchFunc func(pb.WaitTxn, *waiters) (bool, error)
+	waitTxnsFetchFunc func(context.Context, pb.WaitTxn, *waiters) (bool, error)
 	waitTxnAbortFunc  func(pb.WaitTxn, error)
 	ignoreTxns        sync.Map // txnID -> any
 	stopper           *stopper.Stopper
@@ -56,7 +56,7 @@ type detector struct {
 // txn.
 func newDeadlockDetector(
 	logger *log.MOLogger,
-	waitTxnsFetchFunc func(pb.WaitTxn, *waiters) (bool, error),
+	waitTxnsFetchFunc func(context.Context, pb.WaitTxn, *waiters) (bool, error),
 	waitTxnAbortFunc func(pb.WaitTxn, error),
 ) *detector {
 	d := &detector{
@@ -82,6 +82,9 @@ func (d *detector) close() {
 	d.mu.closed = true
 	d.mu.Unlock()
 	d.stopper.Stop()
+	d.mu.Lock()
+	clear(d.mu.activeCheckTxn)
+	d.mu.Unlock()
 	close(d.c)
 }
 
@@ -141,13 +144,16 @@ func (d *detector) doCheck(ctx context.Context) {
 
 	w := &waiters{ignoreTxns: &d.ignoreTxns}
 	for {
+		if ctx.Err() != nil {
+			return
+		}
 		select {
 		case <-ctx.Done():
 			return
 		case txn := <-d.c:
 			v2.TxnDeadlockDetectorQueueDepthGauge.Set(float64(len(d.c)))
 			w.reset(txn)
-			hasDeadlock, deadlockTxn, err := d.checkDeadlock(w)
+			hasDeadlock, deadlockTxn, err := d.checkDeadlock(ctx, w)
 			if hasDeadlock {
 				if err == nil {
 					err = ErrDeadLockDetected
@@ -162,11 +168,14 @@ func (d *detector) doCheck(ctx context.Context) {
 	}
 }
 
-func (d *detector) checkDeadlock(w *waiters) (bool, pb.WaitTxn, error) {
+func (d *detector) checkDeadlock(ctx context.Context, w *waiters) (bool, pb.WaitTxn, error) {
 	for {
+		if err := ctx.Err(); err != nil {
+			return false, pb.WaitTxn{}, err
+		}
 		// find deadlock
 		txn := w.getCheckTargetTxn()
-		added, err := d.waitTxnsFetchFunc(txn, w)
+		added, err := d.waitTxnsFetchFunc(ctx, txn, w)
 		if err != nil {
 			logCheckDeadLockFailed(d.logger, txn, w.root.startTxn(), err)
 			return false, pb.WaitTxn{}, err

pkg/lockservice/deadlock_test.go

@@ -15,6 +15,7 @@
 package lockservice
 
 import (
+	"context"
 	"encoding/hex"
 	"testing"
 	"time"
@@ -23,6 +24,7 @@ import (
 	"github.com/matrixorigin/matrixone/pkg/common/runtime"
 	pb "github.com/matrixorigin/matrixone/pkg/pb/lock"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestCheckWithDeadlock(t *testing.T) {
@@ -43,7 +45,7 @@ func TestCheckWithDeadlock(t *testing.T) {
 
 		d := newDeadlockDetector(
 			runtime.DefaultRuntime().Logger(),
-			func(txn pb.WaitTxn, w *waiters) (bool, error) {
+			func(_ context.Context, txn pb.WaitTxn, w *waiters) (bool, error) {
 				for _, v := range m[string(txn.TxnID)] {
 					if !w.add(v, "") {
 						return false, nil
@@ -79,6 +81,34 @@ func TestCheckWithDeadlock(t *testing.T) {
 	})
 }
 
+func TestDeadlockDetectorCloseCancelsCheck(t *testing.T) {
+	started := make(chan struct{}, 1)
+	aborted := make(chan struct{}, 1)
+	d := newDeadlockDetector(
+		runtime.DefaultRuntime().Logger(),
+		func(ctx context.Context, _ pb.WaitTxn, _ *waiters) (bool, error) {
+			select {
+			case started <- struct{}{}:
+			default:
+			}
+			<-ctx.Done()
+			return false, ctx.Err()
+		},
+		func(pb.WaitTxn, error) { aborted <- struct{}{} },
+	)
+	require.NoError(t, d.check([]byte("holder"), pb.WaitTxn{TxnID: []byte("waiter")}))
+	<-started
+	d.close()
+	select {
+	case <-aborted:
+		t.Fatal("deadlock abort callback ran after detector cancellation")
+	default:
+	}
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	require.Empty(t, d.mu.activeCheckTxn)
+}
+
 func TestCheckWithDeadlockWith2Txn(t *testing.T) {
 	reuse.RunReuseTests(func() {
 		txn1 := []byte("t1")
@@ -95,7 +125,7 @@ func TestCheckWithDeadlockWith2Txn(t *testing.T) {
 
 		d := newDeadlockDetector(
 			runtime.DefaultRuntime().Logger(),
-			func(txn pb.WaitTxn, w *waiters) (bool, error) {
+			func(_ context.Context, txn pb.WaitTxn, w *waiters) (bool, error) {
 				for _, v := range depends[string(txn.TxnID)] {
 					if !w.add(v, "") {
 						return false, nil
@@ -282,7 +312,7 @@ func TestCheckWithComplexDeadlock(t *testing.T) {
 		// Create the deadlock detector
 		d := newDeadlockDetector(
 			runtime.DefaultRuntime().Logger(),
-			func(txn pb.WaitTxn, w *waiters) (bool, error) {
+			func(_ context.Context, txn pb.WaitTxn, w *waiters) (bool, error) {
 				for _, v := range depends[string(txn.TxnID)] {
 					if !w.add(v, "") {
 						return false, nil
@@ -363,7 +393,7 @@ func TestCheckDeadlock(t *testing.T) {
 		// Create the deadlock detector
 		d := newDeadlockDetector(
 			runtime.DefaultRuntime().Logger(),
-			func(txn pb.WaitTxn, w *waiters) (bool, error) {
+			func(_ context.Context, txn pb.WaitTxn, w *waiters) (bool, error) {
 				for _, v := range depends[string(txn.TxnID)] {
 					if !w.add(v, "") {
 						return false, nil

pkg/lockservice/lock_table_keeper.go

@@ -266,10 +266,7 @@ func (k *lockTableKeeper) invalidateRemoteBind(
 }
 
 func (k *lockTableKeeper) doKeepLockTableBind(ctx context.Context) {
-	if k.service.isStatus(pb.Status_ServiceLockWaiting) &&
-		k.service.activeTxnHolder.empty() {
-		k.service.setStatus(pb.Status_ServiceUnLockSucc)
-	}
+	k.service.tryCompleteDrain()
 
 	req := acquireRequest()
 	defer releaseRequest(req)

pkg/lockservice/lock_table_local.go

@@ -397,23 +397,34 @@ func (l *localLockTable) unlock(
 }
 
 func (l *localLockTable) getLock(
+	ctx context.Context,
 	key []byte,
 	txn pb.WaitTxn,
-	fn func(Lock)) {
+	fn func(Lock)) error {
+	if err := ctx.Err(); err != nil {
+		return err
+	}
 	l.mu.RLock()
 	defer l.mu.RUnlock()
+	if err := ctx.Err(); err != nil {
+		return err
+	}
 	if l.mu.closed {
-		return
+		return nil
 	}
 	lock, ok := l.mu.store.Get(key)
 	if ok {
 		fn(lock)
 	}
+	return nil
 }
 
 func (l *localLockTable) getLockHolder(ctx context.Context, key []byte) (pb.WaitTxn, bool, error) {
 	l.mu.RLock()
 	defer l.mu.RUnlock()
+	if err := ctx.Err(); err != nil {
+		return pb.WaitTxn{}, false, err
+	}
 	if l.mu.closed {
 		return pb.WaitTxn{}, false, nil
 	}
@@ -461,6 +472,9 @@ func (l *localLockTable) doAcquireLock(c *lockContext) error {
 	if l.mu.closed {
 		return moerr.NewInvalidStateNoCtx("local lock table closed")
 	}
+	if err := c.ctx.Err(); err != nil {
+		return err
+	}
 
 	switch c.opts.Granularity {
 	case pb.Granularity_Row:

pkg/lockservice/lock_table_local_test.go

@@ -109,7 +109,7 @@ func TestCloseLocalLockTableWithBlockedWaiter(t *testing.T) {
 				require.Equal(t, ErrLockTableNotFound, err)
 			}()
 
-			v, err := l.getLockTable(0, tableID)
+			v, err := l.getLockTable(context.Background(), 0, tableID)
 			require.NoError(t, err)
 			lt := v.(*localLockTable)
 			for {
@@ -408,7 +408,7 @@ func TestMergeRangeWithNoConflict(t *testing.T) {
 			table := uint64(10)
 			for _, c := range cases {
 				stopper := stopper.NewStopper("")
-				v, err := l.getLockTableWithCreate(0, table, nil, pb.Sharding_None)
+				v, err := l.getLockTableWithCreate(context.Background(), 0, table, nil, pb.Sharding_None)
 				require.NoError(t, err)
 				lt := v.(*localLockTable)
 
@@ -542,7 +542,7 @@ func TestLocalLockTableMultipleRowLocksCannotMissIfFoundSelfTxn(t *testing.T) {
 			require.NoError(t, l.Unlock(ctx, []byte{2}, timestamp.Timestamp{}))
 
 			wg.Wait()
-			v, err := l.getLockTable(0, tableID)
+			v, err := l.getLockTable(context.Background(), 0, tableID)
 			require.NoError(t, err)
 			lt := v.(*localLockTable)
 			lt.mu.Lock()
@@ -616,7 +616,7 @@ func TestIssue9856(t *testing.T) {
 				json.MustUnmarshal([]byte(r), v)
 				_, err := l.Lock(ctx, tableID, [][]byte{[]byte(v.Start), []byte(v.End)}, []byte("txn1"), option)
 				require.NoError(t, err)
-				vv, err := l.getLockTable(0, tableID)
+				vv, err := l.getLockTable(context.Background(), 0, tableID)
 				require.NoError(t, err)
 				lt := vv.(*localLockTable)
 				lt.mu.Lock()
@@ -783,7 +783,7 @@ func TestLockedTSIsLastCommittedTS(t *testing.T) {
 			defer cancel()
 
 			tableID := uint64(10)
-			v, err := l.getLockTableWithCreate(0, tableID, nil, pb.Sharding_None)
+			v, err := l.getLockTableWithCreate(context.Background(), 0, tableID, nil, pb.Sharding_None)
 			require.NoError(t, err)
 			lt := v.(*localLockTable)
 			lt.mu.Lock()
@@ -854,7 +854,7 @@ func TestLockedTSIsLastCommittedTSWithRange(t *testing.T) {
 			defer cancel()
 
 			tableID := uint64(10)
-			v, err := l.getLockTableWithCreate(0, tableID, nil, pb.Sharding_None)
+			v, err := l.getLockTableWithCreate(context.Background(), 0, tableID, nil, pb.Sharding_None)
 			require.NoError(t, err)
 			lt := v.(*localLockTable)
 			lt.mu.Lock()
@@ -935,7 +935,7 @@ func Test15608(t *testing.T) {
 			_, err := s1.Lock(ctx, table, rows, txn1, option)
 			require.NoError(t, err, err)
 
-			v, err := s1.getLockTable(0, table)
+			v, err := s1.getLockTable(context.Background(), 0, table)
 			require.NoError(t, err)
 			lt := v.(*localLockTable)
 			lt.options.beforeCloseFirstWaiter = func(c *lockContext) {
@@ -1095,7 +1095,7 @@ func TestCannotHungIfRangeConflictWithRowMultiTimes(t *testing.T) {
 			add(txn1, key4, pb.Granularity_Row)
 			close(startTxn3)
 
-			v, err := l.getLockTable(0, tableID)
+			v, err := l.getLockTable(context.Background(), 0, tableID)
 			require.NoError(t, err)
 			lt := v.(*localLockTable)
 			txn3WaitTimes := 0
@@ -1423,7 +1423,7 @@ func TestRangeLockModeUpgradeUpdatesBothEnds(t *testing.T) {
 			require.NoError(t, err)
 
 			// Verify both ends are Shared
-			v, err := l.getLockTable(0, tableID)
+			v, err := l.getLockTable(context.Background(), 0, tableID)
 			require.NoError(t, err)
 			lt := v.(*localLockTable)
 
@@ -1527,7 +1527,7 @@ func TestSetModePairedRangeLockDirect(t *testing.T) {
 			require.NoError(t, err)
 
 			// Get the lock table and directly test setModePairedRangeLock
-			v, err := l.getLockTable(0, tableID)
+			v, err := l.getLockTable(context.Background(), 0, tableID)
 			require.NoError(t, err)
 			lt := v.(*localLockTable)
 
@@ -1585,7 +1585,7 @@ func TestSetModePairedRangeLockFromRangeEnd(t *testing.T) {
 			require.NoError(t, err)
 
 			// Get the lock table and directly test setModePairedRangeLock from range-end
-			v, err := l.getLockTable(0, tableID)
+			v, err := l.getLockTable(context.Background(), 0, tableID)
 			require.NoError(t, err)
 			lt := v.(*localLockTable)
 
@@ -1637,7 +1637,7 @@ func TestSetModePairedRangeLockRowLockNoOp(t *testing.T) {
 			require.NoError(t, err)
 
 			// Get the lock table
-			v, err := l.getLockTable(0, tableID)
+			v, err := l.getLockTable(context.Background(), 0, tableID)
 			require.NoError(t, err)
 			lt := v.(*localLockTable)
 
@@ -1698,7 +1698,7 @@ func TestRangeLockWithInterleavedRowLocks(t *testing.T) {
 			require.NoError(t, err)
 
 			// Verify btree structure: [0:row] [1:range-start] [10:range-end]
-			v, err := l.getLockTable(0, tableID)
+			v, err := l.getLockTable(context.Background(), 0, tableID)
 			require.NoError(t, err)
 			lt := v.(*localLockTable)