From f99363674b6fe13ff14bf3fb93f9a1d1b53caf14 Mon Sep 17 00:00:00 2001
From: SALLEYRON Julien <julien.salleyron@gmail.com>
Date: Wed, 14 Mar 2018 13:14:03 +0100
Subject: [PATCH] Use context in Server

---
 cmd/context.go                   | 22 ++++++++++++++++++++++
 cmd/traefik/traefik.go           |  4 +++-
 server/server.go                 | 22 +++++++++++++++++++---
 server/server_signals.go         | 12 +-----------
 server/server_signals_windows.go | 23 ++---------------------
 5 files changed, 47 insertions(+), 36 deletions(-)
 create mode 100644 cmd/context.go

diff --git a/cmd/context.go b/cmd/context.go
new file mode 100644
index 000000000..002db47ec
--- /dev/null
+++ b/cmd/context.go
@@ -0,0 +1,22 @@
+package cmd
+
+import (
+	"context"
+	"os"
+	"os/signal"
+	"syscall"
+)
+
+// ContextWithSignal create a context cancelled when SIGINT or SIGTERM are notified
+func ContextWithSignal(ctx context.Context) context.Context {
+	newCtx, cancel := context.WithCancel(ctx)
+	signals := make(chan os.Signal)
+	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
+	go func() {
+		select {
+		case <-signals:
+			cancel()
+		}
+	}()
+	return newCtx
+}
diff --git a/cmd/traefik/traefik.go b/cmd/traefik/traefik.go
index 4e6fe83e5..02f98ca16 100644
--- a/cmd/traefik/traefik.go
+++ b/cmd/traefik/traefik.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"encoding/json"
 	fmtlog "log"
 	"net/http"
@@ -173,7 +174,8 @@ func runCmd(globalConfiguration *configuration.GlobalConfiguration, configFile s
 		acme.Get().SetConfigListenerChan(make(chan types.Configuration))
 		svr.AddListener(acme.Get().ListenConfiguration)
 	}
-	svr.Start()
+	ctx := cmd.ContextWithSignal(context.Background())
+	svr.StartWithContext(ctx)
 	defer svr.Close()
 
 	sent, err := daemon.SdNotify(false, "READY=1")
diff --git a/server/server.go b/server/server.go
index 299688e87..7a71a8b38 100644
--- a/server/server.go
+++ b/server/server.go
@@ -205,6 +205,23 @@ func (s *Server) Start() {
 	go s.listenSignals()
 }
 
+// StartWithContext starts the server and Stop/Close it when context is Done
+func (s *Server) StartWithContext(ctx context.Context) {
+	go func() {
+		defer s.Close()
+		<-ctx.Done()
+		log.Info("I have to go...")
+		reqAcceptGraceTimeOut := time.Duration(s.globalConfiguration.LifeCycle.RequestAcceptGraceTimeout)
+		if reqAcceptGraceTimeOut > 0 {
+			log.Infof("Waiting %s for incoming requests to cease", reqAcceptGraceTimeOut)
+			time.Sleep(reqAcceptGraceTimeOut)
+		}
+		log.Info("Stopping server gracefully")
+		s.Stop()
+	}()
+	s.Start()
+}
+
 // Wait blocks until server is shutted down.
 func (s *Server) Wait() {
 	<-s.stopChan
@@ -235,14 +252,13 @@ func (s *Server) Stop() {
 
 // Close destroys the server
 func (s *Server) Close() {
-	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(s.globalConfiguration.LifeCycle.GraceTimeOut))
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	go func(ctx context.Context) {
 		<-ctx.Done()
 		if ctx.Err() == context.Canceled {
 			return
 		} else if ctx.Err() == context.DeadlineExceeded {
-			log.Warn("Timeout while stopping traefik, killing instance ✝")
-			os.Exit(1)
+			panic("Timeout while stopping traefik, killing instance ✝")
 		}
 	}(ctx)
 	stopMetricsClients()
diff --git a/server/server_signals.go b/server/server_signals.go
index 00eb6d4af..d024a627a 100644
--- a/server/server_signals.go
+++ b/server/server_signals.go
@@ -5,13 +5,12 @@ package server
 import (
 	"os/signal"
 	"syscall"
-	"time"
 
 	"github.com/containous/traefik/log"
 )
 
 func (s *Server) configureSignals() {
-	signal.Notify(s.signals, syscall.SIGINT, syscall.SIGTERM, syscall.SIGUSR1)
+	signal.Notify(s.signals, syscall.SIGUSR1)
 }
 
 func (s *Server) listenSignals() {
@@ -30,15 +29,6 @@ func (s *Server) listenSignals() {
 			if err := log.RotateFile(); err != nil {
 				log.Errorf("Error rotating traefik log: %s", err)
 			}
-		default:
-			log.Infof("I have to go... %+v", sig)
-			reqAcceptGraceTimeOut := time.Duration(s.globalConfiguration.LifeCycle.RequestAcceptGraceTimeout)
-			if reqAcceptGraceTimeOut > 0 {
-				log.Infof("Waiting %s for incoming requests to cease", reqAcceptGraceTimeOut)
-				time.Sleep(reqAcceptGraceTimeOut)
-			}
-			log.Info("Stopping server gracefully")
-			s.Stop()
 		}
 	}
 }
diff --git a/server/server_signals_windows.go b/server/server_signals_windows.go
index b420f20b5..674896b43 100644
--- a/server/server_signals_windows.go
+++ b/server/server_signals_windows.go
@@ -2,25 +2,6 @@
 
 package server
 
-import (
-	"os/signal"
-	"syscall"
+func (s *Server) configureSignals() {}
 
-	"github.com/containous/traefik/log"
-)
-
-func (s *Server) configureSignals() {
-	signal.Notify(s.signals, syscall.SIGINT, syscall.SIGTERM)
-}
-
-func (s *Server) listenSignals() {
-	for {
-		sig := <-s.signals
-		switch sig {
-		default:
-			log.Infof("I have to go... %+v", sig)
-			log.Info("Stopping server")
-			s.Stop()
-		}
-	}
-}
+func (s *Server) listenSignals() {}