Fix race in port-forward (#418)

internal/controller/notifier/notifier.go

@@ -31,16 +31,19 @@ func NewNotifier(logger *zap.SugaredLogger) *Notifier {
 func (watcher *Notifier) Register(ctx context.Context, worker string) (chan *rpc.WatchInstruction, func()) {
 	subCtx, cancel := context.WithCancel(ctx)
 	workerCh := make(chan *rpc.WatchInstruction)
-
-	watcher.logger.Debugf("registering worker %s", worker)
-	watcher.workers.Store(worker, &WorkerSlot{
+	slot := &WorkerSlot{
 		ctx: subCtx,
 		ch:  workerCh,
-	})
+	}
+
+	watcher.logger.Debugf("registering worker %s", worker)
+	watcher.workers.Store(worker, slot)
 
 	return workerCh, func() {
 		watcher.logger.Debugf("deleting worker %s", worker)
-		watcher.workers.Delete(worker)
+		watcher.workers.DeleteIf(worker, func(current *WorkerSlot) bool {
+			return current == slot
+		})
 		cancel()
 	}
 }

internal/controller/notifier/notifier_test.go

@@ -47,3 +47,36 @@ func TestNotifier(t *testing.T) {
 
 	wg.Wait()
 }
+
+func TestNotifierReRegisterKeepsNewestSlot(t *testing.T) {
+	ctx := context.Background()
+	watcher := notifier.NewNotifier(zap.NewNop().Sugar())
+
+	const worker = "worker-a"
+
+	_, staleCancel := watcher.Register(ctx, worker)
+	newestCh, newestCancel := watcher.Register(ctx, worker)
+	defer newestCancel()
+
+	// Simulate stale connection cleanup arriving after the worker has already re-registered.
+	staleCancel()
+
+	notifyCtx, notifyCancel := context.WithTimeout(ctx, 300*time.Millisecond)
+	defer notifyCancel()
+
+	notifyErrCh := make(chan error, 1)
+	go func() {
+		notifyErrCh <- watcher.Notify(notifyCtx, worker, nil)
+	}()
+
+	select {
+	case <-newestCh:
+	case err := <-notifyErrCh:
+		require.NoError(t, err)
+		t.Fatal("notify returned before delivering message to newest registration")
+	case <-time.After(time.Second):
+		t.Fatal("timed out waiting for notify delivery")
+	}
+
+	require.NoError(t, <-notifyErrCh)
+}