more on tree

1 files changed, 29 insertions, 6 deletions

diff --git a/internal/eventloop.go b/internal/eventloop.go
index 7a322a7..f6cd410 100644
--- a/internal/eventloop.go
+++ b/internal/eventloop.go
@@ -24,6 +24,7 @@ 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
@@ -36,6 +37,7 @@ 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),
@@ -58,24 +60,39 @@ 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) {
-	e.startTime = time.Now()
+	var recycle bool
+
+	if e.flags.TreeEnable {
+		e.tree.Start()
+	}
 	if e.flags.PprofEnable {
 		fmt.Println("Profiling, press Ctrl+C to stop")
 		fmt.Println(event.EventStreamHeader)
 	}
+
+	e.startTime = time.Now()
 	for ev := range e.events(rawCh) {
 		switch {
 		case e.flags.TreeEnable:
-			// e.tree.Add(ev)
+			e.tree.Add(ev)
+			recycle = false // tree needs to recycle by itself
 		case e.flags.PprofEnable:
+			recycle = true
 		default:
+			recycle = true
 			fmt.Println(ev.String())
 		}
-		if ev.PrevPair != nil {
-			// Only recycle the previous event, as the current event is the previous event of the next event!
-			ev.PrevPair.Recycle()
-			continue
+		if recycle {
+			ev.RecyclePrev()
 		}
 		e.numSyscallsAfterFilter++
 	}
@@ -87,6 +104,12 @@ func (e *eventLoop) events(rawCh <-chan []byte) <-chan *event.Pair {
 	go func() {
 		defer close(ch)
 		for raw := range rawCh {
+			select {
+			case <-e.done:
+				return
+			default:
+			}
+
 			e.numTracepoints++
 			switch EventType(raw[0]) {
 			case ENTER_OPEN_EVENT: