fix: delete stale or invalid resume token for mongo client changeStreams (#1092)

Author: KaivalyaMDabhadkar

fault-remediation/go.mod

@@ -12,6 +12,7 @@ require (
 	github.com/prometheus/client_golang v1.23.2
 	github.com/prometheus/client_model v0.6.2
 	github.com/stretchr/testify v1.11.1
+	golang.org/x/sync v0.20.0
 	google.golang.org/protobuf v1.36.11
 	k8s.io/api v0.35.3
 	k8s.io/apimachinery v0.35.3
@@ -89,7 +90,6 @@ require (
 	golang.org/x/crypto v0.48.0 // indirect
 	golang.org/x/net v0.51.0 // indirect
 	golang.org/x/oauth2 v0.35.0 // indirect
-	golang.org/x/sync v0.20.0 // indirect
 	golang.org/x/sys v0.41.0 // indirect
 	golang.org/x/term v0.40.0 // indirect
 	golang.org/x/text v0.34.0 // indirect

fault-remediation/main.go

@@ -25,6 +25,7 @@ import (
 	"syscall"
 	"time"
 
+	"golang.org/x/sync/errgroup"
 	batchv1 "k8s.io/api/batch/v1"
 	corev1 "k8s.io/api/core/v1"
 	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
@@ -105,12 +106,63 @@ func run() error {
 func setupCtrlRuntimeManagement(ctx context.Context) error {
 	slog.Info("Running in controller runtime managed mode")
 
+	mgr, err := createManager()
+	if err != nil {
+		return err
+	}
+
+	params := initializer.InitializationParams{
+		TomlConfigPath:     tomlConfigPath,
+		DryRun:             dryRun,
+		EnableLogCollector: enableLogCollector,
+		Config:             mgr.GetConfig(),
+	}
+
+	g, gCtx := errgroup.WithContext(ctx)
+
+	// Start the manager first so health/metrics endpoints are live immediately.
+	// This prevents Kubernetes liveness probes from killing the pod while MongoDB
+	// initialization (which may be slow due to stale resume tokens or connectivity
+	// issues) is still in progress.
+	g.Go(func() error {
+		slog.Info("Starting controller runtime controller")
+
+		if err := mgr.Start(gCtx); err != nil {
+			slog.Error("Problem running manager", "error", err)
+			return err
+		}
+
+		return nil
+	})
+
+	// Initialize datastore and reconciler concurrently — the manager is already
+	// serving health probes, so the pod won't be killed during this phase.
+	// cleanupReconciler is set once the reconciler is created so cleanup can run
+	// after g.Wait() (i.e., after the manager has fully drained).
+	var cleanupReconciler func()
+
+	g.Go(func() error {
+		cleanup, initErr := initializeAndWatch(gCtx, params, mgr)
+		cleanupReconciler = cleanup
+
+		return initErr
+	})
+
+	err = g.Wait()
+
+	if cleanupReconciler != nil {
+		cleanupReconciler()
+	}
+
+	return err
+}
+
+func createManager() (ctrl.Manager, error) {
 	cfg := ctrl.GetConfigOrDie()
 	cfg.Wrap(func(rt http.RoundTripper) http.RoundTripper {
 		return auditlogger.NewAuditingRoundTripper(rt)
 	})
 
-	//TODO: setup informers for node and job
 	mgr, err := ctrl.NewManager(cfg, ctrl.Options{
 		Scheme: scheme,
 		Metrics: metricsserver.Options{
@@ -126,52 +178,64 @@ func setupCtrlRuntimeManagement(ctx context.Context) error {
 	})
 	if err != nil {
 		slog.Error("Unable to start manager", "error", err)
-		return err
+		return nil, err
 	}
 
 	if err = mgr.AddHealthzCheck("healthz", healthz.Ping); err != nil {
 		slog.Error("Unable to set up health check", "error", err)
-		return err
+		return nil, err
 	}
 
 	if err := mgr.AddReadyzCheck("readyz", healthz.Ping); err != nil {
 		slog.Error("Unable to set up ready check", "error", err)
-		return err
+		return nil, err
 	}
 
-	params := initializer.InitializationParams{
-		TomlConfigPath:     tomlConfigPath,
-		DryRun:             dryRun,
-		EnableLogCollector: enableLogCollector,
-		Config:             mgr.GetConfig(),
-	}
+	return mgr, nil
+}
+
+const reconcilerCloseTimeout = 30 * time.Second
 
+// initializeAndWatch performs MongoDB initialization, registers the reconciler, and
+// blocks until shutdown or unexpected stream death. It returns a cleanup function that
+// the caller must invoke after the manager has fully stopped (after g.Wait) so that
+// datastore resources are not torn down under in-flight reconciles.
+func initializeAndWatch(
+	ctx context.Context, params initializer.InitializationParams, mgr ctrl.Manager,
+) (cleanup func(), err error) {
 	components, err := initializer.InitializeAll(ctx, params, mgr.GetClient())
 	if err != nil {
-		return fmt.Errorf("initialization failed: %w", err)
+		return nil, fmt.Errorf("initialization failed: %w", err)
 	}
 
 	reconciler := components.FaultRemediationReconciler
 
-	defer func() {
-		if err := reconciler.CloseAll(ctx); err != nil {
-			slog.Error("failed to close datastore components", "error", err)
+	cleanup = func() {
+		closeCtx, cancel := context.WithTimeout(context.Background(), reconcilerCloseTimeout)
+		defer cancel()
+
+		if closeErr := reconciler.CloseAll(closeCtx); closeErr != nil {
+			slog.Error("failed to close datastore components", "error", closeErr)
 		}
-	}()
+	}
 
-	err = components.FaultRemediationReconciler.SetupWithManager(ctx, mgr)
-	if err != nil {
-		return fmt.Errorf("SetupWithManager failed: %w", err)
+	watcherDone, setupErr := reconciler.SetupWithManager(ctx, mgr)
+	if setupErr != nil {
+		return cleanup, fmt.Errorf("SetupWithManager failed: %w", setupErr)
 	}
 
-	slog.Info("Starting controller runtime controller")
+	slog.Info("Initialization completed, reconciler registered with manager")
 
-	if err = mgr.Start(ctx); err != nil {
-		slog.Error("Problem running manager", "error", err)
-		return err
-	}
+	select {
+	case <-ctx.Done():
+		return cleanup, nil
+	case <-watcherDone:
+		if ctx.Err() == nil {
+			return cleanup, fmt.Errorf("change stream watcher terminated unexpectedly, event processing has stopped")
+		}
 
-	return nil
+		return cleanup, nil
+	}
 }
 
 func parseFlags() {

fault-remediation/pkg/reconciler/reconciler.go

@@ -593,14 +593,14 @@ func (r *FaultRemediationReconciler) CloseAll(ctx context.Context) error {
 }
 
 // SetupWithManager configures the reconciler for controller-runtime managed operation.
-// It starts the watcher stream using the provided context and registers the reconciler
-// with the manager using a typed channel source. This method should only be called
-// when running under controller-runtime management.
-func (r *FaultRemediationReconciler) SetupWithManager(ctx context.Context, mgr ctrl.Manager) error {
+// It starts the watcher stream and returns a done channel that is closed when the event
+// adapter goroutine exits. Callers should monitor this channel: if it closes while the
+// context is still active, the change stream died unexpectedly and the pod should exit.
+func (r *FaultRemediationReconciler) SetupWithManager(ctx context.Context, mgr ctrl.Manager) (<-chan struct{}, error) {
 	r.Watcher.Start(ctx)
 
 	reconciler := builder.TypedControllerManagedBy[*datastore.EventWithToken](mgr)
-	typedCh := AdaptEvents(ctx, r.Watcher.Events())
+	typedCh, watcherDone := AdaptEvents(ctx, r.Watcher.Events())
 
 	src := source.TypedChannel[*datastore.EventWithToken, *datastore.EventWithToken](
 		typedCh,
@@ -615,25 +615,31 @@ func (r *FaultRemediationReconciler) SetupWithManager(ctx context.Context, mgr c
 		},
 	)
 
-	return reconciler.
+	err := reconciler.
 		Named("fault-remediation-controller").
 		WatchesRawSource(
 			src,
 		).
 		Complete(r)
+
+	return watcherDone, err
 }
 
 // AdaptEvents transforms a channel of EventWithToken into a channel of controller-runtime
 // TypedGenericEvent. It spawns a goroutine that continuously reads from the input channel
 // until either the context is cancelled or the input channel is closed.
+// The returned done channel is closed when the adapter goroutine exits. If the input channel
+// closed while the context was still active, this indicates the change stream died unexpectedly.
 func AdaptEvents(
 	ctx context.Context,
 	in <-chan datastore.EventWithToken,
-) <-chan event.TypedGenericEvent[*datastore.EventWithToken] {
+) (<-chan event.TypedGenericEvent[*datastore.EventWithToken], <-chan struct{}) {
 	out := make(chan event.TypedGenericEvent[*datastore.EventWithToken])
+	done := make(chan struct{})
 
 	go func() {
 		defer close(out)
+		defer close(done)
 
 		for {
 			select {
@@ -650,5 +656,5 @@ func AdaptEvents(
 		}
 	}()
 
-	return out
+	return out, done
 }

fault-remediation/pkg/reconciler/reconciler_e2e_test.go

@@ -296,7 +296,7 @@ func TestMain(m *testing.M) {
 
 	reconciler = NewFaultRemediationReconciler(nil, mockWatcher, mockStore, cfg, false)
 
-	err = reconciler.SetupWithManager(testContext, mgr)
+	_, err = reconciler.SetupWithManager(testContext, mgr)
 	if err != nil {
 		log.Fatalf("Failed to launch reconciler with mgr %v", err)
 	}

fault-remediation/pkg/reconciler/reconciler_test.go

@@ -1029,3 +1029,55 @@ func TestLogCollectorOnlyCalledWhenShouldCreateCR(t *testing.T) {
 		})
 	}
 }
+
+func TestAdaptEvents_DoneChannelClosesOnInputClose(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	in := make(chan datastore.EventWithToken)
+	_, done := AdaptEvents(ctx, in)
+
+	// Close the input channel to simulate change stream death
+	close(in)
+
+	select {
+	case <-done:
+		// done channel closed as expected
+	case <-time.After(2 * time.Second):
+		t.Fatal("done channel was not closed after input channel closed")
+	}
+}
+
+func TestAdaptEvents_DoneChannelClosesOnContextCancel(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	in := make(chan datastore.EventWithToken)
+	_, done := AdaptEvents(ctx, in)
+
+	cancel()
+
+	select {
+	case <-done:
+		// done channel closed as expected
+	case <-time.After(2 * time.Second):
+		t.Fatal("done channel was not closed after context cancellation")
+	}
+}
+
+func TestAdaptEvents_ForwardsEvents(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	in := make(chan datastore.EventWithToken, 1)
+	out, _ := AdaptEvents(ctx, in)
+
+	testEvent := datastore.EventWithToken{}
+	in <- testEvent
+
+	select {
+	case <-out:
+		// Event forwarded successfully
+	case <-time.After(2 * time.Second):
+		t.Fatal("event was not forwarded through AdaptEvents")
+	}
+}

node-drainer/main.go

@@ -112,48 +112,50 @@ func run() error {
 		DryRun:                      *dryRun,
 	}
 
-	components, err := initializer.InitializeAll(ctx, params)
+	// Create and start the health/metrics server BEFORE the potentially slow MongoDB
+	// initialization. This ensures Kubernetes liveness probes get HTTP 200 responses
+	// immediately, preventing the pod from being killed during initialization.
+	srv, err := createMetricsServer(*metricsPort)
+	if err != nil {
+		return err
+	}
+
+	g, gCtx := errgroup.WithContext(ctx)
+
+	startMetricsServer(g, gCtx, srv)
+
+	// Initialize components (may block on MongoDB connectivity / change stream setup)
+	components, err := initializer.InitializeAll(gCtx, params)
 	if err != nil {
 		return fmt.Errorf("failed to initialize components: %w", err)
 	}
 
 	// Informers must sync before processing events
 	slog.Info("Starting Kubernetes informers")
 
-	if err := components.Informers.Run(ctx); err != nil {
+	if err := components.Informers.Run(gCtx); err != nil {
 		return fmt.Errorf("failed to start informers: %w", err)
 	}
 
 	slog.Info("Kubernetes informers started and synced")
 
 	slog.Info("Starting queue worker")
-	components.QueueManager.Start(ctx)
+	components.QueueManager.Start(gCtx)
 
 	// Handle cold start - re-process any events that were in-progress during restart
 	slog.Info("Handling cold start")
 
-	if err := handleColdStart(ctx, components); err != nil {
+	if err := handleColdStart(gCtx, components); err != nil {
 		slog.Error("Cold start handling failed", "error", err)
 	}
 
 	slog.Info("Starting database event watcher")
 
 	criticalError := make(chan error)
-	startEventWatcher(ctx, components, criticalError)
+	startEventWatcher(gCtx, components, criticalError)
 
 	slog.Info("All components started successfully")
 
-	srv, err := createMetricsServer(*metricsPort)
-	if err != nil {
-		return err
-	}
-
-	// Start server in errgroup alongside event watcher monitoring
-	g, gCtx := errgroup.WithContext(ctx)
-
-	// Start the metrics/health server
-	startMetricsServer(g, gCtx, srv)
-
 	// Monitor for critical errors or graceful shutdown signals.
 	g.Go(func() error {
 		select {
@@ -176,7 +178,7 @@ func run() error {
 		return shutdownComponents(ctx, components)
 	})
 
-	// Wait for both goroutines to finish
+	// Wait for all goroutines to finish
 	return g.Wait()
 }
 
@@ -242,7 +244,16 @@ func startEventWatcher(ctx context.Context, components *initializer.Components,
 			}
 		}
 
-		slog.Info("Event watcher stopped")
+		// The event channel closed. If the context is still active, this means the
+		// change stream died unexpectedly (e.g., MongoDB error). Signal a critical
+		// failure so the pod exits and Kubernetes restarts it.
+		if ctx.Err() == nil {
+			slog.Error("Event watcher channel closed unexpectedly, event processing has stopped")
+
+			criticalError <- fmt.Errorf("event watcher channel closed unexpectedly")
+		} else {
+			slog.Info("Event watcher stopped")
+		}
 	}()
 }