新增 Mark 方法

eaglexiang · Dec 13, 2024 · c11b9be · c11b9be
1 parent 9af9e43
commit c11b9be
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 8 deletions.
diff --git a/costwhere.go b/costwhere.go
@@ -33,6 +33,16 @@ func (s *CostWhere) End() (stacks []string) {
 	return
 }
 
+func Mark(ctx *context.Context, topic string) (end func()) {
+	if ctx == nil {
+		end = func() {}
+		return
+	}
+	newCtx, end := Begin(*ctx, topic)
+	*ctx = newCtx
+	return
+}
+
 // StartStack 开始一个栈帧
 func Begin(ctx context.Context, topic string) (newCtx context.Context, end func()) {
 	// 读取父级栈帧

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