Gracefully shut down the server

1 files changed, 29 insertions, 6 deletions

diff --git a/cmd/gcp/main.go b/cmd/gcp/main.go
index 19823ba..80f3114 100644
--- a/cmd/gcp/main.go
+++ b/cmd/gcp/main.go
@@ -6,11 +6,16 @@
 package main
 
 import (
-	"go.awhk.org/go-import-redirect/internal"
+	"context"
 	"log"
 	"net/http"
 	"os"
+	"os/signal"
 	"path"
+	"syscall"
+	"time"
+
+	"go.awhk.org/go-import-redirect/internal"
 )
 
 func redirect(resp http.ResponseWriter, req *http.Request) {
@@ -32,10 +37,28 @@ func redirect(resp http.ResponseWriter, req *http.Request) {
 }
 
 func main() {
-	port := os.Getenv("PORT")
-	if port == "" {
-		port = "8080"
+	mux := http.NewServeMux()
+	mux.HandleFunc("/", redirect)
+	srv := http.Server{Handler: mux}
+
+	done := make(chan os.Signal, 1)
+	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
+
+	go func() {
+		port := os.Getenv("PORT")
+		if port == "" {
+			port = "8080"
+		}
+		srv.Addr = os.Getenv("ADDR") + ":" + port
+		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
+			log.Fatalln("server.ListenAndServe:", err)
+		}
+	}()
+
+	<-done
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	if err := srv.Shutdown(ctx); err != nil && err != http.ErrServerClosed {
+		log.Fatalln("server.Shutdown:", err)
 	}
-	http.HandleFunc("/", redirect)
-	log.Fatal(http.ListenAndServe(os.Getenv("ADDR")+":"+port, nil))
 }