Impersonator server should take in a cancellable context instead of a stop channel

internal/concierge/impersonator/impersonator.go

@@ -61,7 +61,7 @@ import (
 
 // FactoryFunc is a function which can create an impersonator server.
 // It returns a function which will start the impersonator server.
-// That start function takes a stopCh which can be used to stop the server.
+// That start function takes a context which can be cancelled to stop the server.
 // Once a server has been stopped, don't start it again using the start function.
 // Instead, call the factory function again to get a new start function.
 type FactoryFunc func(
@@ -69,17 +69,19 @@ type FactoryFunc func(
 	dynamicCertProvider dynamiccert.Private,
 	impersonationProxySignerCA dynamiccert.Public,
 	impersonationProxyTokenCache tokenclient.ExpiringSingletonTokenCacheGet,
-) (func(stopCh <-chan struct{}) error, error)
+) (func(ctx context.Context) error, error)
 
 func New(
 	port int,
 	dynamicCertProvider dynamiccert.Private,
 	impersonationProxySignerCA dynamiccert.Public,
 	impersonationProxyTokenCache tokenclient.ExpiringSingletonTokenCacheGet,
-) (func(stopCh <-chan struct{}) error, error) {
+) (func(ctx context.Context) error, error) {
 	return newInternal(port, dynamicCertProvider, impersonationProxySignerCA, kubeclient.Secure, impersonationProxyTokenCache, nil, nil, nil)
 }
 
+var _ FactoryFunc = New
+
 //nolint:funlen // It is definitely too complicated. New calls newInternal, which makes another function.
 func newInternal(
 	port int,
@@ -90,7 +92,7 @@ func newInternal(
 	baseConfig *rest.Config, // for unit testing, should always be nil in production
 	recOpts func(*genericoptions.RecommendedOptions), // for unit testing, should always be nil in production
 	recConfig func(*genericapiserver.RecommendedConfig), // for unit testing, should always be nil in production
-) (func(stopCh <-chan struct{}) error, error) {
+) (func(ctx context.Context) error, error) {
 	var listener net.Listener
 	var err error
 
@@ -101,7 +103,7 @@ func newInternal(
 		}
 	}
 
-	constructServer := func() (func(stopCh <-chan struct{}) error, error) {
+	constructServer := func() (func(ctx context.Context) error, error) {
 		// Bare minimum server side scheme to allow for status messages to be encoded.
 		scheme := runtime.NewScheme()
 		metav1.AddToGroupVersion(scheme, metav1.Unversioned)
@@ -339,7 +341,7 @@ func newInternal(
 			return nil, fmt.Errorf("invalid mutation of impersonation authorizer detected: %#v", preparedRun.Authorizer)
 		}
 
-		return preparedRun.Run, nil
+		return preparedRun.RunWithContext, nil
 	}
 
 	result, err := constructServer()

internal/concierge/impersonator/impersonator_test.go

@@ -876,16 +876,16 @@ func TestImpersonator(t *testing.T) {
 			require.NotNil(t, runner)
 
 			// Start the impersonator.
-			stopCh := make(chan struct{})
+			runnerCtx, runnerCancel := context.WithCancel(context.Background())
 			errCh := make(chan error)
 			go func() {
-				stopErr := runner(stopCh)
+				stopErr := runner(runnerCtx)
 				errCh <- stopErr
 			}()
 
 			// Stop the impersonator server at the end of the test, even if it fails.
 			t.Cleanup(func() {
-				close(stopCh)
+				runnerCancel()
 				exitErr := <-errCh
 				require.NoError(t, exitErr)
 			})

internal/controller/impersonatorconfig/impersonator_config.go

@@ -81,7 +81,7 @@ type impersonatorConfigController struct {
 	impersonatorFunc                 impersonator.FactoryFunc
 
 	hasControlPlaneNodes              *bool
-	serverStopCh                      chan struct{}
+	serverCancelFunc                  context.CancelFunc
 	errorCh                           chan error
 	tlsServingCertDynamicCertProvider dynamiccert.Private
 	log                               plog.Logger
@@ -461,7 +461,7 @@ func (c *impersonatorConfigController) tlsSecretExists() (bool, *corev1.Secret,
 }
 
 func (c *impersonatorConfigController) ensureImpersonatorIsStarted(syncCtx controllerlib.Context) error {
-	if c.serverStopCh != nil {
+	if c.serverCancelFunc != nil {
 		// The server was already started, but it could have died in the background, so make a non-blocking
 		// check to see if it has sent any errors on the errorCh.
 		select {
@@ -495,7 +495,8 @@ func (c *impersonatorConfigController) ensureImpersonatorIsStarted(syncCtx contr
 		return err
 	}
 
-	c.serverStopCh = make(chan struct{})
+	var serverCtx context.Context
+	serverCtx, c.serverCancelFunc = context.WithCancel(context.Background())
 	// use a buffered channel so that startImpersonatorFunc can send
 	// on it without coordinating with the main controller go routine
 	c.errorCh = make(chan error, 1)
@@ -509,26 +510,26 @@ func (c *impersonatorConfigController) ensureImpersonatorIsStarted(syncCtx contr
 		defer syncCtx.Queue.AddRateLimited(syncCtx.Key)
 
 		// Forward any errors returned by startImpersonatorFunc on the errorCh.
-		c.errorCh <- startImpersonatorFunc(c.serverStopCh)
+		c.errorCh <- startImpersonatorFunc(serverCtx)
 	}()
 
 	return nil
 }
 
 func (c *impersonatorConfigController) ensureImpersonatorIsStopped(shouldCloseErrChan bool) error {
-	if c.serverStopCh == nil {
+	if c.serverCancelFunc == nil {
 		return nil
 	}
 
 	c.log.Info("stopping impersonation proxy", "port", c.impersonationProxyPort)
-	close(c.serverStopCh)
+	c.serverCancelFunc()
 	stopErr := <-c.errorCh
 
 	if shouldCloseErrChan {
 		close(c.errorCh)
 	}
 
-	c.serverStopCh = nil
+	c.serverCancelFunc = nil
 	c.errorCh = nil
 
 	return stopErr

internal/controller/impersonatorconfig/impersonator_config_test.go

@@ -316,7 +316,7 @@ func TestImpersonatorConfigControllerSync(t *testing.T) {
 			dynamicCertProvider dynamiccert.Private,
 			impersonationProxySignerCAProvider dynamiccert.Public,
 			expiringSingletonTokenCacheGet tokenclient.ExpiringSingletonTokenCacheGet,
-		) (func(stopCh <-chan struct{}) error, error) {
+		) (func(ctx context.Context) error, error) {
 			impersonatorFuncWasCalled++
 			r.Equal(8444, port)
 			r.NotNil(dynamicCertProvider)
@@ -376,7 +376,7 @@ func TestImpersonatorConfigControllerSync(t *testing.T) {
 			// This fake server is enough like the real impersonation proxy server for this unit test because it
 			// uses the supplied providers to serve TLS. The goal of this unit test is to make sure that the server
 			// was started/stopped/configured correctly, not to test the actual impersonation behavior.
-			return func(stopCh <-chan struct{}) error {
+			return func(ctx context.Context) error {
 				if impersonatorFuncReturnedFuncError != nil {
 					return impersonatorFuncReturnedFuncError
 				}
@@ -406,7 +406,7 @@ func TestImpersonatorConfigControllerSync(t *testing.T) {
 				if testHTTPServerInterruptCh == nil {
 					// Wait in the foreground for the stopCh to be closed, and kill the server when that happens.
 					// This is similar to the behavior of the real impersonation server.
-					<-stopCh
+					<-ctx.Done()
 				} else {
 					// The test supplied an interrupt channel because it wants to test unexpected termination
 					// of the server, so wait for that channel to close instead of waiting for the one that