优化 End 的调用方式

eaglexiang · Dec 12, 2024 · adc8554 · adc8554
1 parent 099dbef
commit adc8554
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 10 deletions.
diff --git a/costwhere.go b/costwhere.go
@@ -34,11 +34,12 @@ func (s *CostWhere) End() (stacks []string) {
 }
 
 // StartStack 开始一个栈帧
-func Begin(ctx context.Context, topic string) (newCtx context.Context) {
+func Begin(ctx context.Context, topic string) (newCtx context.Context, end func()) {
 	// 读取父级栈帧
 	parent, ok := readThis(ctx)
 	if !ok {
 		newCtx = ctx
+		end = func() {}
 		return
 	}
 
@@ -48,11 +49,12 @@ func Begin(ctx context.Context, topic string) (newCtx context.Context) {
 
 	// 将本级栈帧写入 ctx
 	newCtx = writeThis(ctx, newLayer)
+	end = func() { endStack(newCtx) }
 
 	return
 }
 
-func End(ctx context.Context) {
+func endStack(ctx context.Context) {
 	// 读取栈帧
 	stackLayer, ok := readThis(ctx)
 	if !ok {

diff --git a/tests/01/main.go b/tests/01/main.go
@@ -30,30 +30,30 @@ func main() {
 }
 
 func F0(ctx context.Context) {
-	ctx = costwhere.Begin(ctx, "F0")
-	defer costwhere.End(ctx)
+	ctx, end := costwhere.Begin(ctx, "F0")
+	defer end()
 	time.Sleep(100 * time.Millisecond)
 	F2(ctx)
 }
 
 func F1(ctx context.Context) {
-	ctx = costwhere.Begin(ctx, "F1")
-	defer costwhere.End(ctx)
+	ctx, end := costwhere.Begin(ctx, "F1")
+	defer end()
 	time.Sleep(1 * time.Second)
 	F2(ctx)
 }
 
 func F2(ctx context.Context) {
-	ctx = costwhere.Begin(ctx, "F2")
-	defer costwhere.End(ctx)
+	ctx, end := costwhere.Begin(ctx, "F2")
+	defer end()
 	time.Sleep(100 * time.Millisecond)
 	F3(ctx)
 	F3(ctx)
 	F3(ctx)
 }
 
 func F3(ctx context.Context) {
-	ctx = costwhere.Begin(ctx, "F3")
-	defer costwhere.End(ctx)
+	_, end := costwhere.Begin(ctx, "F3")
+	defer end()
 	time.Sleep(300 * time.Millisecond)
 }