initial flames

3 files changed, 55 insertions, 49 deletions

diff --git a/internal/eventloop.go b/internal/eventloop.go
index f6cd410..9d243d9 100644
--- a/internal/eventloop.go
+++ b/internal/eventloop.go
@@ -3,6 +3,7 @@ package internal
 import "C"
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -24,7 +25,6 @@ type eventLoop struct {
 	comms     map[uint32]string      // Program or thread name of the current Tid.
 	prevPairs map[uint32]*event.Pair // Previous event (to calculate time differences between two events)
 	tree      tree.Tree              // Storing all paths in a tree structure for analysis
-	done      chan struct{}
 
 	// Statistics
 	numTracepoints          uint
@@ -37,7 +37,6 @@ type eventLoop struct {
 func newEventLoop(flags flags.Flags) *eventLoop {
 	return &eventLoop{
 		flags:     flags,
-		done:      make(chan struct{}),
 		filter:    newEventFilter(flags),
 		enterEvs:  make(map[uint32]*event.Pair),
 		files:     make(map[int32]file.File),
@@ -60,19 +59,11 @@ func (e *eventLoop) stats() string {
 			e.numSyscallsAfterFilter, float64(e.numSyscallsAfterFilter)/duration.Seconds())
 }
 
-func (e *eventLoop) stop() {
-	close(e.done)
-	if e.flags.TreeEnable {
-		fmt.Println("Waiting for tree to finish")
-		<-e.tree.Finished
-	}
-}
-
-func (e *eventLoop) run(rawCh <-chan []byte) {
+func (e *eventLoop) run(ctx context.Context, rawCh <-chan []byte) {
 	var recycle bool
 
 	if e.flags.TreeEnable {
-		e.tree.Start()
+		e.tree.Start(ctx)
 	}
 	if e.flags.PprofEnable {
 		fmt.Println("Profiling, press Ctrl+C to stop")
@@ -80,7 +71,7 @@ func (e *eventLoop) run(rawCh <-chan []byte) {
 	}
 
 	e.startTime = time.Now()
-	for ev := range e.events(rawCh) {
+	for ev := range e.events(ctx, rawCh) {
 		switch {
 		case e.flags.TreeEnable:
 			e.tree.Add(ev)
@@ -96,16 +87,21 @@ func (e *eventLoop) run(rawCh <-chan []byte) {
 		}
 		e.numSyscallsAfterFilter++
 	}
+
+	if e.flags.TreeEnable {
+		fmt.Println("Waiting for tree")
+		<-e.tree.Done
+	}
 }
 
-func (e *eventLoop) events(rawCh <-chan []byte) <-chan *event.Pair {
+func (e *eventLoop) events(ctx context.Context, rawCh <-chan []byte) <-chan *event.Pair {
 	ch := make(chan *event.Pair)
 
 	go func() {
 		defer close(ch)
 		for raw := range rawCh {
 			select {
-			case <-e.done:
+			case <-ctx.Done():
 				return
 			default:
 			}
diff --git a/internal/ior.go b/internal/ior.go
index 76cf17d..62e25ef 100644
--- a/internal/ior.go
+++ b/internal/ior.go
@@ -3,6 +3,7 @@ package internal
 import "C"
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"os/signal"
@@ -78,9 +79,11 @@ func Run(flags flags.Flags) {
 
 	loop := newEventLoop(flags)
 
+	ctx, cancel := context.WithCancel(context.Background())
 	c := make(chan os.Signal, 1)
 	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
 	go func() {
+		defer cancel()
 		<-c
 		fmt.Println(loop.stats())
 		if flags.PprofEnable {
@@ -88,10 +91,8 @@ func Run(flags flags.Flags) {
 			pprof.StopCPUProfile()
 			pprof.WriteHeapProfile(memProfile)
 		}
-		loop.stop()
-		fmt.Println("Good bye... (unloading BPF tracepoints will take a few seconds...)")
-		os.Exit(0)
 	}()
 
-	loop.run(ch)
+	loop.run(ctx, ch)
+	fmt.Println("Good bye... (unloading BPF tracepoints will take a few seconds...)")
 }
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)
 			}