fix pprof

2 files changed, 6 insertions, 0 deletions

diff --git a/internal/eventloop.go b/internal/eventloop.go
index 0fc7387..08b53f6 100644
--- a/internal/eventloop.go
+++ b/internal/eventloop.go
@@ -61,6 +61,7 @@ func (e *eventLoop) stats() string {
 
 func (e *eventLoop) run(ctx context.Context, rawCh <-chan []byte) {
 	if e.flags.FlamegraphEnable {
+		fmt.Println("Collecting flame graph stats, press Ctrl+C to stop")
 		e.flamegraph.Start(ctx)
 	}
 	if e.flags.PprofEnable {
diff --git a/internal/ior.go b/internal/ior.go
index c7e4bf2..5e78645 100644
--- a/internal/ior.go
+++ b/internal/ior.go
@@ -67,6 +67,7 @@ func Run(flags flags.Flags) {
 	}
 	rb.Poll(300)
 
+	pprofDone := make(chan struct{})
 	var cpuProfile, memProfile *os.File
 	if flags.PprofEnable {
 		if cpuProfile, err = os.Create("ior.cpuprofile"); err != nil {
@@ -76,6 +77,8 @@ func Run(flags flags.Flags) {
 			panic(err)
 		}
 		pprof.StartCPUProfile(cpuProfile)
+	} else {
+		close(pprofDone)
 	}
 
 	loop := newEventLoop(flags)
@@ -95,9 +98,11 @@ func Run(flags flags.Flags) {
 			fmt.Println("Stoppig profiling, writing ior.cpuprofile and ior.memprofile")
 			pprof.StopCPUProfile()
 			pprof.WriteHeapProfile(memProfile)
+			close(pprofDone)
 		}
 	}()
 
 	loop.run(ctx, ch)
+	<-pprofDone
 	fmt.Println("Good bye... (unloading BPF tracepoints will take a few seconds...) after", time.Since(startTime))
 }