From ad5e4627d029c6229e1af28eaf60def59da057da Mon Sep 17 00:00:00 2001
From: Zeger-Jan van de Weg <git@zjvandeweg.nl>
Date: Thu, 24 Jan 2019 11:59:24 +0100
Subject: [PATCH] Gracefully stop the server on signal received

SIGKILL needs to happen right away, but the signals Gitaly handles can
stop gracefully. This makes sure no new connections are excepted while
the current requests are being finished.

Seen at: https://github.com/fabiolb/fabio/blob/master/proxy/grpc_handler.go#L38
and seemed like a good idea
---
 cmd/gitaly/main.go | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/cmd/gitaly/main.go b/cmd/gitaly/main.go
index 3be1e38e8b1..9decd10749e 100644
--- a/cmd/gitaly/main.go
+++ b/cmd/gitaly/main.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"net"
@@ -8,6 +9,7 @@ import (
 	"os"
 	"os/signal"
 	"syscall"
+	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
@@ -21,6 +23,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/internal/tempdir"
 	"gitlab.com/gitlab-org/gitaly/internal/version"
 	"gitlab.com/gitlab-org/labkit/tracing"
+	"google.golang.org/grpc"
 )
 
 var (
@@ -175,11 +178,15 @@ func run(insecureListeners, secureListeners []net.Listener) error {
 	}
 	defer ruby.Stop()
 
+	var servers []*grpc.Server
+
 	serverErrors := make(chan error, len(insecureListeners)+len(secureListeners))
 	if len(insecureListeners) > 0 {
 		insecureServer := server.NewInsecure(ruby)
 		defer insecureServer.Stop()
 
+		servers = append(servers, insecureServer)
+
 		for _, listener := range insecureListeners {
 			// Must pass the listener as a function argument because there is a race
 			// between 'go' and 'for'.
@@ -193,6 +200,8 @@ func run(insecureListeners, secureListeners []net.Listener) error {
 		secureServer := server.NewSecure(ruby)
 		defer secureServer.Stop()
 
+		servers = append(servers, secureServer)
+
 		for _, listener := range secureListeners {
 			go func(l net.Listener) {
 				serverErrors <- secureServer.Serve(l)
@@ -202,9 +211,33 @@ func run(insecureListeners, secureListeners []net.Listener) error {
 
 	select {
 	case s := <-termCh:
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+
 		err = fmt.Errorf("received signal %q", s)
+		for _, srv := range servers {
+			if err := gracefulStopServer(ctx, srv); err != nil {
+				log.Warnf("error while attempting a graceful stop: %v", err)
+			}
+		}
 	case err = <-serverErrors:
 	}
 
 	return err
 }
+
+func gracefulStopServer(ctx context.Context, srv *grpc.Server) error {
+	done := make(chan struct{})
+	go func() {
+		srv.GracefulStop()
+		close(done)
+	}()
+
+	select {
+	case <-ctx.Done():
+		srv.Stop()
+		return ctx.Err()
+	case <-done:
+		return nil
+	}
+}
-- 
GitLab