fix: timer-based cleanup of listenQueues after transient exporter disconnect

When an exporter's Listen() gRPC stream fails with a transient error the
queue is no longer deleted immediately.  Instead a cleanup timer (default
2 min) is scheduled.  If the exporter reconnects before the timer fires,
Listen() cancels the timer and inherits the existing queue — ensuring that
any router token already buffered there by a concurrent Dial() call is
delivered to the reconnected exporter.

On clean shutdown (ctx.Done() — lease ended or server stopping) the timer
is cancelled and the queue is removed straight away, so there is no memory
leak for the normal lifecycle.

Fixes #414

Author: Ambient Code Bot

controller/internal/service/controller_service.go

@@ -69,6 +69,15 @@ import (
 	"google.golang.org/protobuf/proto"
 )
 
+// listenQueueCleanupDelay is how long to keep a listen queue alive after the
+// exporter's stream disconnects with a transient error.  It must be long enough
+// for the exporter to reconnect and consume any Dial token already buffered in
+// the queue, yet short enough to bound the memory overhead of orphaned queues.
+// The exporter's default retry window (5 retries × ~1 s backoff) is well under
+// 30 s, so 2 minutes is a conservative upper bound.
+// Exposed as a var so tests can shorten it without rebuilding.
+var listenQueueCleanupDelay = 2 * time.Minute
+
 // ControllerService exposes a gRPC service
 type ControllerService struct {
 	pb.UnimplementedControllerServiceServer
@@ -79,7 +88,8 @@ type ControllerService struct {
 	Attr          authorization.ContextAttributesGetter
 	ServerOptions []grpc.ServerOption
 	Router        config.Router
-	listenQueues  sync.Map
+	listenQueues  sync.Map // key: leaseName, value: chan *pb.ListenResponse
+	listenTimers  sync.Map // key: leaseName, value: *time.Timer — deferred cleanup after transient disconnect
 }
 
 type wrappedStream struct {
@@ -439,13 +449,34 @@ func (s *ControllerService) Listen(req *pb.ListenRequest, stream pb.ControllerSe
 		return err
 	}
 
+	// Cancel any pending cleanup timer — the exporter is reconnecting and
+	// should inherit the existing queue (which may hold a buffered Dial token).
+	if t, ok := s.listenTimers.LoadAndDelete(leaseName); ok {
+		t.(*time.Timer).Stop()
+	}
+
 	queue, _ := s.listenQueues.LoadOrStore(leaseName, make(chan *pb.ListenResponse, 8))
 	for {
 		select {
 		case <-ctx.Done():
+			// Clean shutdown (lease ended / server stopping): cancel any timer
+			// and remove the queue immediately.
+			if t, ok := s.listenTimers.LoadAndDelete(leaseName); ok {
+				t.(*time.Timer).Stop()
+			}
+			s.listenQueues.Delete(leaseName)
 			return nil
 		case msg := <-queue.(chan *pb.ListenResponse):
 			if err := stream.Send(msg); err != nil {
+				// Transient stream error: schedule deferred cleanup so a
+				// reconnecting exporter can still inherit the queue and consume
+				// any Dial token that was buffered between the error and now.
+				// listenQueueCleanupDelay gives the exporter time to reconnect.
+				t := time.AfterFunc(listenQueueCleanupDelay, func() {
+					s.listenQueues.Delete(leaseName)
+					s.listenTimers.Delete(leaseName)
+				})
+				s.listenTimers.Store(leaseName, t)
 				return err
 			}
 		}

controller/internal/service/controller_service_test.go

@@ -18,6 +18,7 @@ package service
 
 import (
 	"testing"
+	"time"
 
 	jumpstarterdevv1alpha1 "github.com/jumpstarter-dev/jumpstarter-controller/api/v1alpha1"
 	pb "github.com/jumpstarter-dev/jumpstarter-controller/internal/protocol/jumpstarter/v1"
@@ -299,6 +300,106 @@ func TestSyncOnlineConditionWithStatus(t *testing.T) {
 	}
 }
 
+// TestListenQueueTimerCleanup verifies that the listen queue is NOT removed
+// immediately when a transient stream error occurs, giving a reconnecting
+// exporter time to inherit the queue (and any buffered Dial token), and that
+// the queue IS removed once the cleanup timer fires.
+func TestListenQueueTimerCleanup(t *testing.T) {
+	// Shorten the delay so the test completes quickly.
+	original := listenQueueCleanupDelay
+	listenQueueCleanupDelay = 50 * time.Millisecond
+	t.Cleanup(func() { listenQueueCleanupDelay = original })
+
+	svc := &ControllerService{}
+	leaseName := "test-lease"
+
+	// Seed the queue as Listen() would via LoadOrStore.
+	ch := make(chan *pb.ListenResponse, 8)
+	svc.listenQueues.Store(leaseName, ch)
+
+	// Simulate the stream-error path: schedule deferred cleanup.
+	t.Run("queue survives transient error", func(t *testing.T) {
+		timer := time.AfterFunc(listenQueueCleanupDelay, func() {
+			svc.listenQueues.Delete(leaseName)
+			svc.listenTimers.Delete(leaseName)
+		})
+		svc.listenTimers.Store(leaseName, timer)
+
+		// Queue must still be present immediately after the error.
+		if _, ok := svc.listenQueues.Load(leaseName); !ok {
+			t.Fatal("listen queue was removed immediately after stream error — Dial token would be lost")
+		}
+	})
+
+	t.Run("reconnecting exporter cancels cleanup timer", func(t *testing.T) {
+		// Simulate Listen() reconnect: cancel the timer and call LoadOrStore.
+		if raw, ok := svc.listenTimers.LoadAndDelete(leaseName); ok {
+			raw.(*time.Timer).Stop()
+		}
+		got, _ := svc.listenQueues.LoadOrStore(leaseName, make(chan *pb.ListenResponse, 8))
+		if got != ch {
+			t.Fatal("reconnecting Listen() did not inherit the existing queue")
+		}
+
+		// Wait well past the original delay — the queue must still be present
+		// because the timer was stopped.
+		time.Sleep(listenQueueCleanupDelay * 4)
+		if _, ok := svc.listenQueues.Load(leaseName); !ok {
+			t.Fatal("listen queue was removed even though cleanup timer was cancelled")
+		}
+	})
+
+	t.Run("timer fires and removes queue when exporter does not reconnect", func(t *testing.T) {
+		// Re-arm the timer without cancelling it this time.
+		timer := time.AfterFunc(listenQueueCleanupDelay, func() {
+			svc.listenQueues.Delete(leaseName)
+			svc.listenTimers.Delete(leaseName)
+		})
+		svc.listenTimers.Store(leaseName, timer)
+
+		// Wait for the timer to fire.
+		time.Sleep(listenQueueCleanupDelay * 4)
+		if _, ok := svc.listenQueues.Load(leaseName); ok {
+			t.Fatal("listen queue was not removed after cleanup timer fired")
+		}
+	})
+}
+
+// TestListenQueueCleanShutdown verifies that a clean context cancellation
+// (lease end / server stop) removes the queue immediately without waiting for
+// the cleanup timer.
+func TestListenQueueCleanShutdown(t *testing.T) {
+	original := listenQueueCleanupDelay
+	listenQueueCleanupDelay = 2 * time.Minute // keep long — must NOT fire during test
+	t.Cleanup(func() { listenQueueCleanupDelay = original })
+
+	svc := &ControllerService{}
+	leaseName := "test-lease-shutdown"
+
+	ch := make(chan *pb.ListenResponse, 8)
+	svc.listenQueues.Store(leaseName, ch)
+
+	// Arm a timer that should be cancelled before it fires.
+	timer := time.AfterFunc(listenQueueCleanupDelay, func() {
+		svc.listenQueues.Delete(leaseName)
+		svc.listenTimers.Delete(leaseName)
+	})
+	svc.listenTimers.Store(leaseName, timer)
+
+	// Simulate the ctx.Done() path in Listen().
+	if raw, ok := svc.listenTimers.LoadAndDelete(leaseName); ok {
+		raw.(*time.Timer).Stop()
+	}
+	svc.listenQueues.Delete(leaseName)
+
+	if _, ok := svc.listenQueues.Load(leaseName); ok {
+		t.Fatal("listen queue was not removed on clean shutdown")
+	}
+	if _, ok := svc.listenTimers.Load(leaseName); ok {
+		t.Fatal("cleanup timer was not cancelled on clean shutdown")
+	}
+}
+
 // contains checks if substr is contained in s
 func contains(s, substr string) bool {
 	return len(s) >= len(substr) && (s == substr || len(substr) == 0 ||