cmd/kube-rbac-proxy/app: use SafeWaitGroup from k8s.io/apimachinery

cmd/kube-rbac-proxy/app/kube-rbac-proxy.go

@@ -26,11 +26,11 @@ import (
 	"os"
 	"time"
 
-	"github.com/oklog/run"
 	"github.com/spf13/cobra"
 	"golang.org/x/net/http2"
 
 	utilerrors "k8s.io/apimachinery/pkg/util/errors"
+	waitgroup "k8s.io/apimachinery/pkg/util/waitgroup"
 	"k8s.io/apiserver/pkg/authentication/authenticator"
 	"k8s.io/apiserver/pkg/authorization/authorizer"
 	"k8s.io/apiserver/pkg/authorization/union"
@@ -205,24 +205,27 @@ func Run(cfg *server.KubeRBACProxyConfig) error {
 	mux := http.NewServeMux()
 	mux.Handle("/", handler)
 
-	gr := &run.Group{}
+	var wg waitgroup.SafeWaitGroup
 	{
 		// listener for proxying HTTPS with authentication and authorization (on port --secure-port)
-		gr.Add(secureServerRunner(ctx, cfg.SecureServing, mux))
+		runner, interrupter := secureServerRunner(ctx, cfg.SecureServing, mux)
+		if err := addToSafeWaitGroup(&wg, runner, interrupter); err != nil {
+			return err
+		}
 
 		if cfg.KubeRBACProxyInfo.ProxyEndpointsSecureServing != nil {
 			// we need a second listener in order to serve proxy-specific endpoints
 			// on a different port (--proxy-endpoints-port)
 			proxyEndpointsMux := http.NewServeMux()
 			proxyEndpointsMux.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) { _, _ = w.Write([]byte("ok")) })
 
-			gr.Add(secureServerRunner(ctx, cfg.KubeRBACProxyInfo.ProxyEndpointsSecureServing, proxyEndpointsMux))
+			runner, interrupter := secureServerRunner(ctx, cfg.KubeRBACProxyInfo.ProxyEndpointsSecureServing, proxyEndpointsMux)
+			if err := addToSafeWaitGroup(&wg, runner, interrupter); err != nil {
+				return err
+			}
 		}
 	}
-
-	if err := gr.Run(); err != nil {
-		return fmt.Errorf("failed to run groups: %w", err)
-	}
+	wg.Wait()
 
 	return nil
 }
@@ -261,14 +264,17 @@ func createKubeRBACProxyConfig(opts *completedProxyRunOptions) (*server.KubeRBAC
 	return proxyConfig, nil
 }
 
+type runnerFunc func() error
+type interrupterFunc func(error)
+
 func secureServerRunner(
 	ctx context.Context,
 	config *serverconfig.SecureServingInfo,
 	handler http.Handler,
-) (func() error, func(error)) {
+) (runnerFunc, interrupterFunc) {
 	serverStopCtx, serverCtxCancel := context.WithCancel(ctx)
 
-	runner := func() error {
+	runner := runnerFunc(func() error {
 		stoppedCh, listenerStoppedCh, err := config.Serve(handler, 10*time.Second, serverStopCtx.Done())
 		if err != nil {
 			serverCtxCancel()
@@ -278,11 +284,11 @@ func secureServerRunner(
 		<-listenerStoppedCh
 		<-stoppedCh
 		return err
-	}
+	})
 
-	interrupter := func(err error) {
+	interrupter := interrupterFunc(func(err error) {
 		serverCtxCancel()
-	}
+	})
 
 	return runner, interrupter
 }
@@ -312,3 +318,18 @@ func setupAuthorizer(krpInfo *server.KubeRBACProxyInfo, delegatedAuthz *serverco
 
 	return authz, nil
 }
+
+func addToSafeWaitGroup(wg *waitgroup.SafeWaitGroup, runner runnerFunc, interrupter interrupterFunc) error {
+	if err := wg.Add(1); err != nil {
+		return err
+	}
+
+	go func() {
+		defer wg.Done()
+		if err := runner(); err != nil {
+			interrupter(err)
+		}
+	}()
+
+	return nil
+}