Cancel if signal received

cluster/join.go

@@ -18,6 +18,9 @@ var (
 	// ErrJoinFailed is returned when a node fails to join a cluster
 	ErrJoinFailed = errors.New("failed to join cluster")
 
+	// ErrJoinCanceled is returned when a join operation is canceled
+	ErrJoinCanceled = errors.New("join operation canceled")
+
 	// ErrNotifyFailed is returned when a node fails to notify another node
 	ErrNotifyFailed = errors.New("failed to notify node")
 )
@@ -60,7 +63,7 @@ func (j *Joiner) Do(ctx context.Context, targetAddrs []string, id, addr string,
 		for _, ta := range targetAddrs {
 			select {
 			case <-ctx.Done():
-				return "", ctx.Err() // Return the cancellation error
+				return "", ErrJoinCanceled
 			default:
 				joinee, err = j.join(ta, id, addr, suf)
 				if err == nil {
@@ -74,7 +77,7 @@ func (j *Joiner) Do(ctx context.Context, targetAddrs []string, id, addr string,
 			j.logger.Printf("failed to join cluster at %s, sleeping %s before retry", targetAddrs, j.attemptInterval)
 			select {
 			case <-ctx.Done():
-				return "", ctx.Err() // Check cancellation again before sleeping
+				return "", ErrJoinCanceled
 			case <-time.After(j.attemptInterval):
 				continue // Proceed with the next attempt
 			}

cmd/rqlited/main.go

@@ -58,7 +58,7 @@ func init() {
 func main() {
 	// Handle signals first, so signal handling is established before anything else.
 	sigCh := HandleSignals(syscall.SIGINT, syscall.SIGTERM, os.Interrupt)
-	mainCtx, mainCancel := context.WithCancel(context.Background())
+	mainCtx, _ := CreateContext(sigCh)
 
 	cfg, err := ParseFlags(name, desc, &BuildInfo{
 		Version:       cmd.Version,
@@ -188,18 +188,16 @@ func main() {
 	log.Printf("connect using the command-line tool via 'rqlite -H %s -p %s'", h, p)
 
 	// Start any requested auto-backups
-	backupSrvStx, backupSrvCancel := context.WithCancel(mainCtx)
-	backupSrv, err := startAutoBackups(backupSrvStx, cfg, str)
+	backupSrv, err := startAutoBackups(mainCtx, cfg, str)
 	if err != nil {
 		log.Fatalf("failed to start auto-backups: %s", err.Error())
 	}
 	if backupSrv != nil {
 		httpServ.RegisterStatus("auto_backups", backupSrv)
 	}
 
-	// Block until signalled.
-	<-sigCh
-	mainCancel()
+	// Block until done.
+	<-mainCtx.Done()
 
 	// Stop the HTTP server first, so clients get notification as soon as
 	// possible that the node is going away.
@@ -224,7 +222,6 @@ func main() {
 		str.Stepdown(true)
 	}
 
-	backupSrvCancel()
 	if err := str.Close(true); err != nil {
 		log.Printf("failed to close store: %s", err.Error())
 	}

cmd/rqlited/signals.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"log"
 	"os"
 	"os/signal"
@@ -24,3 +25,14 @@ func HandleSignals(sigs ...os.Signal) <-chan os.Signal {
 	}()
 	return ch
 }
+
+// CreateContext creates a context which is canceled if signals are received
+// on the given channel.
+func CreateContext(ch <-chan os.Signal) (context.Context, context.CancelFunc) {
+	ctx, cancel := context.WithCancel(context.Background())
+	go func() {
+		<-ch
+		cancel()
+	}()
+	return ctx, cancel
+}