From b985e5138cf55c0c830af3206a31c59d1cb72393 Mon Sep 17 00:00:00 2001
From: ginuerzh <ginuerzh@gmail.com>
Date: Sun, 21 Jun 2026 21:07:37 +0800
Subject: [PATCH] fix: use atomic int for multi-worker exit code and clean up
 cancellation

- Replace bare int with atomic.Int32 to fix data race on ret
- Replace wg.Add/go/defer wg.Done pattern with wg.Go (Go 1.22+)
- Pass context by value instead of pointer
- Add ctx.Err() guard to suppress fatal on cancellation errors
- Add ProcessState nil check before accessing ExitCode
- Move cancel() after wg.Wait() to avoid race where a successful
  worker's deferred cancel kills siblings and triggers log.Fatal
---
 cmd/gost/main.go | 37 +++++++++++++++++++++----------------
 1 file changed, 21 insertions(+), 16 deletions(-)

diff --git a/cmd/gost/main.go b/cmd/gost/main.go
index b3141e7..7efdcde 100644
--- a/cmd/gost/main.go
+++ b/cmd/gost/main.go
@@ -11,6 +11,7 @@ import (
 	"runtime"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/go-gost/core/logger"
@@ -48,39 +49,43 @@ func init() {
 	if strings.Contains(args, " -- ") {
 		var (
 			wg  sync.WaitGroup
-			ret int
+			ret atomic.Int32
 		)
 
-		ctx, cancel := context.WithCancel(context.Background())
-		defer cancel()
+		wargsList := strings.Split(" "+args+" ", " -- ")
 
-		for wid, wargs := range strings.Split(" "+args+" ", " -- ") {
-			wg.Add(1)
-			go func(wid int, wargs string) {
-				defer wg.Done()
-				defer cancel()
-				worker(wid, strings.Split(wargs, "  "), &ctx, &ret)
-			}(wid, strings.TrimSpace(wargs))
+		ctx, cancel := context.WithCancel(context.Background())
+
+		for wid, wargs := range wargsList {
+			wg.Go(func() {
+				worker(wid, strings.Split(strings.TrimSpace(wargs), "  "), ctx, &ret)
+			})
 		}
 
 		wg.Wait()
+		cancel()
 
-		os.Exit(ret)
+		os.Exit(int(ret.Load()))
 	}
 }
 
-func worker(id int, args []string, ctx *context.Context, ret *int) {
-	cmd := exec.CommandContext(*ctx, os.Args[0], args...)
+func worker(id int, args []string, ctx context.Context, ret *atomic.Int32) {
+	cmd := exec.CommandContext(ctx, os.Args[0], args...)
 
 	cmd.Stdout = os.Stdout
 	cmd.Stderr = os.Stderr
 	cmd.Env = append(os.Environ(), fmt.Sprintf("_GOST_ID=%d", id))
 
 	if err := cmd.Run(); err != nil {
-		log.Fatal(err)
+		// Context cancellation is expected when one worker exits early.
+		// Only log fatal on other errors.
+		if ctx.Err() == nil {
+			log.Fatal(err)
+		}
+		return
 	}
-	if cmd.ProcessState.Exited() {
-		*ret = cmd.ProcessState.ExitCode()
+	if cmd.ProcessState != nil && cmd.ProcessState.Exited() {
+		ret.Store(int32(cmd.ProcessState.ExitCode()))
 	}
 }