Improve chat shutdown sequence and websocket cleanup
Refactors the shutdown logic in main.go to ensure proper order: cancels context, closes client, gracefully shuts down HTTP server, and waits for all websocket handler goroutines to finish. Adds a sync.WaitGroup to the hub in view.go to track active websocket handlers and ensures clean shutdown by waiting for all handlers to exit.
cognitive committed
Oct 21, 2025 at 22:37 UTC
244992fb4d87b716b4055811d40bb4405c1958b9
2 files changed
+30
-6
cmd/example_chat/main.go
+22
-6
@@ -7,6 +7,7 @@ import (
7
"os"
8
"os/signal"
9
"syscall"
10
+ "time"
11
12
"github.com/gosuda/relaydns/relaydns"
13
"github.com/rs/zerolog/log"
@@ -40,7 +41,6 @@ func main() {
41
42
func runChat(cmd *cobra.Command, args []string) error {
43
ctx, cancel := context.WithCancel(context.Background())
43
- defer cancel()
44
45
// 1) start local chat HTTP backend
46
ln, err := net.Listen("tcp", fmt.Sprintf(":%d", flagPort))
@@ -65,16 +65,32 @@ func runChat(cmd *cobra.Command, args []string) error {
65
if err := client.Start(ctx); err != nil {
66
return fmt.Errorf("start client: %w", err)
67
}
68
- defer client.Close()
68
69
sig := make(chan os.Signal, 1)
70
signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
71
<-sig
73
- log.Info().Msg("[chat] shutting down client...")
74
- if err := srv.Shutdown(ctx); err != nil {
75
- log.Error().Err(err).Msg("[chat] server forced to shutdown")
72
+ log.Info().Msg("[chat] shutting down...")
73
+
74
+ // Shutdown sequence:
75
+ // 1. Cancel context to stop client advertising/refresh loops
76
+ cancel()
77
+
78
+ // 2. Close client (waits for goroutines, closes libp2p host)
79
+ if err := client.Close(); err != nil {
80
+ log.Warn().Err(err).Msg("[chat] client close error")
81
}
77
- // ensure any active websocket conns are closed to stop goroutines
82
+
83
+ // 3. Shutdown HTTP server with a fresh context (with timeout)
84
+ shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
85
+ defer shutdownCancel()
86
+ if err := srv.Shutdown(shutdownCtx); err != nil {
87
+ log.Error().Err(err).Msg("[chat] http server shutdown error")
88
+ }
89
+
90
+ // 4. Close all websocket connections and wait for handlers to finish
91
hub.closeAll()
92
+ hub.wait()
93
+
94
+ log.Info().Msg("[chat] shutdown complete")
95
return nil
96
}
cmd/example_chat/view.go
+8
@@ -19,6 +19,7 @@ type hub struct {
19
messages []message
20
conns map[*websocket.Conn]struct{}
21
names map[*websocket.Conn]string
22
+ wg sync.WaitGroup
23
}
24
25
type message struct {
@@ -60,6 +61,11 @@ func (h *hub) closeAll() {
61
}
62
}
63
64
+// wait blocks until all websocket handler goroutines have finished.
65
+func (h *hub) wait() {
66
+ h.wg.Wait()
67
+}
68
+
69
func handleWS(w http.ResponseWriter, r *http.Request, h *hub) {
70
conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{
71
// Allow any origin for demo simplicity. Consider tightening in production.
@@ -80,6 +86,7 @@ func handleWS(w http.ResponseWriter, r *http.Request, h *hub) {
86
for _, m := range backlog {
87
_ = wsjson.Write(connCtx, conn, m)
88
}
89
+ h.wg.Add(1)
90
go func() {
91
defer func() {
92
var leftUser string
@@ -95,6 +102,7 @@ func handleWS(w http.ResponseWriter, r *http.Request, h *hub) {
102
}
103
conn.Close(websocket.StatusNormalClosure, "")
104
cancelConn()
105
+ h.wg.Done()
106
}()
107
for {
108
var req struct {