Fix shutdown when proxy is killed by a signal (#83)

* Fix shutdown when proxy is killed by a signal * Ignore http server closed error
datastax · Feb 15, 2022 · 724d53d · 724d53d
1 parent 7d5dd17
commit 724d53d
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/proxy/run.go b/proxy/run.go
@@ -264,6 +264,11 @@ func listenAndServe(p *Proxy, ctx context.Context, logger *zap.Logger) (err erro
 
 	wg.Add(numServers)
 
+	go func() {
+		wg.Wait()
+		close(ch)
+	}()
+
 	go func() {
 		select {
 		case <-ctx.Done():
@@ -275,7 +280,7 @@ func listenAndServe(p *Proxy, ctx context.Context, logger *zap.Logger) (err erro
 
 	go func() {
 		defer wg.Done()
-		err = p.Serve()
+		err := p.Serve()
 		if err != nil && err != ErrProxyClosed {
 			ch <- err
 		}
@@ -284,8 +289,8 @@ func listenAndServe(p *Proxy, ctx context.Context, logger *zap.Logger) (err erro
 	if cli.HealthCheck {
 		go func() {
 			defer wg.Done()
-			err = server.Serve(listener)
-			if err != nil {
+			err := server.Serve(listener)
+			if err != nil && err != http.ErrServerClosed {
 				ch <- err
 			}
 		}()