Skip to content

Commit f2471ee

Browse files
0pcomcharmcrush
andcommitted
Improve signal handling for graceful shutdown on Ctrl+C
Address issues where DMSG client utilities would not respond to Ctrl+C: - Use signal-aware context for dmsgC.Serve() instead of context.Background() in internal/cli, pkg/direct, and dmsgpty-host - Add context cancellation check in dmsg-socks5 accept loop to properly exit on signal - Implement graceful HTTP server shutdown with 5-second timeout in dmsghttp - Ensure dmsgC.Close() is called via defer in dmsgpty-host for proper cleanup - Add WaitGroup wait in direct client stop function for clean goroutine shutdown - Document temporary kill.go workaround for verification testing These changes ensure client utilities properly respond to interrupt signals without requiring multiple Ctrl+C presses. 💘 Generated with Crush Co-Authored-By: Crush <[email protected]>
1 parent 215a0c3 commit f2471ee

File tree

6 files changed

+43
-8
lines changed

6 files changed

+43
-8
lines changed

cmd/dmsg-socks5/commands/dmsg-socks5.go

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -157,11 +157,25 @@ var serveCmd = &cobra.Command{
157157
dlog.Printf("Error closing listener: %v", err)
158158
}
159159
}()
160+
161+
go func() {
162+
<-ctx.Done()
163+
if err := dmsgL.Close(); err != nil {
164+
dlog.WithError(err).Debug("Error closing listener on context cancellation")
165+
}
166+
}()
167+
160168
for {
161169
respConn, err := dmsgL.Accept()
162170
if err != nil {
163-
dlog.Errorf("Error accepting initiator: %v", err)
164-
continue
171+
select {
172+
case <-ctx.Done():
173+
dlog.Info("Shutting down SOCKS5 server...")
174+
return
175+
default:
176+
dlog.Errorf("Error accepting initiator: %v", err)
177+
continue
178+
}
165179
}
166180
dlog.Infof("Accepted connection from: %s", respConn.RemoteAddr())
167181

cmd/dmsg/commands/kill.go

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Package commands cmd/.../commands/kill.go
1+
// Package commands cmd/dmsg/commands/kill.go
22
package commands
33

44
import (
@@ -8,7 +8,11 @@ import (
88
)
99

1010
func init() {
11-
//the application must stop on ctrl+c
11+
// TEMPORARY WORKAROUND: Force exit on Ctrl+C after 3 attempts
12+
// This can be removed once the proper signal handling fixes are verified:
13+
// - dmsgC.Serve() now uses signal-aware context (not context.Background())
14+
// - Accept loops now check for context cancellation
15+
// - HTTP servers now shutdown gracefully
1216
c := make(chan os.Signal, 1)
1317
signal.Notify(c, os.Interrupt, syscall.SIGTERM)
1418
go func() {

cmd/dmsghttp/commands/dmsghttp.go

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -165,6 +165,16 @@ func server() {
165165
WriteTimeout: 10 * time.Second,
166166
}
167167

168+
// Gracefully shutdown HTTP server on context cancellation
169+
go func() {
170+
<-ctx.Done()
171+
shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
172+
defer cancel()
173+
if err := serve.Shutdown(shutdownCtx); err != nil {
174+
dlog.WithError(err).Warn("Server shutdown error")
175+
}
176+
}()
177+
168178
// Start serving
169179
go func() {
170180
dlog.WithField("dmsg_addr", lis.Addr().String()).Debug("Serving...\n")

cmd/dmsgpty-host/commands/root.go

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,13 @@ var RootCmd = &cobra.Command{
104104
dmsgC := dmsg.NewClient(pk, sk, disc.NewHTTP(conf.DmsgDisc, &http.Client{}, log), &dmsg.Config{
105105
MinSessions: conf.DmsgSessions,
106106
})
107-
go dmsgC.Serve(context.Background())
107+
defer func() {
108+
if err := dmsgC.Close(); err != nil {
109+
log.WithError(err).Warn("Failed to close dmsg client")
110+
}
111+
}()
112+
113+
go dmsgC.Serve(ctx)
108114
select {
109115
case <-ctx.Done():
110116
return fmt.Errorf("failed to wait dmsg client to be ready: %w", ctx.Err())

internal/cli/cli.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -123,8 +123,8 @@ func StartDmsg(ctx context.Context, dlog *logging.Logger, pk cipher.PubKey, sk c
123123
dmsgC = dmsg.NewClient(pk, sk, disc.NewHTTP(dmsgDisc, httpClient, dlog), &dmsg.Config{MinSessions: dmsgSessions})
124124
dlog.Debug("Created dmsg client.")
125125

126-
go dmsgC.Serve(context.Background())
127-
dlog.Debug("dmsgclient.Serve(context.Background())")
126+
go dmsgC.Serve(ctx)
127+
dlog.Debug("dmsgclient.Serve(ctx)")
128128

129129
stop = func() {
130130
err := dmsgC.Close()

pkg/direct/direct.go

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,12 +23,13 @@ func StartDmsg(ctx context.Context, log *logging.Logger, pk cipher.PubKey, sk ci
2323
wg.Add(1)
2424
go func() {
2525
defer wg.Done()
26-
dmsgDC.Serve(context.Background())
26+
dmsgDC.Serve(ctx)
2727
}()
2828

2929
stop = func() {
3030
err := dmsgDC.Close()
3131
log.WithError(err).Debug("Disconnected from dmsg network.\n")
32+
wg.Wait()
3233
}
3334

3435
log.WithField("public_key", pk.String()).Debug("Connecting to dmsg network...\n")

0 commit comments

Comments
 (0)