From 2726af38fd9abc3785795e76e3c6a27546529a1b Mon Sep 17 00:00:00 2001 From: Makonike Date: Tue, 5 Sep 2023 17:58:44 +0800 Subject: [PATCH] fix: using async context instead of CopyContext --- dtmsvr/trans_process.go | 2 +- dtmsvr/trans_type_msg.go | 2 +- dtmsvr/trans_type_saga.go | 4 +-- dtmsvr/utils.go | 58 +++++++++++---------------------------- dtmsvr/utils_test.go | 33 +++++++++++++++++----- test/saga_grpc_test.go | 3 +- 6 files changed, 48 insertions(+), 54 deletions(-) diff --git a/dtmsvr/trans_process.go b/dtmsvr/trans_process.go index d39898b..267b5d6 100644 --- a/dtmsvr/trans_process.go +++ b/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) { diff --git a/dtmsvr/trans_type_msg.go b/dtmsvr/trans_type_msg.go index 0def41d..016d7f5 100644 --- a/dtmsvr/trans_type_msg.go +++ b/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) diff --git a/dtmsvr/trans_type_saga.go b/dtmsvr/trans_type_saga.go index cd60a56..e984b27 100644 --- a/dtmsvr/trans_type_saga.go +++ b/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 } diff --git a/dtmsvr/utils.go b/dtmsvr/utils.go index 16f28f2..fdf7d81 100644 --- a/dtmsvr/utils.go +++ b/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} } diff --git a/dtmsvr/utils_test.go b/dtmsvr/utils_test.go index 3349ad0..060c1c2 100644 --- a/dtmsvr/utils_test.go +++ b/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) } } diff --git a/test/saga_grpc_test.go b/test/saga_grpc_test.go index 17806a3..e97a8d7 100644 --- a/test/saga_grpc_test.go +++ b/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)