add --duration flag

2 files changed, 8 insertions, 2 deletions

diff --git a/internal/flags/flags.go b/internal/flags/flags.go
index da62f34..4d375fc 100644
--- a/internal/flags/flags.go
+++ b/internal/flags/flags.go
@@ -16,12 +16,14 @@ type Flags struct {
 	PathFilter       string
 	PprofEnable      bool
 	FlamegraphEnable bool
+	Duration         int
 }
 
 func New() (flags Flags) {
 	flag.IntVar(&flags.PidFilter, "pid", -1, "Filter for processes ID")
 	flag.IntVar(&flags.TidFilter, "tid", -1, "Filter for thread ID")
 	flag.IntVar(&flags.EventMapSize, "mapSize", 4096*16, "BPF FD event ring buffer map size")
+	flag.IntVar(&flags.Duration, "duration", 60, "Probe duration in seconds")
 	flag.StringVar(&flags.CommFilter, "comm", "", "Command to filter for")
 	flag.StringVar(&flags.PathFilter, "path", "", "Path to filter for")
 	flag.BoolVar(&flags.PprofEnable, "pprof", false, "Enable profiling")
diff --git a/internal/ior.go b/internal/ior.go
index 3aa4679..c7e4bf2 100644
--- a/internal/ior.go
+++ b/internal/ior.go
@@ -9,6 +9,7 @@ import (
 	"os/signal"
 	"runtime/pprof"
 	"syscall"
+	"time"
 
 	"ior/internal/flags"
 	"ior/internal/generated/tracepoints"
@@ -79,7 +80,10 @@ func Run(flags flags.Flags) {
 
 	loop := newEventLoop(flags)
 
-	ctx, cancel := context.WithCancel(context.Background())
+	startTime := time.Now()
+	duration := time.Duration(flags.Duration) * time.Second
+	fmt.Println("Probing for", duration)
+	ctx, cancel := context.WithTimeout(context.Background(), duration)
 	c := make(chan os.Signal, 1)
 	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
 
@@ -95,5 +99,5 @@ func Run(flags flags.Flags) {
 	}()
 
 	loop.run(ctx, ch)
-	fmt.Println("Good bye... (unloading BPF tracepoints will take a few seconds...)")
+	fmt.Println("Good bye... (unloading BPF tracepoints will take a few seconds...) after", time.Since(startTime))
 }