From 5f1d429c1280e4cfd11a436e55e86c7db79c3ea9 Mon Sep 17 00:00:00 2001
From: Andrew Gutekanst <andrew.gutekanst@gmail.com>
Date: Mon, 13 Jan 2020 18:36:55 -0500
Subject: [PATCH] Add graceful shutdown to channel server

---
 main.go                                |  4 +++
 server/channelserver/channel_server.go | 37 +++++++++++++++++++++++---
 server/channelserver/session.go        |  9 ++++---
 3 files changed, 43 insertions(+), 7 deletions(-)

diff --git a/main.go b/main.go
index f197959bf..3490f8ea4 100644
--- a/main.go
+++ b/main.go
@@ -6,6 +6,7 @@ import (
 	"os"
 	"os/signal"
 	"syscall"
+	"time"
 
 	"github.com/Andoryuuta/Erupe/config"
 	"github.com/Andoryuuta/Erupe/server/channelserver"
@@ -113,7 +114,10 @@ func main() {
 	<-c
 
 	logger.Info("Trying to shutdown gracefully.")
+	channelServer.Shutdown()
 	signServer.Shutdown()
 	entranceServer.Shutdown()
 	launcherServer.Shutdown()
+
+	time.Sleep(5 * time.Second)
 }
diff --git a/server/channelserver/channel_server.go b/server/channelserver/channel_server.go
index e70d922fc..5c43e7e55 100644
--- a/server/channelserver/channel_server.go
+++ b/server/channelserver/channel_server.go
@@ -26,8 +26,9 @@ type Server struct {
 	acceptConns chan net.Conn
 	deleteConns chan net.Conn
 	sessions    map[net.Conn]*Session
-	listenAddr  string
 	listener    net.Listener // Listener that is created when Server.Start is called.
+
+	isShuttingDown bool
 }
 
 // NewServer creates a new Server type.
@@ -57,13 +58,30 @@ func (s *Server) Start() error {
 	return nil
 }
 
+// Shutdown tries to shut down the server gracefully.
+func (s *Server) Shutdown() {
+	s.Lock()
+	s.isShuttingDown = true
+	s.Unlock()
+
+	s.listener.Close()
+	close(s.acceptConns)
+}
+
 func (s *Server) acceptClients() {
 	for {
 		conn, err := s.listener.Accept()
 		if err != nil {
-			// TODO(Andoryuuta): Implement shutdown logic to end this goroutine cleanly here.
-			fmt.Println(err)
-			continue
+			s.Lock()
+			shutdown := s.isShuttingDown
+			s.Unlock()
+
+			if shutdown {
+				break
+			} else {
+				s.logger.Warn("Error accepting client", zap.Error(err))
+				continue
+			}
 		}
 		s.acceptConns <- conn
 	}
@@ -73,6 +91,17 @@ func (s *Server) manageSessions() {
 	for {
 		select {
 		case newConn := <-s.acceptConns:
+			// Gracefully handle acceptConns channel closing.
+			if newConn == nil {
+				s.Lock()
+				shutdown := s.isShuttingDown
+				s.Unlock()
+
+				if shutdown {
+					return
+				}
+			}
+
 			session := NewSession(s, newConn)
 
 			s.Lock()
diff --git a/server/channelserver/session.go b/server/channelserver/session.go
index 0e95d1b14..0e0a0e079 100644
--- a/server/channelserver/session.go
+++ b/server/channelserver/session.go
@@ -10,11 +10,13 @@ import (
 	"github.com/Andoryuuta/Erupe/network"
 	"github.com/Andoryuuta/Erupe/network/mhfpacket"
 	"github.com/Andoryuuta/byteframe"
+	"go.uber.org/zap"
 )
 
 // Session holds state for the channel server connection.
 type Session struct {
 	sync.Mutex
+	logger    *zap.Logger
 	server    *Server
 	rawConn   net.Conn
 	cryptConn *network.CryptConn
@@ -23,6 +25,7 @@ type Session struct {
 // NewSession creates a new Session type.
 func NewSession(server *Server, conn net.Conn) *Session {
 	s := &Session{
+		logger:    server.logger,
 		server:    server,
 		rawConn:   conn,
 		cryptConn: network.NewCryptConn(conn),
@@ -33,15 +36,15 @@ func NewSession(server *Server, conn net.Conn) *Session {
 // Start starts the session packet read&handle loop.
 func (s *Session) Start() {
 	go func() {
-		fmt.Println("Channel server got connection!")
+		s.logger.Info("Channel server got connection!")
+
 		// Unlike the sign and entrance server,
 		// the client DOES NOT initalize the channel connection with 8 NULL bytes.
 
 		for {
 			pkt, err := s.cryptConn.ReadPacket()
 			if err != nil {
-				fmt.Println(err)
-				fmt.Println("Error on channel server readpacket")
+				s.logger.Warn("Error on channel server readpacket", zap.Error(err))
 				return
 			}