Improve resource cleanup and error logging
Adds explicit error logging for resource closures (streams, connections, response bodies) throughout relaydns and example clients. Refactors shutdown sequences in example_chat and example_http_client to ensure proper context cancellation and orderly shutdown. Also fixes minor error handling in websocket closure.
cognitive committed
Oct 21, 2025 at 22:51 UTC
144367abf0ad951f77928ccb5e9b9cd4a2032546
6 files changed
+83
-32
cmd/example_chat/main.go
+10
-10
@@ -41,6 +41,7 @@ func main() {
41
42
func runChat(cmd *cobra.Command, args []string) error {
43
ctx, cancel := context.WithCancel(context.Background())
44
+ defer cancel() // Ensure context is cancelled on all exit paths
45
46
// 1) start local chat HTTP backend
47
ln, err := net.Listen("tcp", fmt.Sprintf(":%d", flagPort))
@@ -70,27 +71,26 @@ func runChat(cmd *cobra.Command, args []string) error {
71
signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
72
<-sig
73
log.Info().Msg("[chat] shutting down...")
73
-
74
+
75
// 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)
76
+ // Note: defer cancel() at function start stops client advertising/refresh loops
77
+
78
+ // 1. 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
}
82
-
83
- // 3. Shutdown HTTP server with a fresh context (with timeout)
82
+
83
+ // 2. 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
89
+
90
+ // 3. Close all websocket connections and wait for handlers to finish
91
hub.closeAll()
92
hub.wait()
93
-
93
+
94
log.Info().Msg("[chat] shutdown complete")
95
return nil
96
}
cmd/example_chat/view.go
+1
-1
@@ -100,7 +100,7 @@ func handleWS(w http.ResponseWriter, r *http.Request, h *hub) {
100
if leftUser != "" {
101
h.broadcast(message{TS: time.Now().UTC(), User: leftUser, Event: "left"})
102
}
103
- conn.Close(websocket.StatusNormalClosure, "")
103
+ _ = conn.Close(websocket.StatusNormalClosure, "")
104
cancelConn()
105
h.wg.Done()
106
}()
cmd/example_http_client/main.go
+18
-6
@@ -8,6 +8,7 @@ import (
8
"os"
9
"os/signal"
10
"syscall"
11
+ "time"
12
13
"github.com/gosuda/relaydns/relaydns"
14
"github.com/rs/zerolog/log"
@@ -41,7 +42,7 @@ func main() {
42
43
func runClient(cmd *cobra.Command, args []string) error {
44
ctx, cancel := context.WithCancel(context.Background())
44
- defer cancel()
45
+ defer cancel() // Ensure context is cancelled on all exit paths
46
47
if flagName == "" {
48
hn, err := os.Hostname()
@@ -83,17 +84,28 @@ func runClient(cmd *cobra.Command, args []string) error {
84
return fmt.Errorf("start client: %w", err)
85
}
86
clientRef = client
86
- defer client.Close()
87
88
// wait for termination
89
sig := make(chan os.Signal, 1)
90
signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
91
<-sig
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")
92
+ log.Info().Msg("[client] shutting down...")
93
+
94
+ // Shutdown sequence:
95
+ // Note: defer cancel() at function start stops client advertising/refresh loops
96
+
97
+ // 1. Close client (waits for goroutines, closes libp2p host)
98
+ if err := client.Close(); err != nil {
99
+ log.Warn().Err(err).Msg("[client] client close error")
100
+ }
101
+
102
+ // 2. Shutdown HTTP server with a fresh context (with timeout)
103
+ shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
104
+ defer shutdownCancel()
105
+ if err := srv.Shutdown(shutdownCtx); err != nil {
106
+ log.Error().Err(err).Msg("[client] http server shutdown error")
107
}
96
- log.Info().Msg("[client] http server stopped")
108
109
+ log.Info().Msg("[client] shutdown complete")
110
return nil
111
}
relaydns/client.go
+15
-5
@@ -207,16 +207,26 @@ func (b *RelayClient) setupStreamHandler() error {
207
b.h.SetStreamHandler(b.protoID, b.cfg.Handler)
208
case b.cfg.TargetTCP != "":
209
b.h.SetStreamHandler(b.protoID, func(s network.Stream) {
210
- defer s.Close()
210
+ defer func() {
211
+ if err := s.Close(); err != nil {
212
+ log.Debug().Err(err).Msg("relaydns: stream close")
213
+ }
214
+ }()
215
c, err := net.Dial("tcp", b.cfg.TargetTCP)
216
if err != nil {
217
log.Error().Err(err).Msgf("relaydns: dial %s", b.cfg.TargetTCP)
218
return
219
}
216
- defer c.Close()
217
- // raw byte pipe
218
- go io.Copy(c, s)
219
- io.Copy(s, c)
220
+ defer func() {
221
+ if err := c.Close(); err != nil {
222
+ log.Debug().Err(err).Msg("relaydns: conn close")
223
+ }
224
+ }()
225
+ // raw byte pipe (bidirectional)
226
+ go func() {
227
+ _, _ = io.Copy(c, s)
228
+ }()
229
+ _, _ = io.Copy(s, c)
230
})
231
default:
232
return fmt.Errorf("relaydns: either Handler or TargetTCP must be set")
relaydns/server.go
+33
-9
@@ -193,7 +193,11 @@ func (d *RelayServer) ProxyHTTP(w http.ResponseWriter, r *http.Request, peerID,
193
http.Error(w, "open stream failed", http.StatusBadGateway)
194
return
195
}
196
- defer s.Close()
196
+ defer func() {
197
+ if err := s.Close(); err != nil {
198
+ log.Debug().Err(err).Msg("stream close")
199
+ }
200
+ }()
201
outReq := r.Clone(d.ctx)
202
outReq.URL = &url.URL{Path: pathSuffix, RawQuery: r.URL.RawQuery}
203
outReq.RequestURI = ""
@@ -209,7 +213,11 @@ func (d *RelayServer) ProxyHTTP(w http.ResponseWriter, r *http.Request, peerID,
213
http.Error(w, "bad upstream response", http.StatusBadGateway)
214
return
215
}
212
- defer resp.Body.Close()
216
+ defer func() {
217
+ if err := resp.Body.Close(); err != nil {
218
+ log.Debug().Err(err).Msg("response body close")
219
+ }
220
+ }()
221
222
// Handle WebSocket Upgrade: write 101 response and then raw-tunnel bytes
223
if resp.StatusCode == http.StatusSwitchingProtocols && strings.Contains(strings.ToLower(resp.Header.Get("Upgrade")), "websocket") {
@@ -219,7 +227,11 @@ func (d *RelayServer) ProxyHTTP(w http.ResponseWriter, r *http.Request, peerID,
227
log.Error().Err(err).Msg("hijack client conn")
228
return
229
}
222
- defer clientConn.Close()
230
+ defer func() {
231
+ if err := clientConn.Close(); err != nil {
232
+ log.Debug().Err(err).Msg("client conn close")
233
+ }
234
+ }()
235
// Write upstream 101 response (headers)
236
if err := resp.Write(clientBuf); err == nil {
237
_ = clientBuf.Flush()
@@ -231,9 +243,11 @@ func (d *RelayServer) ProxyHTTP(w http.ResponseWriter, r *http.Request, peerID,
243
_, _ = clientConn.Write(tmp)
244
}
245
}
234
- // Raw byte tunnel between client and upstream stream
235
- go io.Copy(s, clientConn)
236
- io.Copy(clientConn, s)
246
+ // Raw byte tunnel between client and upstream stream (bidirectional)
247
+ go func() {
248
+ _, _ = io.Copy(s, clientConn)
249
+ }()
250
+ _, _ = io.Copy(clientConn, s)
251
return
252
}
253
}
@@ -267,7 +281,11 @@ func (d *RelayServer) ProxyHTTP(w http.ResponseWriter, r *http.Request, peerID,
281
// ProxyTCP opens a libp2p stream to peerID using the Director protocol and
282
// pipes raw bytes between the accepted TCP connection and the libp2p stream.
283
func (d *RelayServer) ProxyTCP(c net.Conn, peerID string) error {
270
- defer c.Close()
284
+ defer func() {
285
+ if err := c.Close(); err != nil {
286
+ log.Debug().Err(err).Msg("tcp conn close")
287
+ }
288
+ }()
289
d.storeMu.Lock()
290
entry, ok := d.store[peerID]
291
d.storeMu.Unlock()
@@ -281,9 +299,15 @@ func (d *RelayServer) ProxyTCP(c net.Conn, peerID string) error {
299
if err != nil {
300
return err
301
}
284
- defer s.Close()
302
+ defer func() {
303
+ if err := s.Close(); err != nil {
304
+ log.Debug().Err(err).Msg("stream close")
305
+ }
306
+ }()
307
// bidirectional copy
286
- go io.Copy(s, c)
308
+ go func() {
309
+ _, _ = io.Copy(s, c)
310
+ }()
311
_, _ = io.Copy(c, s)
312
return nil
313
}
relaydns/utils.go
+6
-1
@@ -10,6 +10,7 @@ import (
10
"time"
11
12
"github.com/libp2p/go-libp2p/core/host"
13
+ "github.com/rs/zerolog/log"
14
)
15
16
func AddrToTarget(listen string) string {
@@ -73,7 +74,11 @@ func fetchMultiaddrsFromHosts(base string, timeout time.Duration) ([]string, err
74
if err != nil {
75
return nil, err
76
}
76
- defer resp.Body.Close()
77
+ defer func() {
78
+ if err := resp.Body.Close(); err != nil {
79
+ log.Debug().Err(err).Msg("response body close")
80
+ }
81
+ }()
82
83
var payload Hosts
84
if err := json.NewDecoder(resp.Body).Decode(&payload); err != nil {