From 71f82674cff5a8a6614b748c7cb89d0581f53186 Mon Sep 17 00:00:00 2001
From: Andrew Harrison <andrew.harrison@returnpath.com>
Date: Mon, 23 Feb 2015 22:40:27 -0600
Subject: [PATCH] First stab at recognizing signals to kill app

---
 go-mta.go         | 17 ++++++++++++-----
 servicehandler.go | 35 +++++++++++++++++++++++++++++++++++
 smtpserver.go     | 31 ++++++++++++++++++++++++++-----
 3 files changed, 73 insertions(+), 10 deletions(-)
 create mode 100644 servicehandler.go

diff --git a/go-mta.go b/go-mta.go
index 43bfd9f..5f3a075 100644
--- a/go-mta.go
+++ b/go-mta.go
@@ -1,12 +1,13 @@
 package main
 
 import (
-	"fmt"
 	"log"
 	"io"
 	"io/ioutil"
 	"os"
+	"os/signal"
 	"runtime"
+	"syscall"
 )
 
 func Init(
@@ -56,13 +57,15 @@ func main() {
 	runtime.GOMAXPROCS(runtime.NumCPU())
 	Info.Println("Starting the Go MTA Server.\n")
 
+	service := NewServiceHandler()
+
 	// Start the Server listener
 	cl := NewConnectionListener(smtpServerChan)
 	go cl.start()
 
 	// Handle new server connections
 	for i := 0; i < SmtpServerConnectionCount; i++ {
-		go handleSmtpServerConnections(smtpServerChan, envelopeChan)
+		go service.handleSmtpServerConnections(smtpServerChan, envelopeChan)
 	}
 	for i := 0; i < DispatcherThreads; i++ {
 		go handleDispatcher(envelopeChan, smtpClientChan)
@@ -71,7 +74,11 @@ func main() {
 		go handleSmtpClientConnections(smtpClientChan)
 	}
 
-	// Die after input is read.
-	var input string
-	fmt.Scanln(&input)
+	// Handle SIGINT and SIGTERM.
+	signalCh := make(chan os.Signal)
+	signal.Notify(signalCh, syscall.SIGINT, syscall.SIGTERM)
+	Info.Println(<-signalCh)
+
+	// Stop the service gracefully.
+	service.Stop()
 }
diff --git a/servicehandler.go b/servicehandler.go
new file mode 100644
index 0000000..3e30de1
--- /dev/null
+++ b/servicehandler.go
@@ -0,0 +1,35 @@
+package main
+
+import (
+	"sync"
+)
+
+type ServiceHandler struct {
+	ch chan bool
+	wg *sync.WaitGroup
+}
+
+func NewServiceHandler() *ServiceHandler {
+	s := &ServiceHandler {
+		ch: make(chan bool),
+		wg: &sync.WaitGroup{},
+    }
+	s.wg.Add(1) // Add main wait. Closed on Stop().
+    return s
+}
+
+func (s *ServiceHandler) addWatchedProcess() {
+	s.wg.Add(1)
+}
+
+func (s *ServiceHandler) finishProcess() {
+	s.wg.Done()
+}
+
+// Stop the service by closing the service's channel.  Block until the service
+// is really stopped.
+func (s *ServiceHandler) Stop() {
+	close(s.ch)
+	s.wg.Done()
+	s.wg.Wait()
+}
diff --git a/smtpserver.go b/smtpserver.go
index 356294a..2b7e59d 100644
--- a/smtpserver.go
+++ b/smtpserver.go
@@ -8,24 +8,45 @@ import (
 	"bytes"
 	"net"
 	"strings"
+//	"time"
 )
 
 type SmtpServer struct {
 	conn net.Conn
 }
 
-func handleSmtpServerConnections(smtpServerChan chan *SmtpServer, envelopeChan chan *envelope) {
+func (sh ServiceHandler) handleSmtpServerConnections(smtpServerChan chan *SmtpServer,
+	                                                 envelopeChan chan *envelope) {
 	Info.Println("SmtpServer Connection Handler Started.")
 	for {
-		// handle an SMTP Conversation
-		server := <- smtpServerChan
-		Info.Println("Received new SmtpServer, processing inbound SMTP Conversation.")
-		receiveSmtp(server.conn, envelopeChan)
+		select {
+		case <-sh.ch:
+			Info.Println("Terminating SmtpServer Connection Handler.")
+			return
+		case server := <- smtpServerChan:
+			// handle an SMTP Conversation
+			sh.addWatchedProcess()
+			Info.Println("Received new SmtpServer, processing inbound SMTP Conversation.")
+			receiveSmtp(server.conn, envelopeChan)
+			sh.finishProcess()
+		}
 	}
 }
 
+//func receiveWithTimeout(conn net.Conn, envelopeChan chan *envelope) {
+//
+//	timeoutChan := make(chan bool, 1)
+//	go func() {
+//		time.Sleep(3 * time.Second)
+//		timeoutChan <- true
+//	}()
+//
+//	receiveSmtp(conn, envelopeChan, timeoutChan)
+//}
+
 func receiveSmtp(conn net.Conn, envelopeChan chan *envelope) {
 
+
 	remoteIp := conn.RemoteAddr()
 	Info.Println("Received Connection from [", remoteIp, "]")