add graceful shutdown
Kim committed
Oct 21, 2025 at 21:19 UTC
5eab26ae546d6212bc77bd0a965509324bc588b8
5 files changed
+97
-27
cmd/example_chat/main.go
+7
-2
@@ -48,7 +48,7 @@ func runChat(cmd *cobra.Command, args []string) error {
48
return fmt.Errorf("listen chat: %w", err)
49
}
50
hub := newHub()
51
- go serveChatHTTP(ctx, ln, flagName, hub, cancel)
51
+ srv := serveChatHTTP(ln, flagName, hub)
52
53
// 2) advertise over RelayDNS (HTTP tunneled via server /peer route)
54
client, err := relaydns.NewClient(ctx, relaydns.ClientConfig{
@@ -70,6 +70,11 @@ func runChat(cmd *cobra.Command, args []string) error {
70
sig := make(chan os.Signal, 1)
71
signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
72
<-sig
73
- log.Info().Msg("[chat] shutting down")
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")
76
+ }
77
+ // ensure any active websocket conns are closed to stop goroutines
78
+ hub.closeAll()
79
return nil
80
}
cmd/example_chat/view.go
+61
-12
@@ -18,16 +18,18 @@ type hub struct {
18
mu sync.RWMutex
19
messages []message
20
conns map[*websocket.Conn]struct{}
21
+ names map[*websocket.Conn]string
22
}
23
24
type message struct {
24
- TS time.Time `json:"ts"`
25
- User string `json:"user"`
26
- Text string `json:"text"`
25
+ TS time.Time `json:"ts"`
26
+ User string `json:"user"`
27
+ Text string `json:"text"`
28
+ Event string `json:"event,omitempty"` // "joined" | "left"
29
}
30
31
func newHub() *hub {
30
- return &hub{conns: map[*websocket.Conn]struct{}{}, messages: make([]message, 0, 64)}
32
+ return &hub{conns: map[*websocket.Conn]struct{}{}, names: map[*websocket.Conn]string{}, messages: make([]message, 0, 64)}
33
}
34
35
func (h *hub) broadcast(m message) {
@@ -45,6 +47,19 @@ func (h *hub) broadcast(m message) {
47
}
48
}
49
50
+// closeAll force-closes all active websocket connections (used during shutdown).
51
+func (h *hub) closeAll() {
52
+ h.mu.Lock()
53
+ conns := make([]*websocket.Conn, 0, len(h.conns))
54
+ for c := range h.conns {
55
+ conns = append(conns, c)
56
+ }
57
+ h.mu.Unlock()
58
+ for _, c := range conns {
59
+ _ = c.Close(websocket.StatusGoingAway, "server shutdown")
60
+ }
61
+}
62
+
63
func handleWS(w http.ResponseWriter, r *http.Request, h *hub) {
64
conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{
65
// Allow any origin for demo simplicity. Consider tightening in production.
@@ -67,9 +82,17 @@ func handleWS(w http.ResponseWriter, r *http.Request, h *hub) {
82
}
83
go func() {
84
defer func() {
85
+ var leftUser string
86
h.mu.Lock()
87
+ if name, ok := h.names[conn]; ok && name != "" {
88
+ leftUser = name
89
+ }
90
+ delete(h.names, conn)
91
delete(h.conns, conn)
92
h.mu.Unlock()
93
+ if leftUser != "" {
94
+ h.broadcast(message{TS: time.Now().UTC(), User: leftUser, Event: "left"})
95
+ }
96
conn.Close(websocket.StatusNormalClosure, "")
97
cancelConn()
98
}()
@@ -84,6 +107,17 @@ func handleWS(w http.ResponseWriter, r *http.Request, h *hub) {
107
if req.User == "" {
108
req.User = "anon"
109
}
110
+ // first frame per connection: remember name and announce join
111
+ var announce bool
112
+ h.mu.Lock()
113
+ if _, ok := h.names[conn]; !ok {
114
+ h.names[conn] = req.User
115
+ announce = true
116
+ }
117
+ h.mu.Unlock()
118
+ if announce {
119
+ h.broadcast(message{TS: time.Now().UTC(), User: req.User, Event: "joined"})
120
+ }
121
if req.Text == "" {
122
continue
123
}
@@ -97,16 +131,22 @@ func serveIndex(w http.ResponseWriter, r *http.Request, name string) {
131
_ = indexTmpl.Execute(w, struct{ Name string }{Name: name})
132
}
133
100
-// serveChatHTTP hosts the chat UI and websocket endpoint on the provided listener.
101
-func serveChatHTTP(ctx context.Context, ln net.Listener, name string, h *hub, cancel context.CancelFunc) {
134
+// serveChatHTTP starts serving the chat UI and websocket endpoint and returns the server.
135
+// Callers are responsible for shutting it down via Server.Shutdown.
136
+func serveChatHTTP(ln net.Listener, name string, h *hub) *http.Server {
137
mux := http.NewServeMux()
138
mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { serveIndex(w, r, name) })
139
mux.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) { handleWS(w, r, h) })
140
+
141
+ srv := &http.Server{Handler: mux, ReadHeaderTimeout: 5 * time.Second, IdleTimeout: 60 * time.Second}
142
log.Info().Msgf("[chat] http listening on %s", ln.Addr().String())
106
- if err := http.Serve(ln, mux); err != nil && err != http.ErrServerClosed {
107
- log.Error().Err(err).Msg("chat http error")
108
- cancel()
109
- }
143
+
144
+ go func() {
145
+ if err := srv.Serve(ln); err != nil && err != http.ErrServerClosed {
146
+ log.Error().Err(err).Msg("chat http error")
147
+ }
148
+ }()
149
+ return srv
150
}
151
152
var indexTmpl = template.Must(template.New("chat").Parse(`<!DOCTYPE html>
@@ -229,8 +269,14 @@ var indexTmpl = template.Must(template.New("chat").Parse(`<!DOCTYPE html>
269
const ts = new Date(msg.ts).toLocaleTimeString();
270
const nick = (msg.user || 'anon');
271
const color = colorFor(nick);
232
- div.innerHTML = '<span class="ts">[' + ts + ']</span> <span class="usr" style="color:' + color + '">' +
233
- nick + '</span>: ' + escapeHTML(msg.text || '');
272
+ if (msg.event === 'joined' || msg.event === 'left') {
273
+ const verb = msg.event === 'joined' ? 'joined' : 'left';
274
+ div.innerHTML = '<span class="ts">[' + ts + ']</span> <span class="usr" style="color:' + color + '">' + nick + '</span> ' + verb;
275
+ div.style.opacity = '0.8';
276
+ } else {
277
+ div.innerHTML = '<span class="ts">[' + ts + ']</span> <span class="usr" style="color:' + color + '">' +
278
+ nick + '</span>: ' + escapeHTML(msg.text || '');
279
+ }
280
log.appendChild(div);
281
log.scrollTop = log.scrollHeight;
282
}
@@ -242,6 +288,9 @@ var indexTmpl = template.Must(template.New("chat").Parse(`<!DOCTYPE html>
288
const basePath = location.pathname.endsWith('/') ? location.pathname : (location.pathname + '/');
289
const ws = new WebSocket(wsProto + '://' + location.host + basePath + 'ws');
290
ws.onmessage = (e) => { try{ append(JSON.parse(e.data)); }catch(_){ } };
291
+ ws.onopen = () => {
292
+ try{ ws.send(JSON.stringify({ user: (user.value || 'anon'), text: '' })); }catch(_){ }
293
+ };
294
function send(){
295
const payload = { user: (user.value || 'anon'), text: cmd.value.trim() };
296
if(!payload.text) return;
cmd/example_http_client/main.go
+9
-6
@@ -8,7 +8,6 @@ import (
8
"os"
9
"os/signal"
10
"syscall"
11
- "time"
11
12
"github.com/gosuda/relaydns/relaydns"
13
"github.com/rs/zerolog/log"
@@ -60,13 +59,13 @@ func runClient(cmd *cobra.Command, args []string) error {
59
}
60
log.Info().Msgf("[client] local backend http listening on %s", ln.Addr().String())
61
63
- // Serve local backend view in a goroutine
64
- go serveClientHTTP(ctx, ln, flagName, func() string {
62
+ // Serve local backend view and keep a server handle for shutdown
63
+ srv := serveClientHTTP(ln, flagName, func() string {
64
if clientRef == nil {
65
return "Starting..."
66
}
67
return clientRef.ServerStatus()
69
- }, cancel)
68
+ })
69
70
// 2) libp2p host
71
client, err := relaydns.NewClient(ctx, relaydns.ClientConfig{
@@ -90,7 +89,11 @@ func runClient(cmd *cobra.Command, args []string) error {
89
sig := make(chan os.Signal, 1)
90
signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
91
<-sig
93
- log.Info().Msg("[client] shutting down")
94
- time.Sleep(200 * time.Millisecond)
92
+ log.Info().Msg("[client] shutting down client...")
93
+ if err := srv.Shutdown(ctx); err != nil {
94
+ log.Error().Err(err).Msg("[client] server forced to shutdown")
95
+ }
96
+ log.Info().Msg("[client] http server stopped")
97
+
98
return nil
99
}
cmd/example_http_client/view.go
+11
-6
@@ -1,7 +1,6 @@
1
package main
2
3
import (
4
- "context"
4
"html/template"
5
"net"
6
"net/http"
@@ -12,7 +11,7 @@ import (
11
12
// serveClientHTTP serves the simple local backend UI and health endpoint.
13
// getStatus should return a short string like "Connected" or "Connecting...".
15
-func serveClientHTTP(ctx context.Context, ln net.Listener, name string, getStatus func() string, cancel context.CancelFunc) {
14
+func serveClientHTTP(ln net.Listener, name string, getStatus func() string) *http.Server {
15
mux := http.NewServeMux()
16
mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
17
status := getStatus()
@@ -40,11 +39,17 @@ func serveClientHTTP(ctx context.Context, ln net.Listener, name string, getStatu
39
_, _ = w.Write([]byte("ok"))
40
})
41
42
+ srv := &http.Server{Handler: mux, ReadHeaderTimeout: 5 * time.Second, IdleTimeout: 60 * time.Second}
43
+
44
log.Info().Msgf("[client] local backend http %s", ln.Addr().String())
44
- if err := http.Serve(ln, mux); err != nil {
45
- log.Error().Err(err).Msg("[client] http backend error")
46
- cancel()
47
- }
45
+
46
+ // Serve in background
47
+ go func() {
48
+ if err := srv.Serve(ln); err != nil && err != http.ErrServerClosed {
49
+ log.Error().Err(err).Msg(" http backend error")
50
+ }
51
+ }()
52
+ return srv
53
}
54
55
var clientPage = template.Must(template.New("index").Parse(`<!DOCTYPE html>
relaydns/client.go
+9
-1
@@ -162,7 +162,15 @@ func (b *RelayClient) Close() error {
162
b.stop()
163
}
164
b.wg.Wait()
165
- // leaving topic is optional; libp2p will clean up on host close
165
+ if b.t != nil {
166
+ _ = b.t.Close()
167
+ }
168
+ if b.h != nil {
169
+ if b.protoID != "" {
170
+ b.h.RemoveStreamHandler(b.protoID)
171
+ }
172
+ return b.h.Close()
173
+ }
174
return nil
175
}
176