From 951ae39eba692606a956e43f8a98324bb5667201 Mon Sep 17 00:00:00 2001
From: Philipp Winter <phw@nymity.ch>
Date: Tue, 18 Aug 2020 10:07:37 -0700
Subject: [PATCH] Let Web server shut down gracefully.

That allows us to cache bridge lines of connections that were still open
while we received the signal to shut down.
---
 main.go | 42 ++++++++++++++++++++++++++++--------------
 1 file changed, 28 insertions(+), 14 deletions(-)

diff --git a/main.go b/main.go
index 62225cc..d03e809 100644
--- a/main.go
+++ b/main.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"io"
 	"log"
@@ -114,23 +115,36 @@ func main() {
 	if err := cache.ReadFromDisk(cacheFile); err != nil {
 		log.Printf("Could not read cache because: %s", err)
 	}
-	// When catching SIGINT or SIGTERM, write our cache to disk before exiting.
-	signalChan := make(chan os.Signal, 1)
-	signal.Notify(signalChan, syscall.SIGINT)
-	signal.Notify(signalChan, syscall.SIGTERM)
+
+	var srv http.Server
+	srv.Addr = addr
+	srv.Handler = NewRouter()
+	log.Printf("Starting service on port %s.", addr)
 	go func() {
-		<-signalChan
-		if err := cache.WriteToDisk(cacheFile); err != nil {
-			log.Printf("Could not write cache because: %s", err)
+		if certFilename != "" && keyFilename != "" {
+			srv.ListenAndServeTLS(certFilename, keyFilename)
+		} else {
+			srv.ListenAndServe()
 		}
-		os.Exit(1)
 	}()
 
-	router := NewRouter()
-	log.Printf("Starting service on port %s.", addr)
-	if certFilename != "" && keyFilename != "" {
-		log.Fatal(http.ListenAndServeTLS(addr, certFilename, keyFilename, router))
-	} else {
-		log.Fatal(http.ListenAndServe(addr, router))
+	signalChan := make(chan os.Signal, 1)
+	signal.Notify(signalChan, syscall.SIGINT)
+	signal.Notify(signalChan, syscall.SIGTERM)
+	log.Printf("Waiting for signal to shut down.")
+	<-signalChan
+
+	log.Printf("Received signal to shut down.")
+	// Give our Web server a maximum of a minute to shut down, and finish
+	// handling open connections.
+	t := time.Now().Add(time.Minute)
+	ctx, cancel := context.WithDeadline(context.Background(), t)
+	defer cancel()
+	if err := srv.Shutdown(ctx); err != nil {
+		log.Printf("Failed to shut down Web server: %s", err)
+	}
+
+	if err := cache.WriteToDisk(cacheFile); err != nil {
+		log.Printf("Failed to write cache to disk: %s", err)
 	}
 }
-- 
GitLab