From 7ccef7b97747325724e2c138ccc52427cbf20ebd Mon Sep 17 00:00:00 2001
From: James Hillyerd <james@hillyerd.com>
Date: Fri, 16 Nov 2012 20:44:41 -0800
Subject: [PATCH] Basic SIGTERM handling

---
 inbucket.go       | 17 ++++++++++++++---
 smtpd/listener.go | 41 +++++++++++++++++++++++++++++++++++------
 web/server.go     | 31 +++++++++++++++++++++++++++----
 3 files changed, 76 insertions(+), 13 deletions(-)

diff --git a/inbucket.go b/inbucket.go
index 81170c0..3958d5d 100644
--- a/inbucket.go
+++ b/inbucket.go
@@ -29,6 +29,8 @@ var startTime = time.Now()
 // The file we send log output to, will be nil for stderr or stdout
 var logf *os.File
 
+var smtpServer *smtpd.Server
+
 func main() {
 	flag.Parse()
 	if *help {
@@ -49,7 +51,7 @@ func main() {
 
 	// Setup signal handler
 	sigChan := make(chan os.Signal)
-	signal.Notify(sigChan, syscall.SIGHUP)
+	signal.Notify(sigChan, syscall.SIGHUP, syscall.SIGTERM)
 	go signalProcessor(sigChan)
 
 	// Configure logging, close std* fds
@@ -91,8 +93,8 @@ func main() {
 	}
 
 	// Startup SMTP server
-	server := smtpd.New()
-	go server.Start()
+	smtpServer = smtpd.New()
+	go smtpServer.Start()
 
 	web.Start()
 }
@@ -130,6 +132,15 @@ func signalProcessor(c <-chan os.Signal) {
 			} else {
 				log.Info("Ignoring SIGHUP, logfile not configured")
 			}
+		case syscall.SIGTERM:
+			// Initiate shutdown
+			log.Info("Received SIGTERM, shutting down")
+			if smtpServer != nil {
+				smtpServer.Stop()
+			} else {
+				log.Error("smtpServer was nil during shutdown")
+			}
+			web.Stop()
 		}
 	}
 }
diff --git a/smtpd/listener.go b/smtpd/listener.go
index bb78afb..58dedb1 100644
--- a/smtpd/listener.go
+++ b/smtpd/listener.go
@@ -20,6 +20,8 @@ type Server struct {
 	maxMessageBytes int
 	dataStore       DataStore
 	storeMessages   bool
+	listener        net.Listener
+	shutdown        bool
 }
 
 // Raw stat collectors
@@ -62,9 +64,9 @@ func (s *Server) Start() {
 	}
 
 	log.Info("SMTP listening on TCP4 %v", addr)
-	ln, err := net.ListenTCP("tcp4", addr)
+	s.listener, err = net.ListenTCP("tcp4", addr)
 	if err != nil {
-		log.Error("Failed to start tcp4 listener: %v", err)
+		log.Error("SMTP failed to start tcp4 listener: %v", err)
 		// TODO More graceful early-shutdown procedure
 		panic(err)
 	}
@@ -79,18 +81,45 @@ func (s *Server) Start() {
 	StartRetentionScanner(s.dataStore)
 
 	// Handle incoming connections
+	var tempDelay time.Duration
 	for sid := 1; ; sid++ {
-		if conn, err := ln.Accept(); err != nil {
-			// TODO Implement a max error counter before shutdown?
-			// or maybe attempt to restart smtpd
-			panic(err)
+		if conn, err := s.listener.Accept(); err != nil {
+			if nerr, ok := err.(net.Error); ok && nerr.Temporary() {
+				// Temporary error, sleep for a bit and try again
+				if tempDelay == 0 {
+					tempDelay = 5 * time.Millisecond
+				} else {
+					tempDelay *= 2
+				}
+				if max := 1 * time.Second; tempDelay > max {
+					tempDelay = max
+				}
+				log.Error("SMTP accept error: %v; retrying in %v", err, tempDelay)
+				time.Sleep(tempDelay)
+				continue
+			} else {
+				if s.shutdown {
+					log.Trace("SMTP listener shutting down on request")
+					return
+				}
+				// TODO Implement a max error counter before shutdown?
+				// or maybe attempt to restart smtpd
+				panic(err)
+			}
 		} else {
+			tempDelay = 0
 			expConnectsTotal.Add(1)
 			go s.startSession(sid, conn)
 		}
 	}
 }
 
+func (s *Server) Stop() {
+	log.Trace("SMTP shutdown requested")
+	s.shutdown = true
+	s.listener.Close()
+}
+
 // When the provided Ticker ticks, we update our metrics history
 func metricsTicker(t *time.Ticker) {
 	ok := true
diff --git a/web/server.go b/web/server.go
index 531e741..0e21b52 100644
--- a/web/server.go
+++ b/web/server.go
@@ -10,6 +10,7 @@ import (
 	"github.com/gorilla/sessions"
 	"github.com/jhillyerd/inbucket/config"
 	"github.com/jhillyerd/inbucket/log"
+	"net"
 	"net/http"
 	"time"
 )
@@ -17,8 +18,9 @@ import (
 type handler func(http.ResponseWriter, *http.Request, *Context) error
 
 var Router *mux.Router
-
+var listener net.Listener
 var sessionStore sessions.Store
+var shutdown bool
 
 func setupRoutes(cfg config.WebConfig) {
 	log.Info("Theme templates mapped to '%v'", cfg.TemplateDir)
@@ -55,20 +57,41 @@ func Start() {
 	sessionStore = sessions.NewCookieStore([]byte("something-very-secret"))
 
 	addr := fmt.Sprintf("%v:%v", cfg.Ip4address, cfg.Ip4port)
-	log.Info("HTTP listening on TCP4 %v", addr)
-	s := &http.Server{
+	server := &http.Server{
 		Addr:         addr,
 		Handler:      nil,
 		ReadTimeout:  60 * time.Second,
 		WriteTimeout: 60 * time.Second,
 	}
 
-	err := s.ListenAndServe()
+	// We don't use ListenAndServe because it lacks a way to close the listener
+	log.Info("HTTP listening on TCP4 %v", addr)
+	var err error
+	listener, err = net.Listen("tcp", addr)
 	if err != nil {
+		log.Error("HTTP failed to start TCP4 listener: %v", err)
+		// TODO More graceful early-shutdown procedure
+		panic(err)
+	}
+
+	err = server.Serve(listener)
+	if shutdown {
+		log.Trace("HTTP server shutting down on request")
+	} else if err != nil {
 		log.Error("HTTP server failed: %v", err)
 	}
 }
 
+func Stop() {
+	log.Trace("HTTP shutdown requested")
+	shutdown = true
+	if listener != nil {
+		listener.Close()
+	} else {
+		log.Error("HTTP listener was nil during shutdown")
+	}
+}
+
 // ServeHTTP builds the context and passes onto the real handler
 func (h handler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 	// Create the context