From ad042dd9aa129f1952507d89da39250c4f4b981b Mon Sep 17 00:00:00 2001
From: Savely Krendelhoff <alicecandyom@proton.me>
Date: Sat, 23 Aug 2025 17:18:22 +0700
Subject: [PATCH] [PHASE-8] Implement proper graceful shutdown

---
 internal/server/tcp.go | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/internal/server/tcp.go b/internal/server/tcp.go
index 33f550e..3c91df0 100644
--- a/internal/server/tcp.go
+++ b/internal/server/tcp.go
@@ -23,7 +23,7 @@ type TCPServer struct {
 	listener          net.Listener
 	logger            *slog.Logger
 	wg                sync.WaitGroup
-	shutdown          chan struct{}
+	cancel            context.CancelFunc
 }
 
 // Option is a functional option for configuring TCPServer
@@ -44,7 +44,6 @@ func NewTCPServer(wisdomService *service.WisdomService, config *Config, opts ...
 		wisdomApplication: application.NewWisdomApplication(wisdomService),
 		decoder:           protocol.NewMessageDecoder(),
 		logger:            slog.Default(),
-		shutdown:          make(chan struct{}),
 	}
 
 	for _, opt := range opts {
@@ -65,14 +64,22 @@ func (s *TCPServer) Start(ctx context.Context) error {
 	s.listener = listener
 	s.logger.Info("tcp server started", "address", s.config.Address)
 
-	go s.acceptLoop(ctx)
+	// Create cancellable context for server lifecycle
+	serverCtx, cancel := context.WithCancel(ctx)
+	s.cancel = cancel
+
+	go s.acceptLoop(serverCtx)
 	return nil
 }
 
 // Stop gracefully stops the server
 func (s *TCPServer) Stop() error {
 	s.logger.Info("stopping tcp server")
-	close(s.shutdown)
+
+	// Cancel server context to stop accept loop and active connections
+	if s.cancel != nil {
+		s.cancel()
+	}
 
 	if s.listener != nil {
 		s.listener.Close()
@@ -95,8 +102,6 @@ func (s *TCPServer) Address() string {
 func (s *TCPServer) acceptLoop(ctx context.Context) {
 	for {
 		select {
-		case <-s.shutdown:
-			return
 		case <-ctx.Done():
 			return
 		default:
@@ -105,7 +110,7 @@ func (s *TCPServer) acceptLoop(ctx context.Context) {
 		rawConn, err := s.listener.Accept()
 		if err != nil {
 			select {
-			case <-s.shutdown:
+			case <-ctx.Done():
 				return
 			default:
 				s.logger.Error("accept error", sl.Err(err))