diff --git a/cmd/relaysrv/listener.go b/cmd/relaysrv/listener.go index 1e6610f2..f73fa04c 100644 --- a/cmd/relaysrv/listener.go +++ b/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) } diff --git a/cmd/relaysrv/main.go b/cmd/relaysrv/main.go index 02340147..66d452ca 100644 --- a/cmd/relaysrv/main.go +++ b/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() {