From 4e16a778f3bd11c04b404ff20abeacc8043009c8 Mon Sep 17 00:00:00 2001
From: lqs <lqs@lqs.me>
Date: Fri, 5 Apr 2024 15:29:02 +0900
Subject: [PATCH] make graceful shutdown really works

---
 handler.go |  8 ++++++--
 server.go  | 26 ++++++++++++++++----------
 2 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/handler.go b/handler.go
index 4797652..606f923 100644
--- a/handler.go
+++ b/handler.go
@@ -3,6 +3,7 @@ package grpcmix
 import (
 	"net/http"
 	"strings"
+	"sync"
 
 	"github.com/improbable-eng/grpc-web/go/grpcweb"
 	"golang.org/x/net/http2"
@@ -90,16 +91,19 @@ func (h mixHandler) handle() {
 	}
 }
 
-func newHandler(grpcServer *grpc.Server, httpHandler http.Handler) http.Handler {
+func newHandler(grpcServer *grpc.Server, http2Server *http2.Server, httpHandler http.Handler) (http.Handler, func()) {
 	grpcWeb := grpcweb.WrapServer(grpcServer, grpcweb.WithOriginFunc(func(origin string) bool {
 		return true
 	}))
+	var wg sync.WaitGroup
 	return h2c.NewHandler(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
 		defer func() {
 			if r := recover(); r != nil {
 				http.Error(writer, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
 			}
 		}()
+		wg.Add(1)
+		defer wg.Done()
 		mixHandler{
 			writer:      writer,
 			request:     request,
@@ -107,5 +111,5 @@ func newHandler(grpcServer *grpc.Server, httpHandler http.Handler) http.Handler
 			grpcWeb:     grpcWeb,
 			httpHandler: httpHandler,
 		}.handle()
-	}), &http2.Server{})
+	}), http2Server), wg.Wait
 }
diff --git a/server.go b/server.go
index 5340527..1b5be57 100644
--- a/server.go
+++ b/server.go
@@ -9,6 +9,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"golang.org/x/net/http2"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/reflection"
 )
@@ -56,16 +57,21 @@ func (s *server) GetConnStateMap() map[net.Conn]http.ConnState {
 func (s *server) StartAndWait(ctx context.Context) error {
 	var connectionClose atomic.Bool
 
-	server := &http.Server{
-		Handler: newHandler(s.grpcServer, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			if connectionClose.Load() {
-				w.Header().Set("Connection", "close")
-			}
-			s.httpHandler.ServeHTTP(w, r)
-		})),
+	http2Server := &http2.Server{}
+	handler, wait := newHandler(s.grpcServer, http2Server, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if connectionClose.Load() {
+			w.Header().Set("Connection", "close")
+		}
+		s.httpHandler.ServeHTTP(w, r)
+	}))
+	httpServer := &http.Server{
+		Handler:        handler,
 		MaxHeaderBytes: s.config.MaxHeaderBytes,
 		ConnState:      s.updateConnState,
 	}
+	if err := http2.ConfigureServer(httpServer, http2Server); err != nil {
+		return fmt.Errorf("failed to configure HTTP/2 server: %v", err)
+	}
 	if ctx == nil {
 		ctx = context.Background()
 	}
@@ -77,7 +83,7 @@ func (s *server) StartAndWait(ctx context.Context) error {
 	done := make(chan error)
 	go func() {
 		defer close(done)
-		done <- server.Serve(listener)
+		done <- httpServer.Serve(listener)
 	}()
 
 	if s.config.OnStarted != nil {
@@ -102,8 +108,8 @@ func (s *server) StartAndWait(ctx context.Context) error {
 				break
 			}
 		}
-		s.grpcServer.GracefulStop()
-		_ = server.Shutdown(context.Background())
+		_ = httpServer.Shutdown(context.Background())
+		wait()
 		<-done
 		return nil
 	case err := <-done: