From 0fb84813a53d4b1f331f6b87b19298a71cfc7743 Mon Sep 17 00:00:00 2001
From: Pim van Pelt <pim@ipng.ch>
Date: Mon, 23 Mar 2026 22:18:05 +0100
Subject: [PATCH] Enable lame-duck on exit, so systemctl can 'restart'
 collectors

---
 cmd/collector/main.go | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/cmd/collector/main.go b/cmd/collector/main.go
index 891e4ec..e6c3229 100644
--- a/cmd/collector/main.go
+++ b/cmd/collector/main.go
@@ -61,7 +61,22 @@ func main() {
 
 	<-ctx.Done()
 	log.Printf("collector: shutting down")
-	grpcServer.GracefulStop()
+
+	// GracefulStop waits for all RPCs to finish. StreamSnapshots subscribers
+	// (e.g. the aggregator) hold a stream open indefinitely, so we give it a
+	// short window and then force-stop to avoid hanging systemctl stop/restart.
+	stopped := make(chan struct{})
+	go func() {
+		grpcServer.GracefulStop()
+		close(stopped)
+	}()
+	select {
+	case <-stopped:
+	case <-time.After(5 * time.Second):
+		log.Printf("collector: graceful stop timed out, forcing stop")
+		grpcServer.Stop()
+	}
+
 	close(ch)
 }