fix: Proper server shutdown process in goroutine wg countings and improve HTTP server handling

cognitive-glitch committed Oct 22, 2025 at 00:10 UTC c0874d801c985138829c92b37b439e6d9249391d
5 files changed +45 -12
.gitignore
+1
@@ -7,6 +7,7 @@
7 *.dll
8 *.so
9 *.dylib
10 +bin/
11
12 # Test binary, built with `go test -c`
13 *.test
cmd/server/main.go
+23 -2
@@ -53,7 +53,7 @@ func runServer(cmd *cobra.Command, args []string) error {
53 }
54
55 // Admin UI + per-peer HTTP proxy served here
56 - go serveHTTP(ctx, fmt.Sprintf(":%d", flagHttpPort), d, h, cancel)
56 + httpServer := serveHTTP(ctx, fmt.Sprintf(":%d", flagHttpPort), d, h, cancel)
57
58 // Optional raw TCP ingress (e.g., SSH)
59 if flagTcpPort > 0 {
@@ -64,8 +64,29 @@ func runServer(cmd *cobra.Command, args []string) error {
64 sig := make(chan os.Signal, 1)
65 signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
66 <-sig
67 + log.Info().Msg("[server] shutting down...")
68 +
69 + // Cancel context to stop all goroutines
70 cancel()
68 - time.Sleep(300 * time.Millisecond)
71 +
72 + // Shutdown HTTP server with timeout
73 + shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
74 + defer shutdownCancel()
75 + if err := httpServer.Shutdown(shutdownCtx); err != nil {
76 + log.Error().Err(err).Msg("[server] http server shutdown error")
77 + }
78 +
79 + // Close relay server (waits for background goroutines)
80 + if err := d.Close(); err != nil {
81 + log.Warn().Err(err).Msg("[server] relay server close error")
82 + }
83 +
84 + // Close libp2p host
85 + if err := h.Close(); err != nil {
86 + log.Warn().Err(err).Msg("[server] libp2p host close error")
87 + }
88 +
89 + log.Info().Msg("[server] shutdown complete")
90 return nil
91 }
92
cmd/server/view.go
+15 -9
@@ -14,11 +14,8 @@ import (
14 "github.com/rs/zerolog/log"
15 )
16
17 -// serveHTTP builds the HTTP mux.
18 -func serveHTTP(ctx context.Context, addr string, d *relaydns.RelayServer, h host.Host, cancel context.CancelFunc) {
19 - if addr == "" {
20 - return
21 - }
17 +// serveHTTP builds the HTTP mux and returns the server.
18 +func serveHTTP(ctx context.Context, addr string, d *relaydns.RelayServer, h host.Host, cancel context.CancelFunc) *http.Server {
19 mux := http.NewServeMux()
20
21 // Index page
@@ -105,11 +102,20 @@ func serveHTTP(ctx context.Context, addr string, d *relaydns.RelayServer, h host
102 _ = json.NewEncoder(w).Encode(resp)
103 })
104
108 - log.Info().Msgf("[server] http: %s", addr)
109 - if err := http.ListenAndServe(addr, mux); err != nil {
110 - log.Error().Err(err).Msg("[server] http error")
111 - cancel()
105 + srv := &http.Server{
106 + Addr: addr,
107 + Handler: mux,
108 }
109 +
110 + go func() {
111 + log.Info().Msgf("[server] http: %s", addr)
112 + if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
113 + log.Error().Err(err).Msg("[server] http error")
114 + cancel()
115 + }
116 + }()
117 +
118 + return srv
119 }
120
121 var adminIndexTmpl = template.Must(template.New("admin-index").Parse(`<!doctype html>
relaydns/client.go
-1
@@ -134,7 +134,6 @@ func (b *RelayClient) Start(ctx context.Context) error {
134 // 4) advertiser loop
135 advCtx, cancel := context.WithCancel(ctx)
136 b.stop = cancel
137 - b.wg.Add(1)
137 b.startAdvertiser(advCtx)
138
139 // 5) background: periodically re-fetch bootstraps from server and reconnect
relaydns/server.go
+6
@@ -33,6 +33,8 @@ type RelayServer struct {
33 store map[string]HostEntry
34 ttl time.Duration
35 deadTTL time.Duration
36 +
37 + wg sync.WaitGroup
38 }
39
40 func NewRelayServer(ctx context.Context, h host.Host, protocol, topic string) (*RelayServer, error) {
@@ -58,6 +60,7 @@ func NewRelayServer(ctx context.Context, h host.Host, protocol, topic string) (*
60 ttl: 15 * time.Second,
61 deadTTL: 10 * time.Minute,
62 }
63 + d.wg.Add(2)
64 go d.collect()
65 go d.gc()
66 return d, nil
@@ -65,10 +68,12 @@ func NewRelayServer(ctx context.Context, h host.Host, protocol, topic string) (*
68
69 func (s *RelayServer) Close() error {
70 s.sub.Cancel()
71 + s.wg.Wait()
72 return nil
73 }
74
75 func (s *RelayServer) collect() {
76 + defer s.wg.Done()
77 for {
78 msg, err := s.sub.Next(s.ctx)
79 if err != nil {
@@ -111,6 +116,7 @@ func (s *RelayServer) collect() {
116 }
117
118 func (s *RelayServer) gc() {
119 + defer s.wg.Done()
120 t := time.NewTicker(5 * time.Second)
121 defer t.Stop()
122 for {