fix: using async context instead of CopyContext

dtmsvr/trans_process.go

@@ -33,7 +33,7 @@ func (t *TransGlobal) process(branches []TransBranch) error {
 		dtmimp.MustUnmarshalString(t.ExtData, &t.Ext)
 	}
 	if !t.WaitResult {
-		ctx := CopyContext(t.Context)
+		ctx := NewAsyncContext(t.Context)
 		go func(ctx context.Context) {
 			err := t.processInner(ctx, branches)
 			if err != nil && !errors.Is(err, dtmimp.ErrOngoing) {

dtmsvr/trans_type_msg.go

@@ -92,7 +92,7 @@ func (t *transMsgProcessor) ProcessOnce(ctx context.Context, branches []TransBra
 			continue
 		}
 		if t.Concurrent {
-			copyCtx := CopyContext(ctx)
+			copyCtx := NewAsyncContext(ctx)
 			started++
 			go func(ctx context.Context, pos int) {
 				resultsChan <- t.execBranch(ctx, b, pos)

dtmsvr/trans_type_saga.go

@@ -163,7 +163,7 @@ func (t *transSagaProcessor) ProcessOnce(ctx context.Context, branches []TransBr
 			if branchResults[b].op == dtmimp.OpAction {
 				rsAStarted++
 			}
-			copyCtx := CopyContext(ctx)
+			copyCtx := NewAsyncContext(ctx)
 			go asyncExecBranch(copyCtx, b)
 		}
 	}
@@ -180,7 +180,7 @@ func (t *transSagaProcessor) ProcessOnce(ctx context.Context, branches []TransBr
 						t.RetryCount++
 						logger.Infof("Retrying branch %s %s %s, t.RetryLimit: %d, t.RetryCount: %d",
 							branches[r.index].BranchID, branches[r.index].Op, branches[r.index].URL, t.RetryLimit, t.RetryCount)
-						copyCtx := CopyContext(ctx)
+						copyCtx := NewAsyncContext(ctx)
 						go asyncExecBranch(copyCtx, r.index)
 						break
 					}

dtmsvr/utils.go

@@ -9,9 +9,7 @@ package dtmsvr
 import (
 	"context"
 	"fmt"
-	"reflect"
 	"time"
-	"unsafe"
 
 	"github.com/dtm-labs/dtm/client/dtmcli/dtmimp"
 	"github.com/dtm-labs/dtm/dtmsvr/config"
@@ -54,55 +52,31 @@ func GetTransGlobal(gid string) *TransGlobal {
 	return &TransGlobal{TransGlobalStore: *trans}
 }
 
-type iface struct {
-	itab, data uintptr
+type asyncCtx struct {
+	parent context.Context
 }
 
-type valueCtx struct {
-	context.Context
-	key, value any
+func (a *asyncCtx) Deadline() (deadline time.Time, ok bool) {
+	return
 }
 
-type cancelCtx struct {
-	context.Context
+func (a *asyncCtx) Done() <-chan struct{} {
+	return nil
 }
 
-type timerCtx struct {
-	cancelCtx
+func (a *asyncCtx) Err() error {
+	return a.parent.Err()
 }
 
-// CopyContext copy context with value and grpc metadata
-// if raw context is nil, return nil
-func CopyContext(ctx context.Context) context.Context {
-	if ctx == nil {
-		return ctx
-	}
-	newCtx := context.Background()
-	kv := make(map[interface{}]interface{})
-	getKeyValues(ctx, kv)
-	for k, v := range kv {
-		newCtx = context.WithValue(newCtx, k, v)
-	}
-	return newCtx
+func (a *asyncCtx) Value(key any) any {
+	return a.parent.Value(key)
 }
 
-func getKeyValues(ctx context.Context, kv map[interface{}]interface{}) {
-	rtType := reflect.TypeOf(ctx).String()
-	if rtType == "*context.emptyCtx" {
-		return
-	}
-	ictx := *(*iface)(unsafe.Pointer(&ctx))
-	if ictx.data == 0 {
-		return
-	}
-	valCtx := (*valueCtx)(unsafe.Pointer(ictx.data))
-	if valCtx.key != nil && valCtx.value != nil && rtType == "*context.valueCtx" {
-		kv[valCtx.key] = valCtx.value
-	}
-	if rtType == "*context.timerCtx" {
-		tCtx := (*timerCtx)(unsafe.Pointer(ictx.data))
-		getKeyValues(tCtx.cancelCtx, kv)
-		return
+// NewAsyncContext create a new async context
+// the context will not be canceled when the parent context is canceled
+func NewAsyncContext(ctx context.Context) context.Context {
+	if ctx == nil {
+		return nil
 	}
-	getKeyValues(valCtx.Context, kv)
+	return &asyncCtx{parent: ctx}
 }

dtmsvr/utils_test.go

@@ -33,21 +33,40 @@ func TestSetNextCron(t *testing.T) {
 	assert.Equal(t, int64(3), tg.getNextCronInterval(cronReset))
 }
 
-func TestCopyContext(t *testing.T) {
+func TestNewAsyncContext(t *testing.T) {
 	var key testContextType = "key"
 	var value testContextType = "value"
 	ctxWithValue := context.WithValue(context.Background(), key, value)
-	newCtx := CopyContext(ctxWithValue)
+	newCtx := NewAsyncContext(ctxWithValue)
 	assert.Equal(t, ctxWithValue.Value(key), newCtx.Value(key))
 
 	var ctx context.Context
-	newCtx = CopyContext(ctx)
+	newCtx = NewAsyncContext(ctx)
 	assert.Nil(t, newCtx)
 }
 
+func TestAsyncContext(t *testing.T) {
+	ctx := context.Background()
+	cancelCtx2, cancel := context.WithCancel(ctx)
+	async := NewAsyncContext(cancelCtx2)
+	cancelCtx3, cancel2 := context.WithCancel(async)
+	defer cancel2()
+	cancel()
+	select {
+	case <-cancelCtx2.Done():
+	default:
+		assert.Failf(t, "context should be canceled", "context should be canceled")
+	}
+	select {
+	case <-cancelCtx3.Done():
+		assert.Failf(t, "context should not be canceled", "context should not be canceled")
+	default:
+	}
+}
+
 type testContextType string
 
-func TestCopyContextRecursive(t *testing.T) {
+func TestAsyncContextRecursive(t *testing.T) {
 	var key testContextType = "key"
 	var key2 testContextType = "key2"
 	var key3 testContextType = "key3"
@@ -64,7 +83,7 @@ func TestCopyContextRecursive(t *testing.T) {
 	defer cancel2()
 	timer2 := context.WithValue(timerCtxx, key2, value2)
 	timer3 := context.WithValue(timer2, key3, value3)
-	newCtx := CopyContext(timer3)
+	newCtx := NewAsyncContext(timer3)
 
 	assert.Equal(t, timer3.Value(nestedKey), newCtx.Value(nestedKey))
 	assert.Equal(t, timer3.Value(key), newCtx.Value(key))
@@ -76,7 +95,7 @@ func TestCopyContextWithMetadata(t *testing.T) {
 	md := metadata.New(map[string]string{"key": "value"})
 	ctx := metadata.NewIncomingContext(context.Background(), md)
 	ctx = metadata.NewOutgoingContext(ctx, md)
-	newCtx := CopyContext(ctx)
+	newCtx := NewAsyncContext(ctx)
 
 	copiedMD, ok := metadata.FromIncomingContext(newCtx)
 	assert.True(t, ok)
@@ -94,6 +113,6 @@ func BenchmarkCopyContext(b *testing.B) {
 	ctx := context.WithValue(context.Background(), key, value)
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		CopyContext(ctx)
+		NewAsyncContext(ctx)
 	}
 }

test/saga_grpc_test.go

@@ -7,6 +7,7 @@
 package test
 
 import (
+	"context"
 	"testing"
 
 	"github.com/dtm-labs/dtm/client/dtmcli"
@@ -94,7 +95,7 @@ func TestSagaGrpcEmptyUrl(t *testing.T) {
 
 // nolint: unparam
 func genSagaGrpc(gid string, outFailed bool, inFailed bool) *dtmgrpc.SagaGrpc {
-	saga := dtmgrpc.NewSagaGrpc(dtmutil.DefaultGrpcServer, gid)
+	saga := dtmgrpc.NewSagaGrpcWithContext(context.Background(), dtmutil.DefaultGrpcServer, gid)
 	req := busi.GenReqGrpc(30, outFailed, inFailed)
 	saga.Add(busi.BusiGrpc+"/busi.Busi/TransOut", busi.BusiGrpc+"/busi.Busi/TransOutRevert", req)
 	saga.Add(busi.BusiGrpc+"/busi.Busi/TransIn", busi.BusiGrpc+"/busi.Busi/TransInRevert", req)