More shutdown work, closes #11

- Drain SMTP connections
- Force exit after 15 seconds of draining

smtpd/handler.go

@@ -90,8 +90,11 @@ func (ss *Session) String() string {
 func (s *Server) startSession(id int, conn net.Conn) {
 	log.Info("Connection from %v, starting session <%v>", conn.RemoteAddr(), id)
 	expConnectsCurrent.Add(1)
-	defer conn.Close()
-	defer expConnectsCurrent.Add(-1)
+	defer func() {
+		conn.Close()
+		s.waitgroup.Done()
+		expConnectsCurrent.Add(-1)
+	}()
 
 	ss := NewSession(s, id, conn)
 	ss.greet()
@@ -300,7 +303,7 @@ func (ss *Session) dataHandler() {
 		i := 0
 		for e := ss.recipients.Front(); e != nil; e = e.Next() {
 			recip := e.Value.(string)
-			if !strings.HasSuffix(strings.ToLower(recip), "@" + ss.server.domainNoStore) {
+			if !strings.HasSuffix(strings.ToLower(recip), "@"+ss.server.domainNoStore) {
 				// Not our "no store" domain, so store the message
 				mb, err := ss.server.dataStore.MailboxFor(recip)
 				if err != nil {

smtpd/listener.go

@@ -8,6 +8,7 @@ import (
 	"github.com/jhillyerd/inbucket/log"
 	"net"
 	"strings"
+	"sync"
 	"time"
 )
 
@@ -22,6 +23,7 @@ type Server struct {
 	storeMessages   bool
 	listener        net.Listener
 	shutdown        bool
+	waitgroup       *sync.WaitGroup
 }
 
 // Raw stat collectors
@@ -49,7 +51,8 @@ func New() *Server {
 	cfg := config.GetSmtpConfig()
 	return &Server{dataStore: ds, domain: cfg.Domain, maxRecips: cfg.MaxRecipients,
 		maxIdleSeconds: cfg.MaxIdleSeconds, maxMessageBytes: cfg.MaxMessageBytes,
-		storeMessages: cfg.StoreMessages, domainNoStore: strings.ToLower(cfg.DomainNoStore)}
+		storeMessages: cfg.StoreMessages, domainNoStore: strings.ToLower(cfg.DomainNoStore),
+		waitgroup: new(sync.WaitGroup)}
 }
 
 // Main listener loop
@@ -109,17 +112,25 @@ func (s *Server) Start() {
 		} else {
 			tempDelay = 0
 			expConnectsTotal.Add(1)
+			s.waitgroup.Add(1)
 			go s.startSession(sid, conn)
 		}
 	}
 }
 
+// Stop requests the SMTP server closes it's listener
 func (s *Server) Stop() {
-	log.Trace("SMTP shutdown requested")
+	log.Trace("SMTP shutdown requested, connections will be drained")
 	s.shutdown = true
 	s.listener.Close()
 }
 
+// Drain causes the caller to block until all active SMTP sessions have finished
+func (s *Server) Drain() {
+	s.waitgroup.Wait()
+	log.Trace("SMTP connections drained")
+}
+
 // When the provided Ticker ticks, we update our metrics history
 func metricsTicker(t *time.Ticker) {
 	ok := true