Repository: buchgr/bazel-remote
Pull request #592 (2022-10-09 22:09:44)

refactor to allow concurrent graceful shutdown of http and grpc servers


Patch:
diff --git a/BUILD.bazel b/BUILD.bazel
index 4205ce4..eba7b65 100644
--- a/BUILD.bazel
+++ b/BUILD.bazel
@@ -24,7 +24,6 @@ go_library(
     importpath = "github.com/buchgr/bazel-remote",
     visibility = ["//visibility:private"],
     deps = [
-        "//cache:go_default_library",
         "//cache/disk:go_default_library",
         "//config:go_default_library",
         "//server:go_default_library",
@@ -40,6 +39,7 @@ go_library(
         "@com_github_urfave_cli_v2//:go_default_library",
         "@org_golang_google_grpc//:go_default_library",
         "@org_golang_google_grpc//credentials:go_default_library",
+        "@org_golang_x_sync//errgroup:go_default_library",
         "@org_golang_x_sync//semaphore:go_default_library",
     ],
 )
diff --git a/main.go b/main.go
index 1078675..4e48a27 100644
--- a/main.go
+++ b/main.go
@@ -15,7 +15,6 @@ import (
 
 	auth "github.com/abbot/go-http-auth"
 
-	"github.com/buchgr/bazel-remote/cache"
 	"github.com/buchgr/bazel-remote/cache/disk"
 
 	"github.com/buchgr/bazel-remote/config"
@@ -34,6 +33,7 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
 
+	"golang.org/x/sync/errgroup"
 	"golang.org/x/sync/semaphore"
 )
 
@@ -95,12 +95,19 @@ func run(ctx *cli.Context) error {
 	httpSem := semaphore.NewWeighted(1)
 	var httpServer *http.Server
 
+	idleTimeoutChan := make(chan struct{}, 1)
+
 	sigChan := make(chan os.Signal, 1)
 	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
 	go func() {
-		for sig := range sigChan {
+		select {
+		case sig := <-sigChan:
 			log.Printf("Received signal: %s, attempting graceful shutdown", sig)
+		case <-idleTimeoutChan:
+			log.Println("Idle timeout reached, attempting graceful shutdown")
+		}
 
+		go func() {
 			if !grpcSem.TryAcquire(1) {
 				if grpcServer != nil {
 					log.Println("Stopping gRPC server")
@@ -108,26 +115,23 @@ func run(ctx *cli.Context) error {
 					log.Println("gRPC server stopped")
 				}
 			}
+		}()
 
-			// The main goroutine exits when httpServer is stopped, so we
-			// need to wait until the gRPC server is stopped first.
-
+		go func() {
 			if !httpSem.TryAcquire(1) {
 				if httpServer != nil {
 					log.Println("Stopping HTTP server")
 					err := httpServer.Shutdown(context.Background())
 					if err != nil {
 						log.Println("Error occurred while stopping HTTP server:", err)
+					} else {
+						log.Println("HTTP server stopped")
 					}
-
-					// The "HTTP server stopped" log comes from the main goroutine.
 				}
 			}
-		}
+		}()
 	}()
 
-	validateAC := !c.DisableHTTPACValidation
-
 	log.Println("Storage mode:", c.StorageMode)
 	if c.StorageMode == "zstd" {
 		log.Println("Zstandard implementation:", c.ZstdImplementation)
@@ -153,29 +157,14 @@ func run(ctx *cli.Context) error {
 	}
 	diskCache.RegisterMetrics()
 
-	mux := http.NewServeMux()
-	httpServer = &http.Server{
-		Handler:      mux,
-		ReadTimeout:  c.HTTPReadTimeout,
-		TLSConfig:    c.TLSConfig,
-		WriteTimeout: c.HTTPWriteTimeout,
-	}
-
-	checkClientCertForWrites := c.TLSCaFile != "" && c.AllowUnauthenticatedReads
-	h := server.NewHTTPCache(diskCache, c.AccessLogger, c.ErrorLogger, validateAC,
-		c.EnableACKeyInstanceMangling, checkClientCertForWrites, gitCommit)
+	servers := new(errgroup.Group)
 
 	var htpasswdSecrets auth.SecretProvider
+
 	authMode := "disabled"
-	cacheHandler := h.CacheHandler
 	if c.HtpasswdFile != "" {
 		authMode = "basic"
 		htpasswdSecrets = auth.HtpasswdFileProvider(c.HtpasswdFile)
-		if c.AllowUnauthenticatedReads {
-			cacheHandler = unauthenticatedReadWrapper(cacheHandler, htpasswdSecrets, c.HTTPAddress)
-		} else {
-			cacheHandler = authWrapper(cacheHandler, htpasswdSecrets, c.HTTPAddress)
-		}
 	} else if c.TLSCaFile != "" {
 		authMode = "mTLS"
 	}
@@ -191,8 +180,7 @@ func run(ctx *cli.Context) error {
 
 	var idleTimer *idle.Timer
 	if c.IdleTimeout > 0 {
-		idleTimer = idle.NewTimer(c.IdleTimeout)
-		cacheHandler = wrapIdleHandler(cacheHandler, idleTimer, c.AccessLogger, httpServer)
+		idleTimer = idle.NewTimer(c.IdleTimeout, idleTimeoutChan)
 	}
 
 	acKeyManglingStatus := "disabled"
@@ -201,116 +189,95 @@ func run(ctx *cli.Context) error {
 	}
 	log.Println("Mangling non-empty instance names with AC keys:", acKeyManglingStatus)
 
-	if c.EnableEndpointMetrics {
-		metricsMdlw := middleware.New(middleware.Config{
-			Recorder: httpmetrics.NewRecorder(httpmetrics.Config{
-				DurationBuckets: c.MetricsDurationBuckets,
-			}),
-		})
-		mux.Handle("/metrics", middlewarestd.Handler("metrics", metricsMdlw, promhttp.Handler()))
-		mux.Handle("/status", middlewarestd.Handler("status", metricsMdlw, http.HandlerFunc(h.StatusPageHandler)))
-		mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
-			middlewarestd.Handler(r.Method, metricsMdlw, http.HandlerFunc(cacheHandler)).ServeHTTP(w, r)
-		})
-	} else {
-		mux.HandleFunc("/status", h.StatusPageHandler)
-		mux.HandleFunc("/", cacheHandler)
-	}
+	servers.Go(func() error {
+		err := startHttpServer(c, &httpServer, htpasswdSecrets, idleTimer, httpSem, diskCache)
+		if err != nil {
+			log.Fatal("HTTP server returned fatal error:", err)
+		}
+		return nil
+	})
 
 	if c.GRPCAddress != "none" {
-		go func() {
-			opts := []grpc.ServerOption{}
-			streamInterceptors := []grpc.StreamServerInterceptor{}
-			unaryInterceptors := []grpc.UnaryServerInterceptor{}
-
-			if c.EnableEndpointMetrics {
-				streamInterceptors = append(streamInterceptors, grpc_prometheus.StreamServerInterceptor)
-				unaryInterceptors = append(unaryInterceptors, grpc_prometheus.UnaryServerInterceptor)
-				grpc_prometheus.EnableHandlingTimeHistogram(grpc_prometheus.WithHistogramBuckets(c.MetricsDurationBuckets))
-			}
-
-			if c.TLSConfig != nil {
-				opts = append(opts, grpc.Creds(credentials.NewTLS(c.TLSConfig)))
-
-				if c.TLSCaFile != "" {
-					streamInterceptors = append(streamInterceptors,
-						server.GRPCmTLSStreamServerInterceptor(c.AllowUnauthenticatedReads))
-					unaryInterceptors = append(unaryInterceptors,
-						server.GRPCmTLSUnaryServerInterceptor(c.AllowUnauthenticatedReads))
-				}
-			}
-
-			if htpasswdSecrets != nil {
-				gba := server.NewGrpcBasicAuth(htpasswdSecrets, c.AllowUnauthenticatedReads)
-				streamInterceptors = append(streamInterceptors, gba.StreamServerInterceptor)
-				unaryInterceptors = append(unaryInterceptors, gba.UnaryServerInterceptor)
-			}
-
-			if idleTimer != nil {
-				it := server.NewGrpcIdleTimer(idleTimer)
-				streamInterceptors = append(streamInterceptors, it.StreamServerInterceptor)
-				unaryInterceptors = append(unaryInterceptors, it.UnaryServerInterceptor)
+		servers.Go(func() error {
+			err := startGrpcServer(c, &grpcServer, htpasswdSecrets, idleTimer, grpcSem, diskCache)
+			if err != nil {
+				log.Fatal("gRPC server returned fatal error:", err)
 			}
+			return nil
+		})
+	}
 
-			opts = append(opts, grpc.ChainStreamInterceptor(streamInterceptors...))
-			opts = append(opts, grpc.ChainUnaryInterceptor(unaryInterceptors...))
+	if c.ProfileAddress != "" {
+		go func() {
+			// Allow access to /debug/pprof/ URLs.
+			log.Printf("Starting HTTP server for profiling on address %s",
+				c.ProfileAddress)
+			log.Fatal(`Failed to listen on address: "`, c.ProfileAddress,
+				`": `, http.ListenAndServe(c.ProfileAddress, nil))
+		}()
+	}
 
-			validateAC := !c.DisableGRPCACDepsCheck
-			validateStatus := "disabled"
-			if validateAC {
-				validateStatus = "enabled"
-			}
-			log.Println("gRPC AC dependency checks:", validateStatus)
+	if idleTimer != nil {
+		log.Printf("Starting idle timer with value %v", c.IdleTimeout)
+		idleTimer.Start()
+	}
 
-			enableRemoteAssetAPI := c.ExperimentalRemoteAssetAPI
-			remoteAssetStatus := "disabled"
-			if enableRemoteAssetAPI {
-				remoteAssetStatus = "enabled"
-			}
-			log.Println("experimental gRPC remote asset API:", remoteAssetStatus)
+	return servers.Wait()
+}
 
-			network := "tcp"
-			addr := c.GRPCAddress
-			if strings.HasPrefix(c.GRPCAddress, "unix://") {
-				network = "unix"
-				addr = c.GRPCAddress[len("unix://"):]
-			}
+func startHttpServer(c *config.Config, httpServer **http.Server,
+	htpasswdSecrets auth.SecretProvider, idleTimer *idle.Timer,
+	httpSem *semaphore.Weighted, diskCache disk.Cache) error {
 
-			log.Printf("Starting gRPC server on address %s", addr)
+	mux := http.NewServeMux()
+	*httpServer = &http.Server{
+		Handler:      mux,
+		ReadTimeout:  c.HTTPReadTimeout,
+		TLSConfig:    c.TLSConfig,
+		WriteTimeout: c.HTTPWriteTimeout,
+	}
 
-			grpcServer = grpc.NewServer(opts...)
+	checkClientCertForWrites := c.TLSCaFile != "" && c.AllowUnauthenticatedReads
+	validateAC := !c.DisableHTTPACValidation
+	h := server.NewHTTPCache(diskCache, c.AccessLogger, c.ErrorLogger, validateAC,
+		c.EnableACKeyInstanceMangling, checkClientCertForWrites, gitCommit)
 
-			if !grpcSem.TryAcquire(1) {
-				log.Println("bazel-remote is shutting down, not starting gRPC server")
-				return
-			}
-			err3 := server.ListenAndServeGRPC(grpcServer,
-				network, addr,
-				validateAC,
-				c.EnableACKeyInstanceMangling,
-				enableRemoteAssetAPI,
-				diskCache, c.AccessLogger, c.ErrorLogger)
-			if err3 != nil {
-				log.Fatal("gRPC server returned fatal error:", err3)
-			}
-		}()
+	cacheHandler := h.CacheHandler
+	if c.HtpasswdFile != "" {
+		if c.AllowUnauthenticatedReads {
+			cacheHandler = unauthenticatedReadWrapper(cacheHandler, htpasswdSecrets, c.HTTPAddress)
+		} else {
+			cacheHandler = authWrapper(cacheHandler, htpasswdSecrets, c.HTTPAddress)
+		}
 	}
 
-	if c.ProfileAddress != "" {
-		go func() {
-			// Allow access to /debug/pprof/ URLs.
-			log.Printf("Starting HTTP server for profiling on address %s",
-				c.ProfileAddress)
-			log.Fatal(`Failed to listen on address: "`, c.ProfileAddress, `": `, http.ListenAndServe(c.ProfileAddress, nil))
-		}()
+	if c.IdleTimeout > 0 {
+		cacheHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			idleTimer.ResetTimer()
+			cacheHandler(w, r)
+		})
 	}
 
-	validateStatus := "disabled"
-	if validateAC {
-		validateStatus = "enabled"
+	if c.EnableEndpointMetrics {
+		metricsMdlw := middleware.New(middleware.Config{
+			Recorder: httpmetrics.NewRecorder(httpmetrics.Config{
+				DurationBuckets: c.MetricsDurationBuckets,
+			}),
+		})
+		mux.Handle("/metrics", middlewarestd.Handler("metrics", metricsMdlw,
+			promhttp.Handler()))
+		mux.Handle("/status", middlewarestd.Handler("status", metricsMdlw,
+			http.HandlerFunc(h.StatusPageHandler)))
+		mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+			middlewarestd.Handler(r.Method, metricsMdlw, http.HandlerFunc(cacheHandler)).ServeHTTP(w, r)
+		})
+	} else {
+		mux.HandleFunc("/status", h.StatusPageHandler)
+		mux.HandleFunc("/", cacheHandler)
 	}
 
 	var ln net.Listener
+	var err error
 	if strings.HasPrefix(c.HTTPAddress, "unix://") {
 		ln, err = net.Listen("unix", c.HTTPAddress[len("unix://"):])
 	} else {
@@ -320,21 +287,20 @@ func run(ctx *cli.Context) error {
 		log.Fatal(`Failed to listen on address: "`, c.HTTPAddress, `": `, err)
 	}
 
-	if idleTimer != nil {
-		log.Printf("Starting idle timer with value %v", c.IdleTimeout)
-		idleTimer.Start()
+	validateStatus := "disabled"
+	if validateAC {
+		validateStatus = "enabled"
 	}
+	log.Println("HTTP AC validation:", validateStatus)
 
 	if len(c.TLSCertFile) > 0 && len(c.TLSKeyFile) > 0 {
-		log.Printf("Starting HTTPS server on address %s", httpServer.Addr)
-		log.Println("HTTP AC validation:", validateStatus)
-
 		if !httpSem.TryAcquire(1) {
 			log.Println("bazel-remote is shutting down, not starting HTTPS server")
 			return nil
 		}
 
-		err = httpServer.ServeTLS(ln, c.TLSCertFile, c.TLSKeyFile)
+		log.Printf("Starting HTTPS server on address %s", (*httpServer).Addr)
+		err = (*httpServer).ServeTLS(ln, c.TLSCertFile, c.TLSKeyFile)
 		if err == http.ErrServerClosed {
 			log.Println("HTTPS server stopped")
 			return nil
@@ -343,41 +309,96 @@ func run(ctx *cli.Context) error {
 		return err
 	}
 
-	log.Printf("Starting HTTP server on address %s", c.HTTPAddress)
-	log.Println("HTTP AC validation:", validateStatus)
-
 	if !httpSem.TryAcquire(1) {
 		log.Println("bazel-remote is shutting down, not starting HTTP server")
 		return nil
 	}
 
-	err = httpServer.Serve(ln)
+	log.Printf("Starting HTTP server on address %s", c.HTTPAddress)
+	err = (*httpServer).Serve(ln)
 	if err == http.ErrServerClosed {
-		log.Println("HTTP server stopped")
 		return nil
 	}
 
 	return err
 }
 
-func wrapIdleHandler(handler http.HandlerFunc, idleTimer *idle.Timer, accessLogger cache.Logger, httpServer *http.Server) http.HandlerFunc {
+func startGrpcServer(c *config.Config, grpcServer **grpc.Server,
+	htpasswdSecrets auth.SecretProvider, idleTimer *idle.Timer,
+	grpcSem *semaphore.Weighted, diskCache disk.Cache) error {
 
-	tearDown := make(chan struct{})
-	idleTimer.Register(tearDown)
+	opts := []grpc.ServerOption{}
+	streamInterceptors := []grpc.StreamServerInterceptor{}
+	unaryInterceptors := []grpc.UnaryServerInterceptor{}
 
-	go func() {
-		<-tearDown
-		accessLogger.Printf("Shutting down after idle timeout")
-		err := httpServer.Shutdown(context.Background())
-		if err != nil {
-			fmt.Printf("Error when shutting down http server: %s", err.Error())
+	if c.EnableEndpointMetrics {
+		streamInterceptors = append(streamInterceptors, grpc_prometheus.StreamServerInterceptor)
+		unaryInterceptors = append(unaryInterceptors, grpc_prometheus.UnaryServerInterceptor)
+		grpc_prometheus.EnableHandlingTimeHistogram(grpc_prometheus.WithHistogramBuckets(c.MetricsDurationBuckets))
+	}
+
+	if c.TLSConfig != nil {
+		opts = append(opts, grpc.Creds(credentials.NewTLS(c.TLSConfig)))
+
+		if c.TLSCaFile != "" {
+			streamInterceptors = append(streamInterceptors,
+				server.GRPCmTLSStreamServerInterceptor(c.AllowUnauthenticatedReads))
+			unaryInterceptors = append(unaryInterceptors,
+				server.GRPCmTLSUnaryServerInterceptor(c.AllowUnauthenticatedReads))
 		}
-	}()
+	}
 
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		idleTimer.ResetTimer()
-		handler(w, r)
-	})
+	if htpasswdSecrets != nil {
+		gba := server.NewGrpcBasicAuth(htpasswdSecrets, c.AllowUnauthenticatedReads)
+		streamInterceptors = append(streamInterceptors, gba.StreamServerInterceptor)
+		unaryInterceptors = append(unaryInterceptors, gba.UnaryServerInterceptor)
+	}
+
+	if idleTimer != nil {
+		it := server.NewGrpcIdleTimer(idleTimer)
+		streamInterceptors = append(streamInterceptors, it.StreamServerInterceptor)
+		unaryInterceptors = append(unaryInterceptors, it.UnaryServerInterceptor)
+	}
+
+	opts = append(opts, grpc.ChainStreamInterceptor(streamInterceptors...))
+	opts = append(opts, grpc.ChainUnaryInterceptor(unaryInterceptors...))
+
+	validateAC := !c.DisableGRPCACDepsCheck
+	validateStatus := "disabled"
+	if validateAC {
+		validateStatus = "enabled"
+	}
+	log.Println("gRPC AC dependency checks:", validateStatus)
+
+	enableRemoteAssetAPI := c.ExperimentalRemoteAssetAPI
+	remoteAssetStatus := "disabled"
+	if enableRemoteAssetAPI {
+		remoteAssetStatus = "enabled"
+	}
+	log.Println("experimental gRPC remote asset API:", remoteAssetStatus)
+
+	network := "tcp"
+	addr := c.GRPCAddress
+	if strings.HasPrefix(c.GRPCAddress, "unix://") {
+		network = "unix"
+		addr = c.GRPCAddress[len("unix://"):]
+	}
+
+	*grpcServer = grpc.NewServer(opts...)
+
+	if !grpcSem.TryAcquire(1) {
+		log.Println("bazel-remote is shutting down, not starting gRPC server")
+		return nil
+	}
+
+	log.Println("Starting gRPC server on address", addr)
+
+	return server.ListenAndServeGRPC(*grpcServer,
+		network, addr,
+		validateAC,
+		c.EnableACKeyInstanceMangling,
+		enableRemoteAssetAPI,
+		diskCache, c.AccessLogger, c.ErrorLogger)
 }
 
 // A http.HandlerFunc wrapper which requires successful basic
diff --git a/utils/idle/idle.go b/utils/idle/idle.go
index 2476213..ab31396 100644
--- a/utils/idle/idle.go
+++ b/utils/idle/idle.go
@@ -10,28 +10,20 @@ import (
 type Timer struct {
 	mu          sync.Mutex
 	timeout     time.Duration
-	notify      []chan struct{}
+	notify      chan struct{}
 	lastRequest time.Time
 }
 
 // NewTimer creates a new Timer that will send notifications on
 // any registered channels once the idle timeout has been reached.
-func NewTimer(timeout time.Duration) *Timer {
+func NewTimer(timeout time.Duration, notificationChan chan struct{}) *Timer {
 	return &Timer{
 		timeout:     timeout,
 		lastRequest: time.Now(),
-		notify:      make([]chan struct{}, 0),
+		notify:      notificationChan,
 	}
 }
 
-// Register adds a channel that will be notified once the idle
-// timeout is reached.
-func (t *Timer) Register(c chan struct{}) {
-	t.mu.Lock()
-	t.notify = append(t.notify, c)
-	t.mu.Unlock()
-}
-
 // Start begins the idle timer, and returns immediately.
 func (t *Timer) Start() {
 	go t.start()
@@ -48,9 +40,7 @@ func (t *Timer) start() {
 
 		if elapsed > t.timeout {
 			ticker.Stop()
-			for _, c := range t.notify {
-				c <- struct{}{}
-			}
+			t.notify <- struct{}{}
 			return
 		}
 	}


Test patch:
diff --git a/utils/idle/idle_test.go b/utils/idle/idle_test.go
index f498667..cfbf715 100644
--- a/utils/idle/idle_test.go
+++ b/utils/idle/idle_test.go
@@ -9,8 +9,7 @@ import (
 
 func TestIdleTimer(t *testing.T) {
 	tearDown := make(chan struct{})
-	it := idle.NewTimer(time.Second)
-	it.Register(tearDown)
+	it := idle.NewTimer(time.Second, tearDown)
 	it.Start()
 
 	for i := 0; i < 5; i++ {