Add watch mode for dynamic flamegraph updates

7 files changed, 214 insertions, 10 deletions

diff --git a/README.md b/README.md
index af824ae..5826ec1 100644
--- a/README.md
+++ b/README.md
@@ -85,6 +85,14 @@ For experimental WebAssembly frontends, you can also emit a flamegraph JSON tree
 
 This writes `<trace>.<fields>-by-<count>.json` next to the SVG.
 
+To keep the served flamegraph changing as the `.ior.zst` file is updated, enable watch mode:
+
+```sh
+./ior -ior=trace.ior.zst -iorWatchInterval=2s
+```
+
+This polls the input file for modifications, regenerates SVG/JSON outputs, and serves an auto-reloading viewer at `/`.
+
 ## Live Flamegraph Mode
 
 Run live mode (requires root privileges):
diff --git a/internal/flags/flags.go b/internal/flags/flags.go
index bc1654b..b2d9dce 100644
--- a/internal/flags/flags.go
+++ b/internal/flags/flags.go
@@ -71,9 +71,10 @@ type Flags struct {
 	TUIExportEnable  bool
 
 	// To convert ior data into native SVG format
-	IorDataFile     string
-	CollapsedFields []string
-	CountField      string
+	IorDataFile      string
+	IorWatchInterval time.Duration
+	CollapsedFields  []string
+	CountField       string
 }
 
 func Get() Flags {
@@ -130,6 +131,8 @@ func parse() error {
 	flag.BoolVar(&singleton.TUIExportEnable, "tuiExport", true, "Enable writing TUI snapshot export files")
 
 	flag.StringVar(&singleton.IorDataFile, "ior", "", "IOR data file to convert into native SVG flamegraph")
+	flag.DurationVar(&singleton.IorWatchInterval, "iorWatchInterval", 0,
+		"In -ior mode, poll input file for changes and regenerate outputs; also enables auto-reloading viewer")
 	fields := flag.String("fields", "",
 		fmt.Sprintf("Comma separated list of fields to collapse, valid are: %v", validCollapsedFields))
 	flag.StringVar(&singleton.CountField, "count", "count",
diff --git a/internal/flags/flags_test.go b/internal/flags/flags_test.go
index 11414d8..81d01ad 100644
--- a/internal/flags/flags_test.go
+++ b/internal/flags/flags_test.go
@@ -114,6 +114,16 @@ func TestParseFlamegraphJSONFlag(t *testing.T) {
 	}
 }
 
+func TestParseIorWatchIntervalFlag(t *testing.T) {
+	cfg, err := parseForTest(t, "-iorWatchInterval", "2s")
+	if err != nil {
+		t.Fatalf("parse returned error: %v", err)
+	}
+	if cfg.IorWatchInterval != 2*time.Second {
+		t.Fatalf("ior watch interval = %v, want %v", cfg.IorWatchInterval, 2*time.Second)
+	}
+}
+
 func TestParseDefaultCollapsedFieldsOrder(t *testing.T) {
 	cfg, err := parseForTest(t)
 	if err != nil {
diff --git a/internal/flamegraph/webserver.go b/internal/flamegraph/webserver.go
index 2bf6286..c472dfb 100644
--- a/internal/flamegraph/webserver.go
+++ b/internal/flamegraph/webserver.go
@@ -43,6 +43,30 @@ func ServeSVG(svgFile string) error {
 	})
 }
 
+// ServeSVGAutoReload serves an SVG viewer page that periodically reloads the SVG.
+//
+// The SVG file itself is still served directly at its absolute URL path, while "/"
+// serves a small HTML wrapper that appends a cache-busting query parameter on each
+// refresh interval to pick up newly written SVG content.
+func ServeSVGAutoReload(svgFile string, refreshInterval time.Duration) error {
+	if refreshInterval <= 0 {
+		return fmt.Errorf("refresh interval must be > 0")
+	}
+
+	absPath, err := filepath.Abs(svgFile)
+	if err != nil {
+		return fmt.Errorf("resolve svg path: %w", err)
+	}
+	urlPath := buildURLPath(absPath)
+	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+	defer stop()
+
+	mux := buildSVGAutoReloadHandler(absPath, urlPath, refreshInterval)
+	return runServer(ctx, mux, defaultServerTimeouts, func(hostname string, port int) {
+		printServerURL(hostname, port, "/")
+	})
+}
+
 func buildURLPath(absPath string) string {
 	urlPath := filepath.ToSlash(absPath)
 	if !strings.HasPrefix(urlPath, "/") {
@@ -62,6 +86,46 @@ func buildSVGHandler(absPath, urlPath string) *http.ServeMux {
 	return mux
 }
 
+func buildSVGAutoReloadHandler(absPath, urlPath string, refreshInterval time.Duration) *http.ServeMux {
+	intervalMs := refreshInterval.Milliseconds()
+	mux := http.NewServeMux()
+	mux.HandleFunc("/", func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Content-Type", "text/html; charset=utf-8")
+		_, _ = fmt.Fprintf(w, `<!doctype html>
+<html>
+<head>
+  <meta charset="utf-8"/>
+  <meta name="viewport" content="width=device-width, initial-scale=1"/>
+  <title>I/O Flamegraph (Auto-Reload)</title>
+  <style>
+    body { margin: 0; font-family: monospace; }
+    .bar { padding: 8px 12px; border-bottom: 1px solid #ddd; }
+    .viewer { width: 100%%; height: calc(100vh - 42px); border: 0; display: block; }
+  </style>
+</head>
+<body>
+  <div class="bar">
+    Auto-refresh every %d ms.
+    <button type="button" onclick="refreshNow()">Refresh now</button>
+  </div>
+  <iframe id="fg" class="viewer" src="%s"></iframe>
+  <script>
+    const base = %q;
+    function refreshNow() {
+      document.getElementById("fg").src = base + "?t=" + Date.now();
+    }
+    setInterval(refreshNow, %d);
+  </script>
+</body>
+</html>
+`, intervalMs, urlPath, urlPath, intervalMs)
+	})
+	mux.HandleFunc(urlPath, func(w http.ResponseWriter, r *http.Request) {
+		http.ServeFile(w, r, absPath)
+	})
+	return mux
+}
+
 func listenRandomPort() (net.Listener, error) {
 	listener, err := net.Listen("tcp", ":0")
 	if err != nil {
diff --git a/internal/flamegraph/webserver_autoreload_test.go b/internal/flamegraph/webserver_autoreload_test.go
new file mode 100644
index 0000000..ed4c907
--- /dev/null
+++ b/internal/flamegraph/webserver_autoreload_test.go
@@ -0,0 +1,37 @@
+package flamegraph
+
+import (
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+)
+
+func TestBuildSVGAutoReloadHandlerServesViewerPage(t *testing.T) {
+	mux := buildSVGAutoReloadHandler("/tmp/fake.svg", "/tmp/fake.svg", 2*time.Second)
+
+	rec := httptest.NewRecorder()
+	req := httptest.NewRequest("GET", "http://localhost/", nil)
+	mux.ServeHTTP(rec, req)
+
+	if rec.Code != 200 {
+		t.Fatalf("status code = %d, want 200", rec.Code)
+	}
+	if got := rec.Header().Get("Content-Type"); !strings.HasPrefix(got, "text/html") {
+		t.Fatalf("content type = %q, want text/html", got)
+	}
+	body := rec.Body.String()
+	if !strings.Contains(body, "Auto-refresh every 2000 ms.") {
+		t.Fatalf("viewer page missing refresh interval, body=%q", body)
+	}
+	if !strings.Contains(body, `id="fg"`) {
+		t.Fatalf("viewer page missing iframe, body=%q", body)
+	}
+}
+
+func TestServeSVGAutoReloadRejectsNonPositiveInterval(t *testing.T) {
+	err := ServeSVGAutoReload("ignored.svg", 0)
+	if err == nil {
+		t.Fatal("expected error for non-positive interval")
+	}
+}
diff --git a/internal/ior.go b/internal/ior.go
index 865e51d..a6fdbc4 100644
--- a/internal/ior.go
+++ b/internal/ior.go
@@ -121,19 +121,25 @@ func Run() error {
 	var noTraceRun bool
 
 	if iorFile != "" {
+		if cfg.IorWatchInterval < 0 {
+			return errors.New("-iorWatchInterval must be >= 0")
+		}
 		noTraceRun = true
 		native := flamegraph.NewNativeSVG(cfg.CollapsedFields, cfg.CountField)
-		svgFile, err := native.WriteSVGFromFile(iorFile)
+		svgFile, err := writeIorOutputs(native, iorFile, cfg.FlamegraphJSON)
 		if err != nil {
 			return err
 		}
-		if cfg.FlamegraphJSON {
-			if _, err := native.WriteJSONFromFile(iorFile); err != nil {
-				return err
-			}
-		}
 
-		if err := flamegraph.ServeSVG(svgFile); err != nil {
+		done := make(chan struct{})
+		defer close(done)
+		if cfg.IorWatchInterval > 0 {
+			go watchIorOutputs(done, cfg.IorWatchInterval, iorFile, native, cfg.FlamegraphJSON)
+			err = flamegraph.ServeSVGAutoReload(svgFile, cfg.IorWatchInterval)
+		} else {
+			err = flamegraph.ServeSVG(svgFile)
+		}
+		if err != nil {
 			return err
 		}
 	}
@@ -144,6 +150,54 @@ func Run() error {
 	return dispatchRun(cfg)
 }
 
+func writeIorOutputs(native flamegraph.NativeSVG, iorFile string, writeJSON bool) (string, error) {
+	svgFile, err := native.WriteSVGFromFile(iorFile)
+	if err != nil {
+		return "", err
+	}
+	if !writeJSON {
+		return svgFile, nil
+	}
+	if _, err := native.WriteJSONFromFile(iorFile); err != nil {
+		return "", err
+	}
+	return svgFile, nil
+}
+
+func watchIorOutputs(done <-chan struct{}, interval time.Duration, iorFile string,
+	native flamegraph.NativeSVG, writeJSON bool) {
+
+	ticker := time.NewTicker(interval)
+	defer ticker.Stop()
+	lastMod := fileModTime(iorFile)
+
+	for {
+		select {
+		case <-done:
+			return
+		case <-ticker.C:
+			mod := fileModTime(iorFile)
+			if !mod.After(lastMod) {
+				continue
+			}
+			if _, err := writeIorOutputs(native, iorFile, writeJSON); err != nil {
+				_, _ = fmt.Printf("Failed to refresh flamegraph outputs: %v\n", err)
+				continue
+			}
+			lastMod = mod
+			_, _ = fmt.Printf("Refreshed flamegraph outputs at %s\n", time.Now().Format(time.RFC3339))
+		}
+	}
+}
+
+func fileModTime(path string) time.Time {
+	stat, err := os.Stat(path)
+	if err != nil {
+		return time.Time{}
+	}
+	return stat.ModTime()
+}
+
 func dispatchRun(cfg flags.Flags) error {
 	if err := validateRunConfig(cfg); err != nil {
 		return err
@@ -158,6 +212,12 @@ func validateRunConfig(cfg flags.Flags) error {
 	if cfg.LiveFlamegraph && cfg.FlamegraphEnable {
 		return errors.New("-live and -flamegraph are mutually exclusive")
 	}
+	if cfg.IorWatchInterval > 0 && cfg.IorDataFile == "" {
+		return errors.New("-iorWatchInterval requires -ior")
+	}
+	if cfg.IorWatchInterval < 0 {
+		return errors.New("-iorWatchInterval must be >= 0")
+	}
 	return nil
 }
 
diff --git a/internal/ior_mode_test.go b/internal/ior_mode_test.go
index bac54bd..bbca555 100644
--- a/internal/ior_mode_test.go
+++ b/internal/ior_mode_test.go
@@ -176,6 +176,28 @@ func TestDispatchRunRejectsLiveAndFlamegraph(t *testing.T) {
 	}
 }
 
+func TestValidateRunConfigRejectsIorWatchWithoutIor(t *testing.T) {
+	cfg := flags.Flags{IorWatchInterval: time.Second}
+	err := validateRunConfig(cfg)
+	if err == nil {
+		t.Fatalf("expected error for -iorWatchInterval without -ior")
+	}
+	if err.Error() != "-iorWatchInterval requires -ior" {
+		t.Fatalf("unexpected error: %v", err)
+	}
+}
+
+func TestValidateRunConfigRejectsNegativeIorWatchInterval(t *testing.T) {
+	cfg := flags.Flags{IorWatchInterval: -time.Second}
+	err := validateRunConfig(cfg)
+	if err == nil {
+		t.Fatalf("expected error for negative -iorWatchInterval")
+	}
+	if err.Error() != "-iorWatchInterval must be >= 0" {
+		t.Fatalf("unexpected error: %v", err)
+	}
+}
+
 func TestRunTraceWithContextRequiresRoot(t *testing.T) {
 	origGetEUID := getEUID
 	defer func() { getEUID = origGetEUID }()