[ws-proxy] Implement graceful shutdown

Author: aledbf

components/ws-proxy/cmd/run.go

@@ -13,7 +13,6 @@ import (
 	"time"
 
 	"github.com/bombsimon/logrusr/v2"
-	"github.com/gitpod-io/golang-crypto/ssh"
 	"github.com/spf13/cobra"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
@@ -36,6 +35,7 @@ import (
 	"github.com/gitpod-io/gitpod/ws-proxy/pkg/config"
 	"github.com/gitpod-io/gitpod/ws-proxy/pkg/proxy"
 	"github.com/gitpod-io/gitpod/ws-proxy/pkg/sshproxy"
+	"github.com/gitpod-io/golang-crypto/ssh"
 )
 
 var (
@@ -159,11 +159,15 @@ var runCmd = &cobra.Command{
 			}
 		}
 
-		go proxy.NewWorkspaceProxy(cfg.Ingress, cfg.Proxy, proxy.HostBasedRouter(cfg.Ingress.Header, cfg.Proxy.GitpodInstallation.WorkspaceHostSuffix, cfg.Proxy.GitpodInstallation.WorkspaceHostSuffixRegex), infoprov, signers).MustServe()
-		log.Infof("started proxying on %s", cfg.Ingress.HTTPAddress)
+		ctrlCtx := ctrl.SetupSignalHandler()
+
+		go func() {
+			proxy.NewWorkspaceProxy(cfg.Ingress, cfg.Proxy, proxy.HostBasedRouter(cfg.Ingress.Header, cfg.Proxy.GitpodInstallation.WorkspaceHostSuffix, cfg.Proxy.GitpodInstallation.WorkspaceHostSuffixRegex), infoprov, signers).MustServe(ctrlCtx)
+			log.Infof("started proxying on %s", cfg.Ingress.HTTPAddress)
+		}()
 
 		log.Info("🚪 ws-proxy is up and running")
-		if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil {
+		if err := mgr.Start(ctrlCtx); err != nil {
 			log.WithError(err).Fatal(err, "problem starting ws-proxy")
 		}
 

components/ws-proxy/pkg/proxy/proxy.go

@@ -6,11 +6,14 @@ package proxy
 
 import (
 	"bytes"
+	"context"
 	"crypto/tls"
+	"errors"
 	stdlog "log"
 	"net/http"
 	"os"
 	"path/filepath"
+	"time"
 
 	"github.com/gorilla/mux"
 	"github.com/klauspost/cpuid/v2"
@@ -49,13 +52,20 @@ func redirectToHTTPS(w http.ResponseWriter, r *http.Request) {
 }
 
 // MustServe starts the proxy and ends the process if doing so fails.
-func (p *WorkspaceProxy) MustServe() {
+func (p *WorkspaceProxy) MustServe(ctx context.Context) {
 	handler, err := p.Handler()
 	if err != nil {
 		log.WithError(err).Fatal("cannot initialize proxy - this is likely a configuration issue")
 		return
 	}
-	srv := &http.Server{
+
+	httpServer := &http.Server{
+		Addr:     p.Ingress.HTTPAddress,
+		Handler:  http.HandlerFunc(redirectToHTTPS),
+		ErrorLog: stdlog.New(logrusErrorWriter{}, "", 0),
+	}
+
+	httpsServer := &http.Server{
 		Addr:    p.Ingress.HTTPSAddress,
 		Handler: handler,
 		TLSConfig: &tls.Config{
@@ -77,17 +87,35 @@ func (p *WorkspaceProxy) MustServe() {
 		crt = filepath.Join(tproot, crt)
 		key = filepath.Join(tproot, key)
 	}
+
 	go func() {
-		err := http.ListenAndServe(p.Ingress.HTTPAddress, http.HandlerFunc(redirectToHTTPS))
-		if err != nil {
+		err := httpServer.ListenAndServe()
+		if err != nil && !errors.Is(err, http.ErrServerClosed) {
 			log.WithError(err).Fatal("cannot start http proxy")
 		}
 	}()
 
-	err = srv.ListenAndServeTLS(crt, key)
+	go func() {
+		err = httpsServer.ListenAndServeTLS(crt, key)
+		if err != nil && !errors.Is(err, http.ErrServerClosed) {
+			log.WithError(err).Fatal("cannot start proxy")
+			return
+		}
+	}()
+
+	<-ctx.Done()
+
+	shutDownCtx, cancel := context.WithTimeout(ctx, 5*time.Minute)
+	defer cancel()
+
+	err = httpServer.Shutdown(shutDownCtx)
 	if err != nil {
-		log.WithError(err).Fatal("cannot start proxy")
-		return
+		log.WithError(err).Fatal("cannot stop HTTP server")
+	}
+
+	err = httpsServer.Shutdown(shutDownCtx)
+	if err != nil {
+		log.WithError(err).Fatal("cannot stop HTTPS server")
 	}
 }