Add signal handlers (fixes #15)

cmd/relaysrv/listener.go

@@ -252,6 +252,10 @@ func protocolConnectionHandler(tcpConn net.Conn, config *tls.Config) {
 			conn.Close()
 
 		case msg := <-outbox:
+			if msg == nil {
+				conn.Close()
+				return
+			}
 			if debug {
 				log.Printf("Sending message %T to %s", msg, id)
 			}

cmd/relaysrv/main.go

@@ -9,10 +9,13 @@ import (
 	"log"
 	"net"
 	"net/url"
+	"os"
+	"os/signal"
 	"path/filepath"
 	"runtime"
 	"strings"
 	"sync/atomic"
+	"syscall"
 	"time"
 
 	"github.com/juju/ratelimit"
@@ -158,7 +161,32 @@ func main() {
 		}
 	}
 
-	listener(listen, tlsCfg)
+	go listener(listen, tlsCfg)
+
+	sigs := make(chan os.Signal, 1)
+	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
+	<-sigs
+
+	// Gracefully close all connections, hoping that clients will be faster
+	// to realize that the relay is now gone.
+
+	sessionMut.RLock()
+	for _, session := range activeSessions {
+		session.CloseConns()
+	}
+
+	for _, session := range pendingSessions {
+		session.CloseConns()
+	}
+	sessionMut.RUnlock()
+
+	outboxesMut.RLock()
+	for _, outbox := range outboxes {
+		close(outbox)
+	}
+	outboxesMut.RUnlock()
+
+	time.Sleep(500 * time.Millisecond)
 }
 
 func monitorLimits() {