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"
@@ -44,6 +44,7 @@ import (
 	"k8s.io/component-base/logs"
 	"k8s.io/component-base/term"
 	"k8s.io/component-base/version/verflag"
+	"k8s.io/klog/v2"
 
 	"github.com/brancz/kube-rbac-proxy/cmd/kube-rbac-proxy/app/options"
 	"github.com/brancz/kube-rbac-proxy/pkg/authn"
@@ -236,55 +237,59 @@ func Run(cfg *server.KubeRBACProxyConfig) error {
 	handler = kubefilters.WithRequestInfo(handler, &request.RequestInfoFactory{})
 	handler = rewrite.WithKubeRBACProxyParamsHandler(handler, cfg.KubeRBACProxyInfo.Authorization.RewriteAttributesConfig)
 
+	var wg waitgroup.SafeWaitGroup
+	serverCtx, cancel := context.WithCancel(ctx)
+
+	// listener for proxying HTTPS with authentication and authorization (on port --secure-port)
 	mux := http.NewServeMux()
 	mux.Handle("/", handler)
 
-	gr := &run.Group{}
-	{
-		// listener for proxying HTTPS with authentication and authorization (on port --secure-port)
-		gr.Add(secureServerRunner(ctx, cfg.SecureServing, mux))
-
-		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))
-		}
+	if err := wg.Add(1); err != nil {
+		cancel()
+		return err
 	}
 
-	if err := gr.Run(); err != nil {
-		return fmt.Errorf("failed to run groups: %w", err)
-	}
+	go func() {
+		defer wg.Done()
+		stoppedCh, listenerStoppedCh, err := cfg.SecureServing.Serve(mux, 10*time.Second, serverCtx.Done())
+		if err != nil {
+			klog.Errorf("%v", err)
+			cancel()
+			return
+		}
 
-	return nil
-}
+		<-listenerStoppedCh
+		<-stoppedCh
+	}()
 
-func secureServerRunner(
-	ctx context.Context,
-	config *serverconfig.SecureServingInfo,
-	handler http.Handler,
-) (func() error, func(error)) {
-	serverStopCtx, serverCtxCancel := context.WithCancel(ctx)
+	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")) })
 
-	runner := func() error {
-		stoppedCh, listenerStoppedCh, err := config.Serve(handler, 10*time.Second, serverStopCtx.Done())
-		if err != nil {
-			serverCtxCancel()
+		if err := wg.Add(1); err != nil {
+			cancel()
 			return err
 		}
 
-		<-listenerStoppedCh
-		<-stoppedCh
-		return err
-	}
+		go func() {
+			defer wg.Done()
+			stoppedCh, listenerStoppedCh, err := cfg.KubeRBACProxyInfo.ProxyEndpointsSecureServing.Serve(proxyEndpointsMux, 10*time.Second, serverCtx.Done())
+			if err != nil {
+				klog.Errorf("%v", err)
+				cancel()
+				return
+			}
 
-	interrupter := func(err error) {
-		serverCtxCancel()
+			<-listenerStoppedCh
+			<-stoppedCh
+		}()
 	}
 
-	return runner, interrupter
+	wg.Wait()
+
+	return nil
 }
 
 func setupAuthorizer(krbInfo *server.KubeRBACProxyInfo, delegatedAuthz *serverconfig.AuthorizationInfo) (authorizer.Authorizer, error) {