From 0ead6c107ad1a6ac9732a45d346c83810fa148ec Mon Sep 17 00:00:00 2001
From: Alex Guerrieri <klaidliadon@gmail.com>
Date: Thu, 14 Dec 2017 13:54:50 +0000
Subject: [PATCH] Graceful shutdown

---
 tent/cmd/run.go | 23 ++++++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/tent/cmd/run.go b/tent/cmd/run.go
index e38409b..6ea7db0 100644
--- a/tent/cmd/run.go
+++ b/tent/cmd/run.go
@@ -15,10 +15,14 @@
 package cmd
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"log"
+	"net/http"
 	"os"
+	"os/signal"
+	"time"
 
 	"github.com/gin-gonic/gin"
 	"github.com/securityfirst/tent"
@@ -36,6 +40,10 @@ var runCmd = &cobra.Command{
 			os.Exit(1)
 		}
 		e := gin.Default()
+		srv := &http.Server{
+			Addr:    fmt.Sprintf(":%v", config.Port),
+			Handler: e,
+		}
 		r, err := newRepo()
 		if err != nil {
 			log.Fatalf("Repo error: %s", err)
@@ -43,7 +51,20 @@ var runCmd = &cobra.Command{
 
 		o := tent.New(r)
 		o.Register(e.Group("/v2"), config.Config)
-		e.Run(fmt.Sprintf(":%v", config.Port))
+
+		stop := make(chan os.Signal, 1)
+		signal.Notify(stop, os.Interrupt)
+		go func() {
+			if err := e.Run(fmt.Sprintf(":%v", config.Port)); err != nil {
+				log.Fatal(err)
+			}
+		}()
+
+		<-stop
+		log.Println("Shutting down the server...")
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		srv.Shutdown(ctx)
 	},
 }