fix(txn): restore recovery routes and cancellation

pkg/clusterservice/cluster.go

@@ -121,6 +121,33 @@ func GetCNServiceWithoutWorkingStateWithContext(
 	return ctx.Err()
 }
 
+// GetAllTNServicesWithContext returns a TN service snapshot without waiting
+// past ctx for the built-in cluster's initial HAKeeper refresh.
+func GetAllTNServicesWithContext(
+	ctx context.Context,
+	service MOCluster,
+) ([]metadata.TNService, error) {
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	if service == nil {
+		return nil, moerr.NewInternalErrorNoCtx("mocluster service is not initialized")
+	}
+	if builtIn, ok := service.(*cluster); ok {
+		if err := builtIn.waitReadyWithContext(ctx); err != nil {
+			return nil, err
+		}
+		services := builtIn.services.Load()
+		return append([]metadata.TNService(nil), services.tn...), ctx.Err()
+	}
+
+	services := service.GetAllTNServices()
+	return services, ctx.Err()
+}
+
 func lookupMOCluster(service string) (MOCluster, bool, error) {
 	rt := runtime.ServiceRuntime(service)
 	if rt == nil {
@@ -339,9 +366,14 @@ func (c *cluster) Refresh(ctx context.Context) error {
 }
 
 func (c *cluster) Close() {
-	c.waitReady()
 	c.stopper.Stop()
-	close(c.forceRefreshC)
+	// A failed initial refresh leaves readiness waiters blocked. Once the
+	// refresh task has stopped, release them so shutdown does not depend on
+	// HAKeeper becoming available.
+	c.readyOnce.Do(func() {
+		c.ready.Store(true)
+		close(c.readyC)
+	})
 }
 
 // DebugUpdateCNLabel implements the MOCluster interface.

pkg/clusterservice/cluster_test.go

@@ -70,6 +70,16 @@ func TestCNServiceSnapshotHonorsCancellationWhileClusterStarts(t *testing.T) {
 	require.ErrorIs(t, err, context.DeadlineExceeded)
 }
 
+func TestTNServiceSnapshotHonorsCancellationWhileClusterStarts(t *testing.T) {
+	c := &cluster{readyC: make(chan struct{})}
+	c.services.Store(&services{})
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
+	defer cancel()
+
+	_, err := GetAllTNServicesWithContext(ctx, c)
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+}
+
 func TestClusterForceRefresh(t *testing.T) {
 	runClusterTest(
 		time.Hour,

pkg/tnservice/replica.go

@@ -35,6 +35,7 @@ type replica struct {
 	logger       *log.MOLogger
 	shard        metadata.TNShard
 	service      service.TxnService
+	serviceC     chan struct{}
 	startedC     chan struct{}
 	createCtx    context.Context
 	cancelCreate context.CancelFunc
@@ -81,6 +82,7 @@ func newReplica(shard metadata.TNShard, rt runtime.Runtime) *replica {
 		rt:           rt,
 		shard:        shard,
 		logger:       rt.Logger().With(util.TxnTNShardField(shard)),
+		serviceC:     make(chan struct{}),
 		startedC:     make(chan struct{}),
 		createCtx:    ctx,
 		cancelCreate: cancel,
@@ -114,6 +116,7 @@ func (r *replica) startReserved(txnService service.TxnService) error {
 	}
 	r.service = txnService
 	r.mu.Unlock()
+	close(r.serviceC)
 
 	err := txnService.Start()
 	r.finishStart(err)
@@ -159,13 +162,39 @@ func (r *replica) close(destroy bool) error {
 	return r.closeErr
 }
 
+func (r *replica) cancelRecovery() {
+	r.mu.RLock()
+	starting := r.mu.starting
+	txnService := r.service
+	r.mu.RUnlock()
+	if !starting {
+		return
+	}
+	if txnService == nil {
+		// Once start is reserved, startReserved either publishes the service or
+		// finishStart reports that startup ended without one.
+		select {
+		case <-r.serviceC:
+		case <-r.startedC:
+		}
+		r.mu.RLock()
+		txnService = r.service
+		r.mu.RUnlock()
+	}
+	if txnService != nil {
+		txnService.CancelRecovery()
+	}
+}
+
 func (r *replica) closeOnceFn() error {
 	r.mu.RLock()
 	starting := r.mu.starting
 	r.mu.RUnlock()
 	if !starting {
 		return nil
 	}
+	// Recovery may block Start indefinitely while waiting for a participant.
+	r.cancelRecovery()
 
 	r.waitStartCompleted()
 	r.mu.Lock()

pkg/tnservice/replica_test.go

@@ -17,13 +17,20 @@ package tnservice
 import (
 	"context"
 	"errors"
+	"sync"
 	"testing"
 	"time"
 
+	"github.com/matrixorigin/matrixone/pkg/clusterservice"
 	"github.com/matrixorigin/matrixone/pkg/common/runtime"
+	"github.com/matrixorigin/matrixone/pkg/defines"
+	"github.com/matrixorigin/matrixone/pkg/fileservice"
+	logpb "github.com/matrixorigin/matrixone/pkg/pb/logservice"
+	"github.com/matrixorigin/matrixone/pkg/pb/metadata"
 	"github.com/matrixorigin/matrixone/pkg/pb/txn"
 	"github.com/matrixorigin/matrixone/pkg/txn/service"
 	"github.com/matrixorigin/matrixone/pkg/txn/storage"
+	"github.com/matrixorigin/matrixone/pkg/txn/storage/mem"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -39,6 +46,39 @@ type startErrorStorage struct {
 	destroyCalls int
 }
 
+type closeUnblocksStartTxnService struct {
+	service.TxnService
+	started   chan struct{}
+	closed    chan struct{}
+	startOnce sync.Once
+	closeOnce sync.Once
+}
+
+type signalingRecoveryCluster struct {
+	clusterservice.MOCluster
+	entered chan struct{}
+	once    sync.Once
+}
+
+func (c *signalingRecoveryCluster) GetAllTNServices() []metadata.TNService {
+	c.once.Do(func() { close(c.entered) })
+	return c.MOCluster.GetAllTNServices()
+}
+
+func (s *closeUnblocksStartTxnService) Start() error {
+	s.startOnce.Do(func() { close(s.started) })
+	<-s.closed
+	return context.Canceled
+}
+
+func (s *closeUnblocksStartTxnService) CancelRecovery() {
+	s.closeOnce.Do(func() { close(s.closed) })
+}
+
+func (s *closeUnblocksStartTxnService) Close(bool) error {
+	return nil
+}
+
 type closeTrackingTxnService struct {
 	service.TxnService
 	closeCalls int
@@ -210,6 +250,133 @@ func TestCloseFailedStartReplica(t *testing.T) {
 	}
 }
 
+func TestCloseCancelsBlockedReplicaStart(t *testing.T) {
+	txnService := &closeUnblocksStartTxnService{
+		started: make(chan struct{}),
+		closed:  make(chan struct{}),
+	}
+	r := newReplica(newTestTNShard(1, 2, 3), runtime.DefaultRuntime())
+	startResult := make(chan error, 1)
+	go func() {
+		startResult <- r.start(txnService)
+	}()
+
+	select {
+	case <-txnService.started:
+	case <-time.After(time.Second):
+		t.Fatal("replica start did not begin")
+	}
+
+	closeResult := make(chan error, 1)
+	go func() {
+		closeResult <- r.close(false)
+	}()
+
+	select {
+	case err := <-closeResult:
+		require.ErrorIs(t, err, context.Canceled)
+	case <-time.After(time.Second):
+		t.Fatal("close did not cancel blocked replica start")
+	}
+	require.ErrorIs(t, <-startResult, context.Canceled)
+}
+
+func TestRemoveReplicaCancelsBlockedStartBeforeWaiting(t *testing.T) {
+	txnService := &closeUnblocksStartTxnService{
+		started: make(chan struct{}),
+		closed:  make(chan struct{}),
+	}
+	r := newReplica(newTestTNShard(1, 2, 3), runtime.DefaultRuntime())
+	startResult := make(chan error, 1)
+	go func() { startResult <- r.start(txnService) }()
+	select {
+	case <-txnService.started:
+	case <-time.After(time.Second):
+		t.Fatal("replica start did not begin")
+	}
+
+	fs, err := fileservice.NewMemoryFS(
+		defines.LocalFileServiceName,
+		fileservice.DisabledCacheConfig,
+		nil,
+	)
+	require.NoError(t, err)
+	s := &store{
+		cfg:                 &Config{UUID: "test"},
+		rt:                  runtime.DefaultRuntime(),
+		metadataFileService: fs,
+		replicas:            &sync.Map{},
+	}
+	s.replicas.Store(r.shard.ShardID, r)
+
+	removed := make(chan error, 1)
+	go func() { removed <- s.removeReplicaLocked(r.shard.ShardID) }()
+	select {
+	case err := <-removed:
+		require.ErrorIs(t, err, context.Canceled)
+	case <-time.After(time.Second):
+		t.Fatal("removeReplicaLocked waited for Start before canceling recovery")
+	}
+	require.ErrorIs(t, <-startResult, context.Canceled)
+	require.Nil(t, s.getReplica(r.shard.ShardID))
+}
+
+func TestCloseCancelsReplicaBlockedInRecovery(t *testing.T) {
+	meta := service.NewTestTxn(1, 1, 1)
+	meta.Status = txn.TxnStatus_Prepared
+	meta.PreparedTS = service.NewTestTimestamp(2)
+	meta.TNShards = append(meta.TNShards, metadata.TNShard{
+		TNShardRecord: metadata.TNShardRecord{ShardID: 99},
+	})
+	mlog := mem.NewMemLog()
+	data := (&mem.KVLog{Txn: meta}).MustMarshal()
+	record := mlog.GetLogRecord(len(data))
+	record.Type = logpb.UserRecord
+	record.Data = data
+	_, err := mlog.Append(context.Background(), record)
+	require.NoError(t, err)
+
+	sender := service.NewTestSender()
+	t.Cleanup(func() { require.NoError(t, sender.Close()) })
+	txnService := service.NewTestTxnServiceWithLog(
+		t, 1, sender, service.NewTestClock(0), mlog)
+	baseCluster := clusterservice.NewMOCluster(
+		"dn-uuid", nil, time.Hour,
+		clusterservice.WithDisableRefresh(),
+		clusterservice.WithServices(nil, nil),
+	)
+	t.Cleanup(baseCluster.Close)
+	cluster := &signalingRecoveryCluster{
+		MOCluster: baseCluster,
+		entered:   make(chan struct{}),
+	}
+	runtime.ServiceRuntime("dn-uuid").SetGlobalVariables(runtime.ClusterService, cluster)
+
+	r := newReplica(newTestTNShard(1, 2, 3), runtime.DefaultRuntime())
+	startResult := make(chan error, 1)
+	go func() { startResult <- r.start(txnService) }()
+	select {
+	case <-cluster.entered:
+	case <-time.After(time.Second):
+		t.Fatal("recovery did not reach the missing participant route wait")
+	}
+
+	closeResult := make(chan error, 1)
+	go func() { closeResult <- r.close(false) }()
+	select {
+	case err := <-closeResult:
+		require.NoError(t, err)
+	case <-time.After(time.Second):
+		t.Fatal("replica close did not cancel real transaction recovery")
+	}
+	select {
+	case err := <-startResult:
+		require.NoError(t, err)
+	case <-time.After(time.Second):
+		t.Fatal("replica start remained blocked after recovery cancellation")
+	}
+}
+
 func TestWaitStarted(t *testing.T) {
 	r := newReplica(newTestTNShard(1, 2, 3), runtime.DefaultRuntime())
 	c := make(chan struct{})

pkg/tnservice/store.go

@@ -238,16 +238,20 @@ func (s *store) Start() error {
 }
 
 func (s *store) Close() error {
-	s.stopper.Stop()
-	s.moCluster.Close()
-
-	var err error
 	// Reject new replica calls and cancel active call contexts before waiting
-	// for the RPC server to drain. Storage remains open until the drain ends.
+	// for store tasks. A published service may be blocked in recovery, so its
+	// cancellation must be delivered before joining the store stopper. Storage
+	// remains open until the RPC server drains below.
 	s.replicas.Range(func(_, value any) bool {
-		value.(*replica).cancelStart(false)
+		r := value.(*replica)
+		r.cancelStart(false)
+		r.cancelRecovery()
 		return true
 	})
+	s.stopper.Stop()
+	s.moCluster.Close()
+
+	var err error
 	if s.queryService != nil {
 		err = errors.Join(err, s.queryService.Close())
 	}
@@ -333,9 +337,7 @@ func (s *store) createReplicaLocked(shard metadata.TNShard) error {
 	}
 
 	err := s.stopper.RunTask(func(stopperCtx context.Context) {
-		stopCancelPropagation := context.AfterFunc(stopperCtx, func() {
-			r.cancelStart(false)
-		})
+		stopCancelPropagation := propagateReplicaStopperCancellation(stopperCtx, r)
 		defer stopCancelPropagation()
 
 		for {
@@ -403,6 +405,16 @@ func (s *store) createReplicaLocked(shard metadata.TNShard) error {
 	return nil
 }
 
+func propagateReplicaStopperCancellation(
+	stopperCtx context.Context,
+	r *replica,
+) func() bool {
+	return context.AfterFunc(stopperCtx, func() {
+		r.cancelStart(false)
+		r.cancelRecovery()
+	})
+}
+
 func waitCreateRetry(stopperCtx, createCtx context.Context) error {
 	timer := time.NewTimer(retryCreateStorageInterval)
 	defer timer.Stop()
@@ -418,8 +430,6 @@ func waitCreateRetry(stopperCtx, createCtx context.Context) error {
 
 func (s *store) removeReplicaLocked(tnShardID uint64) error {
 	if r := s.getReplica(tnShardID); r != nil {
-		r.cancelStart(true)
-		r.waitStartCompleted()
 		err := r.close(true)
 		s.replicas.CompareAndDelete(tnShardID, r)
 		s.removeTNShardLocked(tnShardID)

pkg/tnservice/store_rpc_handler_test.go

@@ -47,6 +47,8 @@ func (s *leaseCancelReadTxnService) Start() error {
 	return nil
 }
 
+func (s *leaseCancelReadTxnService) CancelRecovery() {}
+
 func (s *leaseCancelReadTxnService) Close(bool) error {
 	return nil
 }