优化展示

costs.go

@@ -34,11 +34,11 @@ func (c *Costs) addCostWithPath(path string, cost time.Duration) {
 		c.costs[path] = costObj
 	}
 	costObj.CalledCount++
-	costObj.TotalCost += cost
+	costObj.Cost += cost
 }
 
 type Cost struct {
 	Path        string
 	CalledCount int
-	TotalCost   time.Duration
+	Cost        time.Duration
 }

costwhere.go

@@ -8,15 +8,17 @@ import (
 	"github.com/pkg/errors"
 )
 
-func Init(ctx context.Context, topic string) (newCtx context.Context, c *CostWhere) {
+func Init(ctx context.Context) (newCtx context.Context, c *CostWhere) {
 	costs := newCosts()
 
 	newCtx = writeThis(ctx, costs)
 
 	startAt := time.Now()
+	path := getStackInfo(3)
+
 	end := func() {
 		cost := time.Since(startAt)
-		costs.addCost(5, cost)
+		costs.addCostWithPath(path, cost)
 	}
 	c = newCostWhere(costs, end)
 

format.go

@@ -2,17 +2,99 @@ package costwhere
 
 import (
 	"fmt"
+	"strings"
+	"time"
+
+	"github.com/samber/lo"
 )
 
-func formatCost(c *Cost) (cost string) {
-	cost = fmt.Sprintf("%s(%d) %d", c.Path, c.CalledCount, c.TotalCost.Milliseconds())
+func formatCosts(c *Costs) (costs []string) {
+	// 重建栈结构
+	s := newCostStack()
+	for _, cost := range c.costs {
+		path := strings.Split(cost.Path, ";")
+		s.add(path, cost.Cost)
+	}
+
+	costs = s.root.format("")
+
 	return
 }
 
-func formatCosts(c *Costs) (costs []string) {
-	costs = make([]string, 0, len(c.costs))
-	for _, cost := range c.costs {
-		costs = append(costs, formatCost(cost))
+type costStack struct {
+	root *stackFrame
+}
+
+func newCostStack() *costStack {
+	return &costStack{}
+}
+
+func (c *costStack) add(path []string, cost time.Duration) {
+	if len(path) == 0 {
+		return
+	}
+
+	root := path[0]
+
+	if c.root == nil {
+		c.root = newStackFrame(root)
+		return
 	}
+
+	c.root.add(path[1:], cost)
+}
+
+type stackFrame struct {
+	path     string
+	cost     time.Duration
+	children []*stackFrame
+}
+
+func newStackFrame(path string) *stackFrame {
+	return &stackFrame{
+		path:     path,
+		children: make([]*stackFrame, 0),
+	}
+}
+
+func (s *stackFrame) add(childPath []string, cost time.Duration) {
+	if len(childPath) == 0 {
+		s.cost += cost
+		return
+	}
+
+	childRoot := childPath[0]
+
+	child, ok := lo.Find(s.children, func(item *stackFrame) bool {
+		return item.path == childRoot
+	})
+	if ok {
+		child.add(childPath[1:], cost)
+		return
+	}
+
+	newChild := newStackFrame(childRoot)
+	s.children = append(s.children, newChild)
+	newChild.add(childPath[1:], cost)
+
+	if len(childPath) == 1 {
+		s.cost -= cost
+	}
+}
+
+func (s *stackFrame) format(head string) (costs []string) {
+	path := s.path
+	if head != "" {
+		path = head + ";" + path
+	}
+
+	cost := fmt.Sprintf("%s %d", head, s.cost.Milliseconds())
+	costs = append(costs, cost)
+
+	for _, child := range s.children {
+		childCosts := child.format(path)
+		costs = append(costs, childCosts...)
+	}
+
 	return
 }

go.mod

@@ -2,4 +2,9 @@ module github.com/eaglexiang/costwhere-go
 
 go 1.21
 
-require github.com/pkg/errors v0.9.1
+require (
+	github.com/pkg/errors v0.9.1
+	github.com/samber/lo v1.47.0
+)
+
+require golang.org/x/text v0.16.0 // indirect

go.sum

@@ -1,2 +1,6 @@
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
+github.com/samber/lo v1.47.0 h1:z7RynLwP5nbyRscyvcD043DWYoOcYRv3mV8lBeqOCLc=
+github.com/samber/lo v1.47.0/go.mod h1:RmDH9Ct32Qy3gduHQuKJ3gW1fMHAnE/fAzQuf6He5cU=
+golang.org/x/text v0.16.0 h1:a94ExnEXNtEwYLGJSIUxnWoxoRz/ZcCsV63ROupILh4=
+golang.org/x/text v0.16.0/go.mod h1:GhwF1Be+LQoKShO3cGOHzqOgRrGaYc9AvblQOmPVHnI=

runtime.go

@@ -1,8 +1,6 @@
 package costwhere
 
 import (
-	"fmt"
-	"path/filepath"
 	"runtime"
 	"slices"
 	"strings"
@@ -16,15 +14,13 @@ func getStackInfo(skip int) (pathText string) {
 	frames := runtime.CallersFrames(pcs[:depth])
 
 	stacks := make([]runtime.Frame, 0)
-	for f, again := frames.Next(); again; f, again = frames.Next() {
+	for f, ok := frames.Next(); ok; f, ok = frames.Next() {
 		stacks = append(stacks, f)
 	}
 
 	path := make([]string, 0, len(stacks))
-	for _, stack := range stacks {
-		filename := filepath.Base(stack.File)
-		frame := fmt.Sprintf("%s/%s", filename, stack.Function)
-		path = append(path, frame)
+	for _, frame := range stacks {
+		path = append(path, frame.Function)
 	}
 
 	slices.Reverse(path)

tests/01/main.go → tests/01/cmd/main.go

@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/eaglexiang/costwhere-go"
+	"github.com/eaglexiang/costwhere-go/tests/01/x"
 )
 
 func main() {
@@ -19,7 +20,7 @@ func f(sig chan struct{}) {
 	defer func() { sig <- struct{}{} }()
 
 	ctx := context.Background()
-	ctx, cw := costwhere.Init(ctx, "main")
+	ctx, cw := costwhere.Init(ctx)
 	defer func() {
 		stacks, err := cw.EndWithJSON()
 		if err != nil {
@@ -40,27 +41,12 @@ func F0(ctx context.Context) {
 	defer costwhere.Mark(ctx)()
 
 	time.Sleep(100 * time.Millisecond)
-	F2(ctx)
+	x.F2(ctx)
 }
 
 func F1(ctx context.Context) {
 	defer costwhere.Mark(ctx)()
 
 	time.Sleep(1 * time.Second)
-	F2(ctx)
-}
-
-func F2(ctx context.Context) {
-	defer costwhere.Mark(ctx)()
-
-	time.Sleep(100 * time.Millisecond)
-	for i := 0; i < 3; i++ {
-		F3(ctx)
-	}
-}
-
-func F3(ctx context.Context) {
-	defer costwhere.Mark(ctx)()
-
-	time.Sleep(300 * time.Millisecond)
+	x.F2(ctx)
 }

tests/01/x/x.go

@@ -0,0 +1,18 @@
+package x
+
+import (
+	"context"
+	"time"
+
+	"github.com/eaglexiang/costwhere-go"
+	"github.com/eaglexiang/costwhere-go/tests/01/y"
+)
+
+func F2(ctx context.Context) {
+	defer costwhere.Mark(ctx)()
+
+	time.Sleep(100 * time.Millisecond)
+	for i := 0; i < 3; i++ {
+		y.F3(ctx)
+	}
+}

tests/01/y/y.go

@@ -0,0 +1,14 @@
+package y
+
+import (
+	"context"
+	"time"
+
+	"github.com/eaglexiang/costwhere-go"
+)
+
+func F3(ctx context.Context) {
+	defer costwhere.Mark(ctx)()
+
+	time.Sleep(300 * time.Millisecond)
+}