fix context handling so informers are active while webhook is draining

Author: dprotaso

vendor/knative.dev/pkg/injection/sharedmain/main.go

@@ -240,7 +240,11 @@ func IsHADisabled(ctx context.Context) bool {
 
 // MainWithConfig runs the generic main flow for controllers and webhooks
 // with the given config.
-func MainWithConfig(ctx context.Context, component string, cfg *rest.Config, ctors ...injection.ControllerConstructor) {
+func MainWithConfig(parentCtx context.Context, component string, cfg *rest.Config, ctors ...injection.ControllerConstructor) {
+	ctx := context.WithoutCancel(parentCtx)
+	ctx, ctxCancel := context.WithCancel(ctx)
+	defer ctxCancel()
+
 	log.Printf("Registering %d clients", len(injection.Default.GetClients()))
 	log.Printf("Registering %d informer factories", len(injection.Default.GetInformerFactories()))
 	log.Printf("Registering %d informers", len(injection.Default.GetInformers()))
@@ -315,13 +319,23 @@ func MainWithConfig(ctx context.Context, component string, cfg *rest.Config, cto
 	// and pass them in.
 	var wh *webhook.Webhook
 	if len(webhooks) > 0 {
+		// we use ctx that has the informers etc.
 		wh, err = webhook.New(ctx, webhooks)
 		if err != nil {
 			logger.Fatalw("Failed to create webhook", zap.Error(err))
 		}
 		eg.Go(func() error {
-			return wh.Run(ctx.Done())
+			// we use the parent context because we want the webhook
+			// to stop first prior to stopping the informers
+			defer ctxCancel()
+			return wh.Run(parentCtx.Done())
 		})
+	} else {
+		// no webhooks cancel the context when the parent is cancelled
+		go func() {
+			<-parentCtx.Done()
+			ctxCancel()
+		}()
 	}
 
 	// Start the injection clients and informers.
@@ -522,7 +536,8 @@ func WatchObservabilityConfigOrDie(
 
 // ControllersAndWebhooksFromCtors returns a list of the controllers and a list
 // of the webhooks created from the given constructors.
-func ControllersAndWebhooksFromCtors(ctx context.Context,
+func ControllersAndWebhooksFromCtors(
+	ctx context.Context,
 	cmw *cminformer.InformedWatcher,
 	ctors ...injection.ControllerConstructor,
 ) ([]*controller.Impl, []any) {