Add graceful shutdown to channel server

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()

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
 			}