fix: add more graceful shutdown

Signed-off-by: Timur Tuktamyshev <timur.tuktamyshev@flant.com>

Author: timmilesdw

pkg/kube_events_manager/factory.go

@@ -36,38 +36,44 @@ type Factory struct {
 	handlerRegistrations map[string]cache.ResourceEventHandlerRegistration
 	ctx                  context.Context
 	cancel               context.CancelFunc
+	// done is closed when the underlying informer.Run returns
+	done chan struct{}
 }
 
 type FactoryStore struct {
 	mu   sync.Mutex
-	data map[FactoryIndex]Factory
+	cond *sync.Cond
+	data map[FactoryIndex]*Factory
 }
 
 func NewFactoryStore() *FactoryStore {
-	return &FactoryStore{
-		data: make(map[FactoryIndex]Factory),
+	fs := &FactoryStore{
+		data: make(map[FactoryIndex]*Factory),
 	}
+	fs.cond = sync.NewCond(&fs.mu)
+	return fs
 }
 
 func (c *FactoryStore) Reset() {
 	c.mu.Lock()
 	defer c.mu.Unlock()
-	c.data = make(map[FactoryIndex]Factory)
+	c.data = make(map[FactoryIndex]*Factory)
 }
 
 func (c *FactoryStore) add(index FactoryIndex, f dynamicinformer.DynamicSharedInformerFactory) {
 	ctx, cancel := context.WithCancel(context.Background())
-	c.data[index] = Factory{
+	c.data[index] = &Factory{
 		shared:               f,
 		handlerRegistrations: make(map[string]cache.ResourceEventHandlerRegistration),
 		ctx:                  ctx,
 		cancel:               cancel,
+		done:                 nil,
 	}
 	log.Debug("Factory store: added a new factory for index",
 		slog.String("namespace", index.Namespace), slog.String("gvr", index.GVR.String()))
 }
 
-func (c *FactoryStore) get(client dynamic.Interface, index FactoryIndex) Factory {
+func (c *FactoryStore) get(client dynamic.Interface, index FactoryIndex) *Factory {
 	f, ok := c.data[index]
 	if ok {
 		log.Debug("Factory store: the factory with index found",
@@ -115,8 +121,18 @@ func (c *FactoryStore) Start(ctx context.Context, informerId string, client dyna
 		slog.Int("value", len(factory.handlerRegistrations)),
 		slog.String("namespace", index.Namespace), slog.String("gvr", index.GVR.String()))
 
+	// Ensure informer.Run is started once and tracked
+	if factory.done == nil {
+		factory.done = make(chan struct{})
+		go func() {
+			informer.Run(factory.ctx.Done())
+			close(factory.done)
+			log.Debug("Factory store: informer goroutine exited",
+				slog.String("namespace", index.Namespace), slog.String("gvr", index.GVR.String()))
+		}()
+	}
+
 	if !informer.HasSynced() {
-		go informer.Run(factory.ctx.Done())
 
 		if err := wait.PollUntilContextCancel(ctx, DefaultSyncTime, true, func(_ context.Context) (bool, error) {
 			return informer.HasSynced(), nil
@@ -131,11 +147,10 @@ func (c *FactoryStore) Start(ctx context.Context, informerId string, client dyna
 
 func (c *FactoryStore) Stop(informerId string, index FactoryIndex) {
 	c.mu.Lock()
-	defer c.mu.Unlock()
-
 	f, ok := c.data[index]
 	if !ok {
 		// already deleted
+		c.mu.Unlock()
 		return
 	}
 
@@ -152,10 +167,32 @@ func (c *FactoryStore) Stop(informerId string, index FactoryIndex) {
 			slog.Int("value", len(f.handlerRegistrations)),
 			slog.String("namespace", index.Namespace), slog.String("gvr", index.GVR.String()))
 		if len(f.handlerRegistrations) == 0 {
+			log.Debug("Factory store: last handler removed, canceling shared informer",
+				slog.String("namespace", index.Namespace), slog.String("gvr", index.GVR.String()))
+			done := f.done
 			f.cancel()
+			c.mu.Unlock()
+			if done != nil {
+				<-done
+			}
+			c.mu.Lock()
 			delete(c.data, index)
 			log.Debug("Factory store: deleted factory",
 				slog.String("namespace", index.Namespace), slog.String("gvr", index.GVR.String()))
+			c.cond.Broadcast()
+		}
+	}
+	c.mu.Unlock()
+}
+
+// WaitStopped blocks until there is no factory for the index
+func (c *FactoryStore) WaitStopped(index FactoryIndex) {
+	c.mu.Lock()
+	for {
+		if _, ok := c.data[index]; !ok {
+			c.mu.Unlock()
+			return
 		}
+		c.cond.Wait()
 	}
 }

pkg/kube_events_manager/kube_events_manager.go

@@ -25,6 +25,7 @@ type KubeEventsManager interface {
 
 	Ch() chan kemtypes.KubeEvent
 	Stop()
+	Wait()
 }
 
 // kubeEventsManager is a main implementation of KubeEventsManager.
@@ -143,6 +144,18 @@ func (mgr *kubeEventsManager) Stop() {
 	mgr.cancel()
 }
 
+func (mgr *kubeEventsManager) Wait() {
+	mgr.m.RLock()
+	monitors := make([]Monitor, 0, len(mgr.Monitors))
+	for _, mon := range mgr.Monitors {
+		monitors = append(monitors, mon)
+	}
+	mgr.m.RUnlock()
+	for _, mon := range monitors {
+		mon.Wait()
+	}
+}
+
 func (mgr *kubeEventsManager) MetricStorage() metric.Storage {
 	return mgr.metricStorage
 }

pkg/kube_events_manager/monitor.go

@@ -19,6 +19,7 @@ type Monitor interface {
 	CreateInformers() error
 	Start(context.Context)
 	Stop()
+	Wait()
 	Snapshot() []kemtypes.ObjectAndFilterResult
 	EnableKubeEventCb()
 	GetConfig() *MonitorConfig
@@ -372,6 +373,21 @@ func (m *monitor) Stop() {
 	}
 }
 
+// Wait waits for all started informers to stop
+func (m *monitor) Wait() {
+	for _, informer := range m.ResourceInformers {
+		informer.wait()
+	}
+	m.VaryingInformers.RangeValue(func(value []*resourceInformer) {
+		for _, informer := range value {
+			informer.wait()
+		}
+	})
+	if m.NamespaceInformer != nil {
+		m.NamespaceInformer.wait()
+	}
+}
+
 func (m *monitor) SnapshotOperations() (*CachedObjectsInfo /*total*/, *CachedObjectsInfo /*last*/) {
 	total := &CachedObjectsInfo{}
 	last := &CachedObjectsInfo{}

pkg/kube_events_manager/namespace_informer.go

@@ -21,6 +21,7 @@ type namespaceInformer struct {
 	ctx     context.Context
 	cancel  context.CancelFunc
 	stopped bool
+	done    chan struct{}
 
 	KubeClient     *klient.Client
 	Monitor        *MonitorConfig
@@ -128,13 +129,18 @@ func (ni *namespaceInformer) start() {
 		return
 	}
 	cctx, cancel := context.WithCancel(ni.ctx)
+	ni.done = make(chan struct{})
 	go func() {
 		<-ni.ctx.Done()
 		ni.stopped = true
 		cancel()
 	}()
 
-	go ni.SharedInformer.Run(cctx.Done())
+	go func() {
+		ni.SharedInformer.Run(cctx.Done())
+		close(ni.done)
+		log.Debug("Namespace informer goroutine exited", slog.String("name", ni.Monitor.Metadata.DebugName))
+	}()
 
 	if err := wait.PollUntilContextCancel(cctx, DefaultSyncTime, true, func(_ context.Context) (bool, error) {
 		return ni.SharedInformer.HasSynced(), nil
@@ -145,3 +151,9 @@ func (ni *namespaceInformer) start() {
 
 	log.Debug("Informer is ready", slog.String("debugName", ni.Monitor.Metadata.DebugName))
 }
+
+func (ni *namespaceInformer) wait() {
+	if ni.done != nil {
+		<-ni.done
+	}
+}

pkg/kube_events_manager/resource_informer.go

@@ -63,6 +63,9 @@ type resourceInformer struct {
 	stopped bool
 
 	logger *log.Logger
+
+	// wait until underlying client-go informer stops
+	// implemented via FactoryStore.WaitStopped on FactoryIndex
 }
 
 // resourceInformer should implement ResourceInformer
@@ -469,6 +472,11 @@ func (ei *resourceInformer) start() {
 	log.Debug("informer is ready", slog.String("debugName", ei.Monitor.Metadata.DebugName))
 }
 
+// wait blocks until the underlying shared informer for this FactoryIndex is stopped
+func (ei *resourceInformer) wait() {
+	DefaultFactoryStore.WaitStopped(ei.FactoryIndex)
+}
+
 // CachedObjectsInfo returns info accumulated from start.
 func (ei *resourceInformer) getCachedObjectsInfo() CachedObjectsInfo {
 	ei.cacheLock.RLock()

pkg/shell-operator/operator.go

@@ -960,9 +960,20 @@ func (op *ShellOperator) runMetrics() {
 
 // Shutdown pause kubernetes events handling and stop queues. Wait for queues to stop.
 func (op *ShellOperator) Shutdown() {
+	log.Info("shutdown: begin")
 	op.ScheduleManager.Stop()
+	log.Info("shutdown: schedule manager stopped")
+
 	op.KubeEventsManager.Stop()
+	log.Info("shutdown: kube events manager canceled, waiting for informers")
+	if kem, ok := op.KubeEventsManager.(interface{ Wait() }); ok {
+		kem.Wait()
+	}
+	log.Info("shutdown: kube events manager done")
+
 	op.TaskQueues.Stop()
+	log.Info("shutdown: task queues stop signaled, waiting")
 	// Wait for queues to stop, but no more than 10 seconds
 	op.TaskQueues.WaitStopWithTimeout(WaitQueuesTimeout)
+	log.Info("shutdown: task queues stopped")
 }