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

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,51 @@ 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.
+				//
+				// Known limitation: there is a narrow race at timer expiry --
+				// if Dial() calls LoadOrStore and obtains the queue just before
+				// this callback deletes it, the buffered token is lost.  The
+				// window is microseconds wide and only opens after
+				// listenQueueCleanupDelay has fully elapsed (i.e. the exporter
+				// has been gone for 2+ minutes), at which point the lease is
+				// typically being reclaimed anyway.  Replacing sync.Map with a
+				// mutex-protected map would close this race by holding the lock
+				// across both the timer check and the queue access.
+				if old, ok := s.listenTimers.LoadAndDelete(leaseName); ok {
+					old.(*time.Timer).Stop()
+				}
+				t := time.AfterFunc(listenQueueCleanupDelay, func() {
+					if q, ok := s.listenQueues.LoadAndDelete(leaseName); ok {
+						logger.Info("listen queue cleanup timer fired",
+							"lease", leaseName,
+							"bufferedTokens", len(q.(chan *pb.ListenResponse)))
+					}
+					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,151 @@ 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) {
+	original := listenQueueCleanupDelay
+	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)
+
+	// Use a long delay for the first two subtests so the timer cannot fire
+	// between sequential subtests under CI load (fixes flake when >50ms
+	// elapses between subtest boundaries).
+	listenQueueCleanupDelay = 10 * time.Second
+
+	// 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("buffered token survives disconnect and is readable after reconnect", func(t *testing.T) {
+		// Write a Dial token into the queue while the exporter is disconnected.
+		token := &pb.ListenResponse{}
+		ch <- token
+
+		// 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))
+		inherited := got.(chan *pb.ListenResponse)
+		if inherited != ch {
+			t.Fatal("reconnecting Listen() did not inherit the existing queue")
+		}
+
+		// Read the token back -- this is the primary scenario this fix targets.
+		select {
+		case msg := <-inherited:
+			if msg != token {
+				t.Fatal("token read from inherited queue does not match the one written during disconnect")
+			}
+		default:
+			t.Fatal("no token available in inherited queue -- Dial token was lost")
+		}
+	})
+
+	t.Run("reconnecting exporter cancels cleanup timer", func(t *testing.T) {
+		// Re-arm a timer to simulate another transient error.
+		timer := time.AfterFunc(listenQueueCleanupDelay, func() {
+			svc.listenQueues.Delete(leaseName)
+			svc.listenTimers.Delete(leaseName)
+		})
+		svc.listenTimers.Store(leaseName, timer)
+
+		// 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")
+		}
+
+		// Verify the queue is still present -- the stopped timer must not
+		// have fired.
+		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) {
+		// Shorten the delay so this subtest completes quickly.
+		listenQueueCleanupDelay = 50 * time.Millisecond
+
+		// Use a channel to detect when the timer callback completes,
+		// avoiding time.Sleep-based flakiness under CI load.
+		done := make(chan struct{})
+		timer := time.AfterFunc(listenQueueCleanupDelay, func() {
+			svc.listenQueues.Delete(leaseName)
+			svc.listenTimers.Delete(leaseName)
+			close(done)
+		})
+		svc.listenTimers.Store(leaseName, timer)
+
+		// Wait for the timer callback to signal completion.
+		select {
+		case <-done:
+		case <-time.After(5 * time.Second):
+			t.Fatal("cleanup timer did not fire within timeout")
+		}
+		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 ||