兼容火焰图栈协议

README.md

@@ -1,2 +1,3 @@
 # costwhere
 
+show where your project cost time

cmd/main.go

@@ -0,0 +1,23 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+	"os/exec"
+)
+
+func main() {
+	src := flag.String("if", "", "source file")
+	flag.Parse()
+
+	if *src == "" {
+		fmt.Println("source file cannot be empty")
+		return
+	}
+
+	cmd := fmt.Sprintf("cat %s | flamegraph.pl > flamegraph.svg", *src)
+	err := exec.Command("bash", "-c", cmd).Run()
+	if err != nil {
+		fmt.Println(err)
+	}
+}

costwhere.go

@@ -0,0 +1,63 @@
+package costwhere
+
+import (
+	"context"
+)
+
+func Init(ctx context.Context, topic string) (newCtx context.Context, s *CostWhere) {
+	// 构造根帧
+	newLayer := newStackLayer(topic)
+	// 构造栈
+	s = newCostWhere(newLayer)
+
+	// 将栈帧写入 ctx
+	newCtx = writeThis(ctx, newLayer)
+
+	return
+}
+
+type CostWhere struct {
+	Root *StackLayer
+}
+
+func newCostWhere(root *StackLayer) *CostWhere {
+	s := &CostWhere{
+		Root: root,
+	}
+	return s
+}
+
+func (s *CostWhere) End() (stacks []string) {
+	s.Root.Stop()
+	stacks = s.Root.Format("")
+	return
+}
+
+// StartStack 开始一个栈帧
+func Begin(ctx context.Context, topic string) (newCtx context.Context) {
+	// 读取父级栈帧
+	parent, ok := readThis(ctx)
+	if !ok {
+		return
+	}
+
+	// 写入本级栈帧
+	newLayer := newStackLayer(topic)
+	parent.AddChild(newLayer)
+
+	// 将本级栈帧写入 ctx
+	newCtx = writeThis(ctx, newLayer)
+
+	return
+}
+
+func End(ctx context.Context) {
+	// 读取栈帧
+	stackLayer, ok := readThis(ctx)
+	if !ok {
+		return
+	}
+
+	// 结束栈帧
+	stackLayer.Stop()
+}

ctx.go

@@ -0,0 +1,20 @@
+package costwhere
+
+import "context"
+
+type ctxKeyType string
+
+const (
+	thisLayerKey ctxKeyType = "thisLayerCtx"
+)
+
+func readThis(ctx context.Context) (s *StackLayer, ok bool) {
+	val := ctx.Value(thisLayerKey)
+	s, ok = val.(*StackLayer)
+	return
+}
+
+func writeThis(ctx context.Context, s *StackLayer) (dst context.Context) {
+	dst = context.WithValue(ctx, thisLayerKey, s)
+	return
+}

go.mod

@@ -0,0 +1,3 @@
+module costwhere
+
+go 1.21.3

stack.go

@@ -0,0 +1,54 @@
+package costwhere
+
+import (
+	"strconv"
+	"sync"
+	"time"
+)
+
+type StackLayer struct {
+	l *sync.Mutex
+
+	name  string
+	times TimeSeg
+
+	children []*StackLayer
+}
+
+func newStackLayer(name string) *StackLayer {
+	return &StackLayer{
+		l: new(sync.Mutex),
+
+		name:  name,
+		times: newTimeSeg(),
+
+		children: make([]*StackLayer, 0),
+	}
+}
+
+func (s *StackLayer) AddChild(child *StackLayer) {
+	s.l.Lock()
+	defer s.l.Unlock()
+	s.children = append(s.children, child)
+}
+
+func (s *StackLayer) Stop() (end time.Time) {
+	end = time.Now()
+	s.times.End = end
+	return
+}
+
+func (s *StackLayer) Format(head string) (lines []string) {
+	if head != "" {
+		head += ";"
+	}
+	head += s.name
+	line := head + " " + strconv.FormatInt(s.times.Cost().Milliseconds(), 10)
+	lines = append(lines, line)
+
+	for _, chiled := range s.children {
+		lines = append(lines, chiled.Format(head)...)
+	}
+
+	return
+}

tests/01/main.go

@@ -0,0 +1,56 @@
+package main
+
+import (
+	"bytes"
+	"context"
+	"costwhere"
+	"fmt"
+	"os"
+	"time"
+)
+
+func main() {
+	ctx := context.Background()
+	ctx, cw := costwhere.Init(ctx, "main")
+	defer func() {
+		stacks := cw.End()
+		buf := bytes.NewBuffer(nil)
+		for _, line := range stacks {
+			buf.WriteString(line + "\n")
+		}
+		err := os.WriteFile("costwhere.log", buf.Bytes(), 0644)
+		if err != nil {
+			fmt.Println(err)
+		}
+	}()
+
+	F0(ctx)
+	F1(ctx)
+}
+
+func F0(ctx context.Context) {
+	ctx = costwhere.Begin(ctx, "F0")
+	defer costwhere.End(ctx)
+	time.Sleep(100 * time.Millisecond)
+	F2(ctx)
+}
+
+func F1(ctx context.Context) {
+	ctx = costwhere.Begin(ctx, "F1")
+	defer costwhere.End(ctx)
+	time.Sleep(1 * time.Second)
+	F2(ctx)
+}
+
+func F2(ctx context.Context) {
+	ctx = costwhere.Begin(ctx, "F2")
+	defer costwhere.End(ctx)
+	time.Sleep(100 * time.Millisecond)
+	F3(ctx)
+}
+
+func F3(ctx context.Context) {
+	ctx = costwhere.Begin(ctx, "F3")
+	defer costwhere.End(ctx)
+	time.Sleep(300 * time.Millisecond)
+}

timeseg.go

@@ -0,0 +1,20 @@
+package costwhere
+
+import "time"
+
+type TimeSeg struct {
+	Start time.Time
+	End   time.Time
+}
+
+func newTimeSeg() TimeSeg {
+	now := time.Now()
+	return TimeSeg{
+		Start: now,
+		End:   now,
+	}
+}
+func (t *TimeSeg) Cost() (d time.Duration) {
+	d = t.End.Sub(t.Start)
+	return
+}