1 files changed, 39 insertions, 30 deletions
diff --git a/internal/tree/tree.go b/internal/tree/tree.go
index a7c5c20..9b7ac16 100644
--- a/internal/tree/tree.go
+++ b/internal/tree/tree.go
@@ -1,6 +1,7 @@
 package tree
 
 import (
+	"context"
 	"fmt"
 	"ior/internal/event"
 	"ior/internal/generated/types"
@@ -19,49 +20,54 @@ type Tree struct {
 	// Collapsed flamegraph stats collector
 	collapsed map[string]map[types.TraceId]counter
 	inCh      chan *event.Pair
-	Finished  chan struct{}
+	Done      chan struct{}
 }
 
 func New() Tree {
 	return Tree{
 		collapsed: make(map[string]map[types.TraceId]counter),
 		inCh:      make(chan *event.Pair, 4096),
-		Finished:  make(chan struct{}),
+		Done:      make(chan struct{}),
 	}
 }
 
-func (t Tree) Add(ev *event.Pair) {
-	t.inCh <- ev
-}
-
-func (t Tree) Start() {
+func (t Tree) Start(ctx context.Context) {
 	go func() {
-		for ev := range t.inCh {
-			pathname := path.Dir(ev.File.Name())
-			pathMap, ok := t.collapsed[pathname]
-			if !ok {
-				pathMap = make(map[types.TraceId]counter)
-			}
+		for {
+			select {
+			case ev := <-t.inCh:
+				pathname := path.Dir(ev.File.Name())
+				pathMap, ok := t.collapsed[pathname]
+				if !ok {
+					pathMap = make(map[types.TraceId]counter)
+				}
 
-			traceId := ev.EnterEv.GetTraceId()
-			cnt := pathMap[traceId]
-			cnt.count++
-			cnt.duration += ev.Duration
-			pathMap[traceId] = cnt
+				traceId := ev.EnterEv.GetTraceId()
+				cnt := pathMap[traceId]
+				cnt.count++
+				cnt.duration += ev.Duration
+				pathMap[traceId] = cnt
 
-			t.collapsed[pathname] = pathMap
-			ev.RecyclePrev()
+				t.collapsed[pathname] = pathMap
+				ev.RecyclePrev()
+			default:
+				select {
+				case <-ctx.Done():
+					fmt.Println("Tree processed last event")
+					t.mustDump("ior.collapsed")
+					close(t.Done)
+					return
+				default:
+				}
+			}
 		}
-		fmt.Println("Tree processed last event")
-		t.mustDump("ior.collapsed")
-		close(t.Finished)
-		fmt.Println("DEBUG")
 	}()
 }
 
-// TODO: By count and by duration
-// TODO: Reverse by syscall vs path
-// TODO: Allow full path?
+func (t Tree) Add(ev *event.Pair) {
+	t.inCh <- ev
+}
+
 func (t Tree) mustDump(outfile string) {
 	fmt.Println("Writing", outfile)
 	file, err := os.Create(outfile)
@@ -73,13 +79,16 @@ func (t Tree) mustDump(outfile string) {
 	for path, value := range t.collapsed {
 		var sb strings.Builder
 
-		for _, part := range strings.Split(path, "/") {
-			sb.WriteString("/")
+		for i, part := range strings.Split(path, "/") {
+			if i > 1 {
+				sb.WriteString(";")
+				sb.WriteString("/")
+			}
 			sb.WriteString(part)
 		}
 
 		for traceId, cnt := range value {
-			_, err := fmt.Fprintf(file, "%s;%s %v\n", sb.String(), traceId, cnt.count)
+			_, err := fmt.Fprintf(file, "%s;[%s] %v\n", sb.String(), traceId, cnt.count)
 			if err != nil {
 				panic(err)
 			}