From 035b0896523873c4ede7df792e2563fdeacad47a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E6=B4=8B?= <liuyang@medlinker.com>
Date: Fri, 10 Apr 2020 15:53:11 +0800
Subject: [PATCH] graceful shutdown

---
 main.go | 41 ++++++++++++++++++++++++++++++++++++++---
 1 file changed, 38 insertions(+), 3 deletions(-)

diff --git a/main.go b/main.go
index d1da3880..f6ea9283 100644
--- a/main.go
+++ b/main.go
@@ -1,12 +1,20 @@
 package main
 
 import (
+	"context"
+	"log"
+	"net/http"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+
 	"github.com/gin-gonic/gin"
+
 	"go-admin/config"
 	orm "go-admin/database"
 	"go-admin/models"
 	"go-admin/router"
-	"log"
 )
 
 // @title go-admin API
@@ -35,8 +43,35 @@ func main() {
 	r := router.InitRouter()
 
 	defer orm.Eloquent.Close()
-	if err := r.Run(config.ApplicationConfig.Host + ":" + config.ApplicationConfig.Port); err != nil {
-		log.Fatal(err)
+
+	server := &http.Server{
+		Addr:              config.ApplicationConfig.Host + ":" + config.ApplicationConfig.Port,
+		Handler:           r,
 	}
 
+	go func() {
+		err := server.ListenAndServe(); if err != nil && err != http.ErrServerClosed {
+			log.Fatalf("listen error: %s\n", err)
+		}
+	}()
+
+	// graceful shutdown
+	c := make(chan os.Signal)
+
+	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
+	<-c
+	log.Println("Shutdown Server ...")
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	if err := server.Shutdown(ctx); err != nil {
+		log.Fatal("Server Shutdown:", err)
+	}
+	// catching ctx.Done(). timeout of 5 seconds.
+	select {
+	case <-ctx.Done():
+		log.Println("timeout of 5 seconds.")
+	}
+	log.Println("Server exiting")
+
 }