From 0a5443f6cf4ff6a1000aaebfbc5fcb787b935ecf Mon Sep 17 00:00:00 2001
From: xfy <i@rua.plus>
Date: Mon, 8 Jun 2026 17:47:37 +0800
Subject: [PATCH] fix(sticky): guard against double Stop, nil fallback, and
 multiple Start calls

- Add sync.Once to prevent double close of stopCh in Stop()
- Add nil fallback guard in NewStickySession (defaults to RoundRobin)
- Add atomic.Bool to make Start() idempotent
- Add tests for double Stop() and nil fallback scenarios
---
 internal/loadbalance/sticky.go      | 16 +++++++++++++-
 internal/loadbalance/sticky_test.go | 33 +++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/internal/loadbalance/sticky.go b/internal/loadbalance/sticky.go
index cc0eccb..2a38a3a 100644
--- a/internal/loadbalance/sticky.go
+++ b/internal/loadbalance/sticky.go
@@ -5,6 +5,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/valyala/fasthttp"
@@ -29,10 +30,15 @@ type StickySession struct {
 	shards   []*stickyShard
 	stopCh   chan struct{}
 	wg       sync.WaitGroup
+	started  atomic.Bool
+	stopOnce sync.Once
 }
 
 // NewStickySession 创建一个新的会话粘性负载均衡器。
 func NewStickySession(config StickyConfig, fallback Balancer) *StickySession {
+	if fallback == nil {
+		fallback = NewRoundRobin()
+	}
 	shards := make([]*stickyShard, stickyShardCount)
 	for i := range shards {
 		shards[i] = &stickyShard{
@@ -50,13 +56,21 @@ func NewStickySession(config StickyConfig, fallback Balancer) *StickySession {
 
 // Start 启动后台清理 goroutine。
 func (s *StickySession) Start() {
+	if s.started.Swap(true) {
+		return
+	}
 	s.wg.Add(1)
 	go s.cleanupLoop()
 }
 
 // Stop 停止后台清理 goroutine。
 func (s *StickySession) Stop() {
-	close(s.stopCh)
+	if !s.started.Swap(false) {
+		return
+	}
+	s.stopOnce.Do(func() {
+		close(s.stopCh)
+	})
 	s.wg.Wait()
 }
 
diff --git a/internal/loadbalance/sticky_test.go b/internal/loadbalance/sticky_test.go
index 2cce1e5..acd3451 100644
--- a/internal/loadbalance/sticky_test.go
+++ b/internal/loadbalance/sticky_test.go
@@ -278,6 +278,39 @@ func TestStickySession_ExpiredCookie(t *testing.T) {
 	}
 }
 
+// TestStickySession_DoubleStop 测试重复调用 Stop() 不会 panic。
+func TestStickySession_DoubleStop(t *testing.T) {
+	t.Parallel()
+	fallback := NewRoundRobin()
+	config := DefaultStickyConfig()
+
+	sticky := NewStickySession(config, fallback)
+	sticky.Start()
+	sticky.Stop()
+	sticky.Stop() // Should not panic
+}
+
+// TestStickySession_NilFallback 测试 nil fallback 会使用默认 RoundRobin。
+func TestStickySession_NilFallback(t *testing.T) {
+	t.Parallel()
+	config := DefaultStickyConfig()
+	config.Enabled = true
+
+	sticky := NewStickySession(config, nil)
+	sticky.Start()
+	defer sticky.Stop()
+
+	targets := []*Target{
+		createHealthyTarget("http://backend1:8080", true),
+	}
+
+	ctx := &fasthttp.RequestCtx{}
+	selected := sticky.Select(ctx, targets)
+	if selected == nil {
+		t.Fatal("expected a target with default fallback")
+	}
+}
+
 // TestStickySession_SelectExcluding 测试排除选择委托给 fallback。
 func TestStickySession_SelectExcluding(t *testing.T) {
 	t.Parallel()