fix: harden TN replica lifecycle

pkg/tnservice/cfg.go

@@ -31,6 +31,7 @@ import (
 	"github.com/matrixorigin/matrixone/pkg/pb/txn"
 	"github.com/matrixorigin/matrixone/pkg/txn/rpc"
 	"github.com/matrixorigin/matrixone/pkg/util/toml"
+	logtailservice "github.com/matrixorigin/matrixone/pkg/vm/engine/tae/logtail/service"
 	"github.com/matrixorigin/matrixone/pkg/vm/engine/tae/options"
 )
 
@@ -282,6 +283,10 @@ func (c *Config) Validate() error {
 	if c.LogtailServer.RpcMaxMessageSize <= 0 {
 		c.LogtailServer.RpcMaxMessageSize = toml.ByteSize(defaultRpcMaxMsgSize)
 	}
+	if err := logtailservice.ValidateRPCMaxMessageSize(
+		int64(c.LogtailServer.RpcMaxMessageSize)); err != nil {
+		return err
+	}
 	if c.LogtailServer.LogtailRPCStreamPoisonTime.Duration <= 0 {
 		c.LogtailServer.LogtailRPCStreamPoisonTime.Duration = defaultRPCStreamPoisonTime
 	}

pkg/tnservice/cfg_test.go

@@ -17,10 +17,19 @@ package tnservice
 import (
 	"testing"
 
+	"github.com/matrixorigin/matrixone/pkg/util/toml"
 	"github.com/matrixorigin/matrixone/pkg/vm/engine/tae/options"
 	"github.com/stretchr/testify/assert"
 )
 
+func TestValidateRejectsInvalidLogtailRPCMessageSize(t *testing.T) {
+	for _, size := range []toml.ByteSize{1, 101 * 1024 * 1024} {
+		c := &Config{UUID: "tn1"}
+		c.LogtailServer.RpcMaxMessageSize = size
+		assert.Error(t, c.Validate())
+	}
+}
+
 func TestValidate(t *testing.T) {
 	c := &Config{}
 	assert.Error(t, c.Validate())

pkg/tnservice/factory.go

@@ -31,7 +31,6 @@ import (
 	"github.com/matrixorigin/matrixone/pkg/vm/engine/memoryengine"
 	"github.com/matrixorigin/matrixone/pkg/vm/engine/tae/logstore/driver/logservicedriver"
 	"github.com/matrixorigin/matrixone/pkg/vm/engine/tae/options"
-	"go.uber.org/zap"
 )
 
 var (
@@ -47,36 +46,20 @@ func (s *store) createTxnStorage(
 	shard metadata.TNShard,
 	txnServer rpc.TxnServer,
 ) (storage.TxnStorage, error) {
-
-	factory := s.createLogServiceClientFactroy(shard)
-	closeLogClientFn := func(logClient logservice.Client) {
-		if err := logClient.Close(); err != nil {
-			s.rt.Logger().Error("close log client failed",
-				zap.Error(err))
-		}
-	}
-
 	switch s.cfg.Txn.Storage.Backend {
 	case StorageMEM:
-		logClient, err := factory()
-		if err != nil {
-			return nil, err
-		}
-		ts, err := s.newMemTxnStorage(shard, logClient, s.hakeeperClient)
-		if err != nil {
-			closeLogClientFn(logClient)
-			return nil, err
-		}
-		return ts, nil
+		return s.newMemTxnStorage(shard, s.hakeeperClient)
 
 	case StorageMEMKV:
+		factory := s.createLogServiceClientFactroy(shard)
 		logClient, err := factory()
 		if err != nil {
 			return nil, err
 		}
 		return s.newMemKVStorage(shard, logClient)
 
 	case StorageTAE:
+		factory := s.createLogServiceClientFactroy(shard)
 		ts, err := s.newTAEStorage(ctx, shard, factory, txnServer)
 		if err != nil {
 			return nil, err
@@ -119,7 +102,6 @@ func (s *store) newLogServiceClient(shard metadata.TNShard) (logservice.Client,
 
 func (s *store) newMemTxnStorage(
 	shard metadata.TNShard,
-	logClient logservice.Client,
 	hakeeper logservice.TNHAKeeperClient,
 ) (storage.TxnStorage, error) {
 	// should it be no fixed or a certain size?

pkg/tnservice/factory_test.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"testing"
 
+	"github.com/matrixorigin/matrixone/pkg/catalog"
 	"github.com/matrixorigin/matrixone/pkg/common/runtime"
 	"github.com/matrixorigin/matrixone/pkg/common/stopper"
 	"github.com/matrixorigin/matrixone/pkg/logservice"
@@ -53,3 +54,26 @@ func TestCreateTxnStorage(t *testing.T) {
 	assert.Error(t, err)
 	assert.Nil(t, v)
 }
+
+func TestCreateMemoryStorageDoesNotCreateUnusedLogClient(t *testing.T) {
+	catalog.SetupDefines("")
+	ctx := context.Background()
+	s := &store{
+		rt:             runtime.DefaultRuntime(),
+		cfg:            &Config{},
+		stopper:        stopper.NewStopper(""),
+		hakeeperClient: newTestHAKeeperClient(),
+	}
+	calls := 0
+	s.options.logServiceClientFactory = func(metadata.TNShard) (logservice.Client, error) {
+		calls++
+		return mem.NewMemLog(), nil
+	}
+	s.cfg.Txn.Storage.Backend = StorageMEM
+
+	v, err := s.createTxnStorage(ctx, metadata.TNShard{}, nil)
+	assert.NoError(t, err)
+	assert.NotNil(t, v)
+	assert.Zero(t, calls)
+	assert.NoError(t, v.Close(ctx))
+}

pkg/tnservice/replica.go

@@ -39,26 +39,54 @@ type replica struct {
 	createCtx    context.Context
 	cancelCreate context.CancelFunc
 	startedOnce  sync.Once
+	closeOnce    sync.Once
+	closeErr     error
 
 	mu struct {
 		sync.RWMutex
+		cond            *sync.Cond
 		starting        bool
 		cancelled       bool
 		destroyOnCancel bool
 		startErr        error
+		activeCalls     int
 	}
 }
 
+type txnServiceLease struct {
+	replica               *replica
+	service               service.TxnService
+	ctx                   context.Context
+	cancel                context.CancelFunc
+	stopCancelPropagation func() bool
+	releaseOnce           sync.Once
+}
+
+func (l *txnServiceLease) release() {
+	l.releaseOnce.Do(func() {
+		l.stopCancelPropagation()
+		l.cancel()
+		l.replica.mu.Lock()
+		l.replica.mu.activeCalls--
+		if l.replica.mu.activeCalls == 0 {
+			l.replica.mu.cond.Broadcast()
+		}
+		l.replica.mu.Unlock()
+	})
+}
+
 func newReplica(shard metadata.TNShard, rt runtime.Runtime) *replica {
 	ctx, cancel := context.WithCancel(context.Background())
-	return &replica{
+	r := &replica{
 		rt:           rt,
 		shard:        shard,
 		logger:       rt.Logger().With(util.TxnTNShardField(shard)),
 		startedC:     make(chan struct{}),
 		createCtx:    ctx,
 		cancelCreate: cancel,
 	}
+	r.mu.cond = sync.NewCond(&r.mu.RWMutex)
+	return r
 }
 
 func (r *replica) start(txnService service.TxnService) error {
@@ -124,49 +152,131 @@ func (r *replica) waitStartCompleted() {
 }
 
 func (r *replica) close(destroy bool) error {
+	r.cancelStart(destroy)
+	r.closeOnce.Do(func() {
+		r.closeErr = r.closeOnceFn()
+	})
+	return r.closeErr
+}
+
+func (r *replica) closeOnceFn() error {
 	r.mu.RLock()
 	starting := r.mu.starting
 	r.mu.RUnlock()
 	if !starting {
 		return nil
 	}
-	startErr := r.waitStarted(context.Background())
-	r.mu.RLock()
+
+	r.waitStartCompleted()
+	r.mu.Lock()
+	for r.mu.activeCalls > 0 {
+		r.mu.cond.Wait()
+	}
+	startErr := r.mu.startErr
 	txnService := r.service
-	r.mu.RUnlock()
+	destroy := r.mu.destroyOnCancel
+	r.mu.Unlock()
 	if txnService == nil {
 		return startErr
 	}
 	return errors.Join(startErr, txnService.Close(destroy))
 }
 
+func (r *replica) started() bool {
+	select {
+	case <-r.startedC:
+		r.mu.RLock()
+		defer r.mu.RUnlock()
+		return !r.mu.cancelled && r.mu.startErr == nil && r.service != nil
+	default:
+		return false
+	}
+}
+
 func (r *replica) handleLocalRequest(ctx context.Context, request *txn.TxnRequest, response *txn.TxnResponse) error {
-	if err := r.waitStarted(ctx); err != nil {
+	lease, err := r.acquireService(ctx)
+	if err != nil {
 		return err
 	}
+	defer lease.release()
 	prepareResponse(request, response)
 
 	switch request.Method {
 	case txn.TxnMethod_GetStatus:
-		return r.service.GetStatus(ctx, request, response)
+		return lease.service.GetStatus(lease.ctx, request, response)
 	case txn.TxnMethod_Prepare:
-		return r.service.Prepare(ctx, request, response)
+		return lease.service.Prepare(lease.ctx, request, response)
 	case txn.TxnMethod_CommitTNShard:
-		return r.service.CommitTNShard(ctx, request, response)
+		return lease.service.CommitTNShard(lease.ctx, request, response)
 	case txn.TxnMethod_RollbackTNShard:
-		return r.service.RollbackTNShard(ctx, request, response)
+		return lease.service.RollbackTNShard(lease.ctx, request, response)
 	default:
-		return moerr.NewNotSupportedf(ctx, "unknown txn request method: %s", request.Method.String())
+		return moerr.NewNotSupportedf(lease.ctx, "unknown txn request method: %s", request.Method.String())
 	}
 }
 
 func (r *replica) waitStarted(ctx context.Context) error {
+	if err := context.Cause(ctx); err != nil {
+		return err
+	}
+	if err := context.Cause(r.createCtx); err != nil {
+		return err
+	}
 	select {
 	case <-ctx.Done():
-		return ctx.Err()
+		return context.Cause(ctx)
+	case <-r.createCtx.Done():
+		return context.Cause(r.createCtx)
 	case <-r.startedC:
+		if err := context.Cause(ctx); err != nil {
+			return err
+		}
 		r.mu.RLock()
 		defer r.mu.RUnlock()
+		if r.mu.cancelled {
+			return context.Canceled
+		}
 		return r.mu.startErr
 	}
 }
+
+// acquireService pins the transaction service until release is called. Closing
+// a replica cancels the call context and waits for all acquired services to be
+// released before closing the underlying storage.
+func (r *replica) acquireService(ctx context.Context) (*txnServiceLease, error) {
+	if err := r.waitStarted(ctx); err != nil {
+		return nil, err
+	}
+
+	r.mu.Lock()
+	if err := context.Cause(ctx); err != nil {
+		r.mu.Unlock()
+		return nil, err
+	}
+	if r.mu.cancelled {
+		r.mu.Unlock()
+		return nil, context.Canceled
+	}
+	if r.mu.startErr != nil {
+		err := r.mu.startErr
+		r.mu.Unlock()
+		return nil, err
+	}
+	if r.service == nil {
+		r.mu.Unlock()
+		return nil, context.Canceled
+	}
+	txnService := r.service
+	r.mu.activeCalls++
+	r.mu.Unlock()
+
+	callCtx, cancel := context.WithCancel(ctx)
+	stopCancelPropagation := context.AfterFunc(r.createCtx, cancel)
+	return &txnServiceLease{
+		replica:               r,
+		service:               txnService,
+		ctx:                   callCtx,
+		cancel:                cancel,
+		stopCancelPropagation: stopCancelPropagation,
+	}, nil
+}