1 files changed, 15 insertions, 3 deletions
diff --git a/internal/ior.go b/internal/ior.go
index a910fc0..a19d2f7 100644
--- a/internal/ior.go
+++ b/internal/ior.go
@@ -272,9 +272,17 @@ func runTraceWithContext(parentCtx context.Context, started chan<- struct{}, con
 			origPrintCb(ep)
 		}
 	}
-	duration := time.Duration(flags.Get().Duration) * time.Second
-	logln("Probing for", duration)
-	ctx, cancel := context.WithTimeout(parentCtx, duration)
+	cfg := flags.Get()
+	ctx := parentCtx
+	cancel := func() {}
+	if shouldAutoStopByDuration(cfg) {
+		duration := time.Duration(cfg.Duration) * time.Second
+		logln("Probing for", duration)
+		ctx, cancel = context.WithTimeout(parentCtx, duration)
+	} else {
+		logln("Probing until stopped...")
+		ctx, cancel = context.WithCancel(parentCtx)
+	}
 	defer cancel()
 
 	signalCh := make(chan os.Signal, 1)
@@ -318,3 +326,7 @@ func signalTraceStarted(started chan<- struct{}) {
 	}
 	close(started)
 }
+
+func shouldAutoStopByDuration(cfg flags.Flags) bool {
+	return cfg.PlainMode || cfg.FlamegraphEnable || cfg.PprofEnable
+}