xfy
0a5443f6cf
- 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
340 lines
9.3 KiB
Go
340 lines
9.3 KiB
Go
package loadbalance
|
||
|
||
import (
|
||
"sync"
|
||
"testing"
|
||
"time"
|
||
|
||
"github.com/valyala/fasthttp"
|
||
)
|
||
|
||
// TestStickySession_BasicRoute 测试基本的会话粘性路由。
|
||
// 第一次请求应设置 cookie，第二次携带相同 cookie 应路由到同一目标。
|
||
func TestStickySession_BasicRoute(t *testing.T) {
|
||
t.Parallel()
|
||
t.Run("首次请求设置cookie并路由", func(_ *testing.T) {
|
||
config := DefaultStickyConfig()
|
||
config.Enabled = true
|
||
fallback := NewRoundRobin()
|
||
sticky := NewStickySession(config, fallback)
|
||
sticky.Start()
|
||
defer sticky.Stop()
|
||
|
||
targets := []*Target{
|
||
createHealthyTarget("http://backend1:8080", true),
|
||
createHealthyTarget("http://backend2:8080", true),
|
||
}
|
||
|
||
ctx := &fasthttp.RequestCtx{}
|
||
got := sticky.Select(ctx, targets)
|
||
if got == nil {
|
||
t.Fatal("Select() = nil, want non-nil")
|
||
}
|
||
|
||
// 验证设置了 cookie
|
||
cookieValue := ctx.Response.Header.PeekCookie(config.Name)
|
||
if len(cookieValue) == 0 {
|
||
t.Error("首次请求未设置 cookie")
|
||
}
|
||
})
|
||
|
||
t.Run("相同cookie路由到同一目标", func(_ *testing.T) {
|
||
config := DefaultStickyConfig()
|
||
config.Enabled = true
|
||
fallback := NewRoundRobin()
|
||
sticky := NewStickySession(config, fallback)
|
||
sticky.Start()
|
||
defer sticky.Stop()
|
||
|
||
targets := []*Target{
|
||
createHealthyTarget("http://backend1:8080", true),
|
||
createHealthyTarget("http://backend2:8080", true),
|
||
}
|
||
|
||
// 第一次请求
|
||
ctx1 := &fasthttp.RequestCtx{}
|
||
got1 := sticky.Select(ctx1, targets)
|
||
if got1 == nil {
|
||
t.Fatal("第一次 Select() = nil")
|
||
}
|
||
|
||
// 提取 cookie
|
||
cookie := &fasthttp.Cookie{}
|
||
cookie.SetKey(config.Name)
|
||
if err := cookie.ParseBytes(ctx1.Response.Header.PeekCookie(config.Name)); err != nil {
|
||
t.Fatalf("解析 cookie 失败: %v", err)
|
||
}
|
||
|
||
// 第二次请求携带相同 cookie
|
||
ctx2 := &fasthttp.RequestCtx{}
|
||
ctx2.Request.Header.SetCookie(config.Name, string(cookie.Value()))
|
||
got2 := sticky.Select(ctx2, targets)
|
||
if got2 == nil {
|
||
t.Fatal("第二次 Select() = nil")
|
||
}
|
||
|
||
if got2.URL != got1.URL {
|
||
t.Errorf("相同 cookie 路由到不同目标: 第一次=%q, 第二次=%q", got1.URL, got2.URL)
|
||
}
|
||
})
|
||
|
||
t.Run("禁用时不设置cookie", func(_ *testing.T) {
|
||
config := DefaultStickyConfig()
|
||
config.Enabled = false
|
||
fallback := NewRoundRobin()
|
||
sticky := NewStickySession(config, fallback)
|
||
sticky.Start()
|
||
defer sticky.Stop()
|
||
|
||
targets := []*Target{
|
||
createHealthyTarget("http://backend1:8080", true),
|
||
}
|
||
|
||
ctx := &fasthttp.RequestCtx{}
|
||
got := sticky.Select(ctx, targets)
|
||
if got == nil {
|
||
t.Fatal("Select() = nil")
|
||
}
|
||
|
||
cookieValue := ctx.Response.Header.PeekCookie(config.Name)
|
||
if len(cookieValue) > 0 {
|
||
t.Error("禁用时不应设置 cookie")
|
||
}
|
||
})
|
||
}
|
||
|
||
// TestStickySession_TargetUnavailable 测试目标不可用时回退到 fallback。
|
||
func TestStickySession_TargetUnavailable(t *testing.T) {
|
||
t.Parallel()
|
||
t.Run("目标不健康时回退", func(_ *testing.T) {
|
||
config := DefaultStickyConfig()
|
||
config.Enabled = true
|
||
fallback := NewRoundRobin()
|
||
sticky := NewStickySession(config, fallback)
|
||
sticky.Start()
|
||
defer sticky.Stop()
|
||
|
||
targets := []*Target{
|
||
createHealthyTarget("http://backend1:8080", true),
|
||
createHealthyTarget("http://backend2:8080", true),
|
||
}
|
||
|
||
// 第一次请求，记录会话
|
||
ctx1 := &fasthttp.RequestCtx{}
|
||
got1 := sticky.Select(ctx1, targets)
|
||
if got1 == nil {
|
||
t.Fatal("第一次 Select() = nil")
|
||
}
|
||
|
||
// 提取 cookie
|
||
cookie := &fasthttp.Cookie{}
|
||
cookie.SetKey(config.Name)
|
||
if err := cookie.ParseBytes(ctx1.Response.Header.PeekCookie(config.Name)); err != nil {
|
||
t.Fatalf("解析 cookie 失败: %v", err)
|
||
}
|
||
|
||
// 使之前选中的目标不健康
|
||
for _, target := range targets {
|
||
if target.URL == got1.URL {
|
||
target.Healthy.Store(false)
|
||
break
|
||
}
|
||
}
|
||
|
||
// 第二次请求，应回退到其他目标
|
||
ctx2 := &fasthttp.RequestCtx{}
|
||
ctx2.Request.Header.SetCookie(config.Name, string(cookie.Value()))
|
||
got2 := sticky.Select(ctx2, targets)
|
||
if got2 == nil {
|
||
t.Fatal("第二次 Select() = nil")
|
||
}
|
||
|
||
if got2.URL == got1.URL {
|
||
t.Errorf("不健康目标未回退: %q", got2.URL)
|
||
}
|
||
})
|
||
}
|
||
|
||
// TestStickySession_CookieEncodeDecode 测试 cookie 编解码。
|
||
func TestStickySession_CookieEncodeDecode(t *testing.T) {
|
||
t.Parallel()
|
||
t.Run("编码解码round-trip", func(_ *testing.T) {
|
||
url := "http://backend1:8080"
|
||
expires := time.Now().Add(time.Hour)
|
||
encoded := encodeStickyCookie(url, expires)
|
||
if encoded == "" {
|
||
t.Fatal("encodeStickyCookie() 返回空字符串")
|
||
}
|
||
|
||
decodedURL, decodedExpires, ok := decodeStickyCookie(encoded)
|
||
if !ok {
|
||
t.Fatal("decodeStickyCookie() returned ok=false")
|
||
}
|
||
|
||
if decodedURL != url {
|
||
t.Errorf("解码后 URL = %q, want %q", decodedURL, url)
|
||
}
|
||
if decodedExpires.Unix() != expires.Unix() {
|
||
t.Errorf("解码后 expires = %v, want %v", decodedExpires, expires)
|
||
}
|
||
})
|
||
|
||
t.Run("空URL编码解码", func(_ *testing.T) {
|
||
expires := time.Now().Add(time.Hour)
|
||
encoded := encodeStickyCookie("", expires)
|
||
decodedURL, decodedExpires, ok := decodeStickyCookie(encoded)
|
||
if !ok {
|
||
t.Fatal("decodeStickyCookie() returned ok=false")
|
||
}
|
||
if decodedURL != "" {
|
||
t.Errorf("解码后 URL = %q, want 空字符串", decodedURL)
|
||
}
|
||
if decodedExpires.Unix() != expires.Unix() {
|
||
t.Errorf("解码后 expires = %v, want %v", decodedExpires, expires)
|
||
}
|
||
})
|
||
|
||
t.Run("无效编码", func(_ *testing.T) {
|
||
decodedURL, decodedExpires, ok := decodeStickyCookie("invalid-base64!!!")
|
||
if ok {
|
||
t.Errorf("decodeStickyCookie() = (%q, %v, %v), want ok=false", decodedURL, decodedExpires, ok)
|
||
}
|
||
})
|
||
}
|
||
|
||
// TestStickySession_Concurrent 测试并发安全。
|
||
// 100 个 goroutine 同时访问会话存储。
|
||
func TestStickySession_Concurrent(t *testing.T) {
|
||
t.Parallel()
|
||
config := DefaultStickyConfig()
|
||
config.Enabled = true
|
||
fallback := NewRoundRobin()
|
||
sticky := NewStickySession(config, fallback)
|
||
sticky.Start()
|
||
defer sticky.Stop()
|
||
|
||
targets := []*Target{
|
||
createHealthyTarget("http://backend1:8080", true),
|
||
createHealthyTarget("http://backend2:8080", true),
|
||
createHealthyTarget("http://backend3:8080", true),
|
||
}
|
||
|
||
var wg sync.WaitGroup
|
||
for i := 0; i < 100; i++ {
|
||
wg.Add(1)
|
||
go func(idx int) {
|
||
defer wg.Done()
|
||
ctx := &fasthttp.RequestCtx{}
|
||
// 交替使用有 cookie 和没有 cookie 的请求
|
||
if idx%2 == 0 {
|
||
ctx.Request.Header.SetCookie(config.Name, encodeStickyCookie("http://backend1:8080", time.Now().Add(time.Hour)))
|
||
}
|
||
got := sticky.Select(ctx, targets)
|
||
if got == nil {
|
||
t.Error("并发 Select() = nil")
|
||
}
|
||
}(i)
|
||
}
|
||
wg.Wait()
|
||
}
|
||
|
||
// TestStickySession_ExpiredCookie 测试过期 cookie 会导致回退到 fallback。
|
||
func TestStickySession_ExpiredCookie(t *testing.T) {
|
||
t.Parallel()
|
||
fallback := NewRoundRobin()
|
||
config := DefaultStickyConfig()
|
||
config.Enabled = true
|
||
config.Expires = -time.Hour // Negative = already expired
|
||
|
||
sticky := NewStickySession(config, fallback)
|
||
sticky.Start()
|
||
defer sticky.Stop()
|
||
|
||
// Make backend1 unavailable so we know for sure fallback picks backend2
|
||
targets := []*Target{
|
||
createHealthyTarget("http://backend1:8080", false),
|
||
createHealthyTarget("http://backend2:8080", true),
|
||
}
|
||
|
||
// Create an expired cookie manually
|
||
expiredCookie := encodeStickyCookie("http://backend1:8080", time.Now().Add(-time.Hour))
|
||
|
||
ctx := &fasthttp.RequestCtx{}
|
||
ctx.Request.Header.SetCookie(config.Name, expiredCookie)
|
||
|
||
// Should fallback because cookie is expired (and even if not, backend1 is unavailable)
|
||
selected := sticky.Select(ctx, targets)
|
||
if selected == nil {
|
||
t.Fatal("expected a target")
|
||
}
|
||
// Should not route to backend1 because cookie expired / unavailable
|
||
if selected.URL == "http://backend1:8080" {
|
||
t.Error("should not route using expired cookie")
|
||
}
|
||
// Should set a new cookie
|
||
newCookie := ctx.Response.Header.PeekCookie(config.Name)
|
||
if len(newCookie) == 0 {
|
||
t.Error("expected new cookie to be set")
|
||
}
|
||
}
|
||
|
||
// 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()
|
||
t.Run("SelectExcluding委托给fallback", func(_ *testing.T) {
|
||
config := DefaultStickyConfig()
|
||
config.Enabled = true
|
||
fallback := NewRoundRobin()
|
||
sticky := NewStickySession(config, fallback)
|
||
sticky.Start()
|
||
defer sticky.Stop()
|
||
|
||
targets := []*Target{
|
||
createHealthyTarget("http://backend1:8080", true),
|
||
createHealthyTarget("http://backend2:8080", true),
|
||
}
|
||
|
||
excluded := []*Target{targets[0]}
|
||
got := sticky.SelectExcluding(targets, excluded)
|
||
if got == nil {
|
||
t.Fatal("SelectExcluding() = nil")
|
||
}
|
||
if got.URL != "http://backend2:8080" {
|
||
t.Errorf("SelectExcluding() = %q, want %q", got.URL, "http://backend2:8080")
|
||
}
|
||
})
|
||
}
|