openteam/server/internal/proxy/affinity_test.go

package proxy

import (
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/openteam/server/internal/channel"
	"github.com/openteam/server/internal/store"
)

// 会话亲和：同一会话（同密钥+模型+同首条消息）粘住上次应答的渠道——
// 即使加权随机此时会选别的渠道；不同会话不受影响，照常负载均衡。
func TestAffinitySticksToAnsweringChannel(t *testing.T) {
	var ch1Hits, ch2Hits atomic.Int64
	ch1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ch1Hits.Add(1)
		_, _ = io.WriteString(w, `{"id":"ok","choices":[{"message":{"role":"assistant","content":"from-1"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)
	}))
	defer ch1.Close()
	ch2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ch2Hits.Add(1)
		_, _ = io.WriteString(w, `{"id":"ok","choices":[{"message":{"role":"assistant","content":"from-2"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)
	}))
	defer ch2.Close()

	h := newHarness(t, ch1.URL, 0, true, nil)
	h.gw.affinityTTL = 5 * time.Minute
	addChannel(t, h, "c2", ch2.URL, 1)
	pushWeightTo(t, h, "c1", 1000)

	sameConv := `{"model":"m1","messages":[{"role":"user","content":"会话的第一条消息"}]}`
	// 第一轮：权重压在渠道1，应答来自渠道1，亲和绑定建立
	if w := h.run(t, nil, sameConv); w.Code != http.StatusOK || !strings.Contains(w.Body.String(), "from-1") {
		t.Fatalf("第一轮应来自渠道1：%d %s", w.Code, w.Body.String())
	}
	// 反转权重：加权随机此刻会选渠道2
	pushWeightTo(t, h, "c1", 1)
	pushWeightTo(t, h, "c2", 1000)

	// 同一会话第二轮：仍应粘在渠道1（缓存值得保留）
	if w := h.run(t, nil, sameConv); !strings.Contains(w.Body.String(), "from-1") {
		t.Fatalf("同一会话应粘住渠道1：%s", w.Body.String())
	}
	// 不同会话（不同首条消息）：走加权随机，落到渠道2
	otherConv := `{"model":"m1","messages":[{"role":"user","content":"另一个会话的开场"}]}`
	if w := h.run(t, nil, otherConv); !strings.Contains(w.Body.String(), "from-2") {
		t.Fatalf("不同会话应走加权随机落到渠道2：%s", w.Body.String())
	}
	if ch1Hits.Load() != 2 || ch2Hits.Load() != 1 {
		t.Fatalf("命中次数不符：ch1=%d ch2=%d", ch1Hits.Load(), ch2Hits.Load())
	}
}

// 亲和渠道进入运行期冷却后自动失效，回退加权随机；不会把请求发给冷却中的渠道。
func TestAffinityRespectsCooldown(t *testing.T) {
	ch1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("冷却中的渠道1 不应收到请求")
		_, _ = io.WriteString(w, `{"id":"bad"}`)
	}))
	defer ch1.Close()
	ch2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = io.WriteString(w, `{"id":"ok","choices":[{"message":{"role":"assistant","content":"from-2"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)
	}))
	defer ch2.Close()

	h := newHarness(t, ch1.URL, 0, true, nil)
	h.gw.affinityTTL = 5 * time.Minute
	addChannel(t, h, "c2", ch2.URL, 1)

	// 预置亲和绑定指向渠道1，随后让渠道1 进运行期冷却
	var ch1row store.Channel
	h.db.Where("name = ?", "c1").First(&ch1row)
	h.gw.aff.put(affinityKey(1, "m1", []byte("会话的第一条消息")), ch1row.ID, 5*time.Minute)
	for i := 0; i < 3; i++ {
		h.gw.ch.ReportFailure(&ch1row, channel.FailRate, 0)
	}

	w := h.run(t, nil, `{"model":"m1","messages":[{"role":"user","content":"会话的第一条消息"}]}`)
	if !strings.Contains(w.Body.String(), "from-2") {
		t.Fatalf("冷却后亲和应失效并切换渠道2：%d %s", w.Code, w.Body.String())
	}
}

// 亲和键：同输入稳定、任一要素不同即不同；请求路径上不分配。
func TestAffinityKeyStable(t *testing.T) {
	prefix := []byte(`{"model":"m1","messages":[{"role":"user","content":"hi"}]}`)
	a := affinityKey(7, "m1", prefix)
	if b := affinityKey(7, "m1", prefix); a != b {
		t.Fatal("同一会话的键必须稳定")
	}
	if affinityKey(8, "m1", prefix) == a {
		t.Fatal("不同密钥应得到不同键")
	}
	if affinityKey(7, "m2", prefix) == a {
		t.Fatal("不同模型应得到不同键")
	}
	if affinityKey(7, "m1", []byte("other")) == a {
		t.Fatal("不同前缀应得到不同键")
	}
	// 前缀只有前 affinityPrefix 字节参与：同一个开头、更长的尾部不影响键
	head := make([]byte, affinityPrefix)
	for i := range head {
		head[i] = byte('a' + i%26)
	}
	tailA := append(append([]byte{}, head...), []byte("AAAA")...)
	tailB := append(append([]byte{}, head...), []byte("BBBB")...)
	if affinityKey(7, "m1", tailA) != affinityKey(7, "m1", tailB) {
		t.Fatal("超过前缀长度的部分不应影响键")
	}
	if affinityKey(7, "m1", head) != affinityKey(7, "m1", tailA) {
		t.Fatal("尾部长短不同但前 512 字节相同，键应一致")
	}
	inner := append([]byte{}, head...)
	inner[len(inner)-1] = 'z' // 改动前缀内最后一个字节
	if affinityKey(7, "m1", head) == affinityKey(7, "m1", inner) {
		t.Fatal("前缀以内的差异必须影响键")
	}
}

func TestAffinityKeyNoAlloc(t *testing.T) {
	prefix := []byte(`{"model":"m1"}`)
	allocs := testing.AllocsPerRun(200, func() { _ = affinityKey(7, "m1", prefix) })
	if allocs != 0 {
		t.Fatalf("affinityKey 分配了 %v 次，期望 0", allocs)
	}
}

// 分片表的过期与并发安全（-race 下验证）。
func TestAffinityTableSharded(t *testing.T) {
	tbl := newAffinityTable()
	tbl.put(1, 10, 50*time.Millisecond)
	tbl.put(2, 11, time.Minute)
	if id, ok := tbl.get(1); !ok || id != 10 {
		t.Fatalf("get(1) = %d/%v", id, ok)
	}
	if tbl.len() != 2 {
		t.Fatalf("len = %d, 期望 2", tbl.len())
	}
	time.Sleep(60 * time.Millisecond)
	if _, ok := tbl.get(1); ok {
		t.Fatal("过期条目应失效")
	}
	var wg sync.WaitGroup
	for i := 0; i < 8; i++ {
		wg.Add(1)
		go func(n int) {
			defer wg.Done()
			for j := 0; j < 200; j++ {
				k := uint64(n*1000 + j)
				tbl.put(k, k, time.Minute)
				tbl.get(k)
			}
		}(i)
	}
	wg.Wait()
}