openai_ws_pool_test.go 77 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006200720082009201020112012201320142015201620172018201920202021202220232024202520262027202820292030203120322033203420352036203720382039204020412042204320442045204620472048204920502051205220532054205520562057205820592060206120622063206420652066206720682069207020712072207320742075207620772078207920802081208220832084208520862087208820892090209120922093209420952096209720982099210021012102210321042105210621072108210921102111211221132114211521162117211821192120212121222123212421252126212721282129213021312132213321342135213621372138213921402141214221432144214521462147214821492150215121522153215421552156215721582159216021612162216321642165216621672168216921702171217221732174217521762177217821792180218121822183218421852186218721882189219021912192219321942195219621972198219922002201220222032204220522062207220822092210221122122213221422152216221722182219222022212222222322242225222622272228222922302231223222332234223522362237223822392240224122422243224422452246224722482249225022512252225322542255225622572258225922602261226222632264226522662267226822692270227122722273227422752276227722782279228022812282228322842285228622872288228922902291229222932294229522962297229822992300230123022303230423052306230723082309231023112312231323142315231623172318231923202321232223232324232523262327232823292330233123322333233423352336233723382339234023412342234323442345234623472348234923502351235223532354235523562357235823592360236123622363236423652366236723682369237023712372237323742375237623772378237923802381238223832384238523862387238823892390239123922393239423952396
  1. package service
  2. import (
  3. "context"
  4. "errors"
  5. "net/http"
  6. "strings"
  7. "sync"
  8. "sync/atomic"
  9. "testing"
  10. "time"
  11. "github.com/Wei-Shaw/sub2api/internal/config"
  12. "github.com/stretchr/testify/require"
  13. )
  14. func TestOpenAIWSConnPool_CleanupStaleAndTrimIdle(t *testing.T) {
  15. cfg := &config.Config{}
  16. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  17. pool := newOpenAIWSConnPool(cfg)
  18. accountID := int64(10)
  19. ap := pool.getOrCreateAccountPool(accountID)
  20. stale := newOpenAIWSConn("stale", accountID, nil, nil)
  21. stale.createdAtNano.Store(time.Now().Add(-2 * time.Hour).UnixNano())
  22. stale.lastUsedNano.Store(time.Now().Add(-2 * time.Hour).UnixNano())
  23. idleOld := newOpenAIWSConn("idle_old", accountID, nil, nil)
  24. idleOld.lastUsedNano.Store(time.Now().Add(-10 * time.Minute).UnixNano())
  25. idleNew := newOpenAIWSConn("idle_new", accountID, nil, nil)
  26. idleNew.lastUsedNano.Store(time.Now().Add(-1 * time.Minute).UnixNano())
  27. ap.conns[stale.id] = stale
  28. ap.conns[idleOld.id] = idleOld
  29. ap.conns[idleNew.id] = idleNew
  30. evicted := pool.cleanupAccountLocked(ap, time.Now(), pool.maxConnsHardCap())
  31. closeOpenAIWSConns(evicted)
  32. require.Nil(t, ap.conns["stale"], "stale connection should be rotated")
  33. require.Nil(t, ap.conns["idle_old"], "old idle should be trimmed by max_idle")
  34. require.NotNil(t, ap.conns["idle_new"], "newer idle should be kept")
  35. }
  36. func TestOpenAIWSConnPool_NextConnIDFormat(t *testing.T) {
  37. pool := newOpenAIWSConnPool(&config.Config{})
  38. id1 := pool.nextConnID(42)
  39. id2 := pool.nextConnID(42)
  40. require.True(t, strings.HasPrefix(id1, "oa_ws_42_"))
  41. require.True(t, strings.HasPrefix(id2, "oa_ws_42_"))
  42. require.NotEqual(t, id1, id2)
  43. require.Equal(t, "oa_ws_42_1", id1)
  44. require.Equal(t, "oa_ws_42_2", id2)
  45. }
  46. func TestOpenAIWSConnPool_AcquireCleanupInterval(t *testing.T) {
  47. require.Equal(t, 3*time.Second, openAIWSAcquireCleanupInterval)
  48. require.Less(t, openAIWSAcquireCleanupInterval, openAIWSBackgroundSweepTicker)
  49. }
  50. func TestNormalizeOpenAIWSRoutingAffinityPrefersCanonicalAndSortsVariants(t *testing.T) {
  51. headers := http.Header{
  52. "X-CODEX-ROUTING-HINT": []string{" variant-uppercase "},
  53. "X-Codex-Routing-Hint": []string{" ", " canonical "},
  54. }
  55. require.Equal(t, "canonical", normalizeOpenAIWSRoutingAffinity(headers))
  56. delete(headers, "X-Codex-Routing-Hint")
  57. require.Equal(t, "variant-uppercase", normalizeOpenAIWSRoutingAffinity(headers))
  58. }
  59. func TestOpenAIWSConnLease_WriteJSONAndGuards(t *testing.T) {
  60. conn := newOpenAIWSConn("lease_write", 1, &openAIWSFakeConn{}, nil)
  61. lease := &openAIWSConnLease{conn: conn}
  62. require.NoError(t, lease.WriteJSON(map[string]any{"type": "response.create"}, 0))
  63. var nilLease *openAIWSConnLease
  64. err := nilLease.WriteJSONWithContextTimeout(context.Background(), map[string]any{"type": "response.create"}, time.Second)
  65. require.ErrorIs(t, err, errOpenAIWSConnClosed)
  66. err = (&openAIWSConnLease{}).WriteJSONWithContextTimeout(context.Background(), map[string]any{"type": "response.create"}, time.Second)
  67. require.ErrorIs(t, err, errOpenAIWSConnClosed)
  68. }
  69. func TestOpenAIWSConn_WriteJSONWithTimeout_NilParentContextUsesBackground(t *testing.T) {
  70. probe := &openAIWSContextProbeConn{}
  71. conn := newOpenAIWSConn("ctx_probe", 1, probe, nil)
  72. require.NoError(t, conn.writeJSONWithTimeout(context.Background(), map[string]any{"type": "response.create"}, 0))
  73. require.NotNil(t, probe.lastWriteCtx)
  74. }
  75. func TestOpenAIWSConnPool_TargetConnCountAdaptive(t *testing.T) {
  76. cfg := &config.Config{}
  77. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 6
  78. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 1
  79. cfg.Gateway.OpenAIWS.PoolTargetUtilization = 0.5
  80. pool := newOpenAIWSConnPool(cfg)
  81. ap := pool.getOrCreateAccountPool(88)
  82. conn1 := newOpenAIWSConn("c1", 88, nil, nil)
  83. conn2 := newOpenAIWSConn("c2", 88, nil, nil)
  84. require.True(t, conn1.tryAcquire())
  85. require.True(t, conn2.tryAcquire())
  86. conn1.waiters.Store(1)
  87. conn2.waiters.Store(1)
  88. ap.conns[conn1.id] = conn1
  89. ap.conns[conn2.id] = conn2
  90. target := pool.targetConnCountLocked(ap, pool.maxConnsHardCap())
  91. require.Equal(t, 6, target, "应按 inflight+waiters 与 target_utilization 自适应扩容到上限")
  92. conn1.release()
  93. conn2.release()
  94. conn1.waiters.Store(0)
  95. conn2.waiters.Store(0)
  96. target = pool.targetConnCountLocked(ap, pool.maxConnsHardCap())
  97. require.Equal(t, 1, target, "低负载时应缩回到最小空闲连接")
  98. }
  99. func TestOpenAIWSConnPool_TargetConnCountMinIdleZero(t *testing.T) {
  100. cfg := &config.Config{}
  101. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 4
  102. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  103. cfg.Gateway.OpenAIWS.PoolTargetUtilization = 0.8
  104. pool := newOpenAIWSConnPool(cfg)
  105. ap := pool.getOrCreateAccountPool(66)
  106. target := pool.targetConnCountLocked(ap, pool.maxConnsHardCap())
  107. require.Equal(t, 0, target, "min_idle=0 且无负载时应允许缩容到 0")
  108. }
  109. func TestOpenAIWSConnPool_EnsureTargetIdleAsync(t *testing.T) {
  110. cfg := &config.Config{}
  111. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 4
  112. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 2
  113. cfg.Gateway.OpenAIWS.PoolTargetUtilization = 0.8
  114. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 1
  115. pool := newOpenAIWSConnPool(cfg)
  116. pool.setClientDialerForTest(&openAIWSFakeDialer{})
  117. accountID := int64(77)
  118. account := &Account{ID: accountID, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  119. ap := pool.getOrCreateAccountPool(accountID)
  120. ap.mu.Lock()
  121. ap.lastAcquire = &openAIWSAcquireRequest{
  122. Account: account,
  123. WSURL: "wss://example.com/v1/responses",
  124. }
  125. ap.mu.Unlock()
  126. pool.ensureTargetIdleAsync(accountID)
  127. require.Eventually(t, func() bool {
  128. ap, ok := pool.getAccountPool(accountID)
  129. if !ok || ap == nil {
  130. return false
  131. }
  132. ap.mu.Lock()
  133. defer ap.mu.Unlock()
  134. return len(ap.conns) >= 2
  135. }, 2*time.Second, 20*time.Millisecond)
  136. metrics := pool.SnapshotMetrics()
  137. require.GreaterOrEqual(t, metrics.ScaleUpTotal, int64(2))
  138. }
  139. func TestOpenAIWSConnPool_EnsureTargetIdleAsyncCooldown(t *testing.T) {
  140. cfg := &config.Config{}
  141. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 4
  142. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 2
  143. cfg.Gateway.OpenAIWS.PoolTargetUtilization = 0.8
  144. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 1
  145. cfg.Gateway.OpenAIWS.PrewarmCooldownMS = 500
  146. pool := newOpenAIWSConnPool(cfg)
  147. dialer := &openAIWSCountingDialer{}
  148. pool.setClientDialerForTest(dialer)
  149. accountID := int64(178)
  150. account := &Account{ID: accountID, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  151. ap := pool.getOrCreateAccountPool(accountID)
  152. ap.mu.Lock()
  153. ap.lastAcquire = &openAIWSAcquireRequest{
  154. Account: account,
  155. WSURL: "wss://example.com/v1/responses",
  156. }
  157. ap.mu.Unlock()
  158. pool.ensureTargetIdleAsync(accountID)
  159. require.Eventually(t, func() bool {
  160. ap, ok := pool.getAccountPool(accountID)
  161. if !ok || ap == nil {
  162. return false
  163. }
  164. ap.mu.Lock()
  165. defer ap.mu.Unlock()
  166. return len(ap.conns) >= 2 && !ap.prewarmActive
  167. }, 2*time.Second, 20*time.Millisecond)
  168. firstDialCount := dialer.DialCount()
  169. require.GreaterOrEqual(t, firstDialCount, 2)
  170. // 人工制造缺口触发新一轮预热需求。
  171. ap, ok := pool.getAccountPool(accountID)
  172. require.True(t, ok)
  173. require.NotNil(t, ap)
  174. ap.mu.Lock()
  175. for id := range ap.conns {
  176. delete(ap.conns, id)
  177. break
  178. }
  179. ap.mu.Unlock()
  180. pool.ensureTargetIdleAsync(accountID)
  181. time.Sleep(120 * time.Millisecond)
  182. require.Equal(t, firstDialCount, dialer.DialCount(), "cooldown 窗口内不应再次触发预热")
  183. time.Sleep(450 * time.Millisecond)
  184. pool.ensureTargetIdleAsync(accountID)
  185. require.Eventually(t, func() bool {
  186. return dialer.DialCount() > firstDialCount
  187. }, 2*time.Second, 20*time.Millisecond)
  188. }
  189. func TestOpenAIWSConnPool_EnsureTargetIdleAsyncFailureSuppress(t *testing.T) {
  190. cfg := &config.Config{}
  191. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
  192. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 1
  193. cfg.Gateway.OpenAIWS.PoolTargetUtilization = 0.8
  194. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 1
  195. cfg.Gateway.OpenAIWS.PrewarmCooldownMS = 0
  196. pool := newOpenAIWSConnPool(cfg)
  197. dialer := &openAIWSAlwaysFailDialer{}
  198. pool.setClientDialerForTest(dialer)
  199. accountID := int64(279)
  200. account := &Account{ID: accountID, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  201. ap := pool.getOrCreateAccountPool(accountID)
  202. ap.mu.Lock()
  203. ap.lastAcquire = &openAIWSAcquireRequest{
  204. Account: account,
  205. WSURL: "wss://example.com/v1/responses",
  206. }
  207. ap.mu.Unlock()
  208. pool.ensureTargetIdleAsync(accountID)
  209. require.Eventually(t, func() bool {
  210. ap, ok := pool.getAccountPool(accountID)
  211. if !ok || ap == nil {
  212. return false
  213. }
  214. ap.mu.Lock()
  215. defer ap.mu.Unlock()
  216. return !ap.prewarmActive
  217. }, 2*time.Second, 20*time.Millisecond)
  218. pool.ensureTargetIdleAsync(accountID)
  219. require.Eventually(t, func() bool {
  220. ap, ok := pool.getAccountPool(accountID)
  221. if !ok || ap == nil {
  222. return false
  223. }
  224. ap.mu.Lock()
  225. defer ap.mu.Unlock()
  226. return !ap.prewarmActive
  227. }, 2*time.Second, 20*time.Millisecond)
  228. require.Equal(t, 2, dialer.DialCount())
  229. // 连续失败达到阈值后,新的预热触发应被抑制,不再继续拨号。
  230. pool.ensureTargetIdleAsync(accountID)
  231. time.Sleep(120 * time.Millisecond)
  232. require.Equal(t, 2, dialer.DialCount())
  233. }
  234. func TestOpenAIWSConnPool_AcquireQueueWaitMetrics(t *testing.T) {
  235. cfg := &config.Config{}
  236. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  237. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  238. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 4
  239. pool := newOpenAIWSConnPool(cfg)
  240. accountID := int64(99)
  241. account := &Account{ID: accountID, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  242. conn := newOpenAIWSConn("busy", accountID, &openAIWSFakeConn{}, nil)
  243. require.True(t, conn.tryAcquire()) // 占用连接,触发后续排队
  244. ap := pool.ensureAccountPoolLocked(accountID)
  245. ap.mu.Lock()
  246. ap.conns[conn.id] = conn
  247. ap.lastAcquire = &openAIWSAcquireRequest{
  248. Account: account,
  249. WSURL: "wss://example.com/v1/responses",
  250. }
  251. ap.mu.Unlock()
  252. go func() {
  253. time.Sleep(60 * time.Millisecond)
  254. conn.release()
  255. }()
  256. lease, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
  257. Account: account,
  258. WSURL: "wss://example.com/v1/responses",
  259. })
  260. require.NoError(t, err)
  261. require.NotNil(t, lease)
  262. require.True(t, lease.Reused())
  263. require.GreaterOrEqual(t, lease.QueueWaitDuration(), 50*time.Millisecond)
  264. lease.Release()
  265. metrics := pool.SnapshotMetrics()
  266. require.GreaterOrEqual(t, metrics.AcquireQueueWaitTotal, int64(1))
  267. require.Greater(t, metrics.AcquireQueueWaitMsTotal, int64(0))
  268. require.GreaterOrEqual(t, metrics.ConnPickTotal, int64(1))
  269. }
  270. func TestOpenAIWSConnPool_DialSuccessWakesTopologyWaiterAndCanceledWaiterDoesNotLoseLease(t *testing.T) {
  271. cfg := &config.Config{}
  272. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  273. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  274. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  275. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 4
  276. pool := newOpenAIWSConnPool(cfg)
  277. dialer := newOpenAIWSFirstDialBlockingCaptureDialer()
  278. pool.setClientDialerForTest(dialer)
  279. account := &Account{ID: 991, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  280. req := openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"}
  281. type result struct {
  282. lease *openAIWSConnLease
  283. err error
  284. }
  285. firstCh := make(chan result, 1)
  286. go func() {
  287. lease, err := pool.Acquire(context.Background(), req)
  288. firstCh <- result{lease: lease, err: err}
  289. }()
  290. <-dialer.firstStarted
  291. waitCtx, cancelWait := context.WithCancel(context.Background())
  292. secondCh := make(chan result, 1)
  293. waitReq := req
  294. waitReq.Headers = http.Header{openAICodexRoutingHintHeader: {"model=gpt-5.6-codex;tier=priority"}}
  295. go func() {
  296. lease, err := pool.Acquire(waitCtx, waitReq)
  297. secondCh <- result{lease: lease, err: err}
  298. }()
  299. close(dialer.releaseFirst)
  300. first := <-firstCh
  301. require.NoError(t, first.err)
  302. require.NotNil(t, first.lease)
  303. // The second acquire initially waits on the account topology channel while
  304. // the first dial is in flight. Dial success must wake it immediately so it
  305. // can queue on the newly-created (still leased) connection.
  306. require.Eventually(t, func() bool {
  307. ap, ok := pool.getAccountPool(account.ID)
  308. if !ok || ap == nil {
  309. return false
  310. }
  311. ap.mu.Lock()
  312. defer ap.mu.Unlock()
  313. for _, conn := range ap.conns {
  314. if conn != nil && conn.waiters.Load() == 1 {
  315. return true
  316. }
  317. }
  318. return false
  319. }, time.Second, 5*time.Millisecond)
  320. cancelWait()
  321. second := <-secondCh
  322. require.ErrorIs(t, second.err, context.Canceled)
  323. require.Nil(t, second.lease)
  324. ap, ok := pool.getAccountPool(account.ID)
  325. require.True(t, ok)
  326. ap.mu.Lock()
  327. require.NotNil(t, ap.lastAcquire)
  328. require.Empty(t, normalizeOpenAIWSRoutingAffinity(ap.lastAcquire.Headers), "a canceled acquire must not replace the successful prewarm target")
  329. ap.mu.Unlock()
  330. first.lease.Release()
  331. third, err := pool.Acquire(context.Background(), req)
  332. require.NoError(t, err)
  333. require.True(t, third.Reused(), "a canceled waiter must not consume the released semaphore token")
  334. require.Equal(t, first.lease.ConnID(), third.ConnID())
  335. third.Release()
  336. }
  337. func TestOpenAIWSConnPool_PrewarmHintChangeDoesNotInvalidateHealthyDial(t *testing.T) {
  338. cfg := &config.Config{}
  339. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  340. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 1
  341. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  342. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 2
  343. pool := newOpenAIWSConnPool(cfg)
  344. dialer := newOpenAIWSFirstDialBlockingCaptureDialer()
  345. pool.setClientDialerForTest(dialer)
  346. account := &Account{ID: 992, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  347. oldHeaders := make(http.Header)
  348. oldHeaders.Set(openAICodexRoutingHintHeader, "model=gpt-5.6-codex")
  349. newHeaders := make(http.Header)
  350. newHeaders.Set(openAICodexRoutingHintHeader, "model=gpt-5.6-codex;tier=priority")
  351. ap := pool.getOrCreateAccountPool(account.ID)
  352. ap.mu.Lock()
  353. ap.lastAcquire = &openAIWSAcquireRequest{
  354. Account: account,
  355. WSURL: "wss://example.com/v1/responses",
  356. Headers: oldHeaders,
  357. }
  358. ap.mu.Unlock()
  359. pool.ensureTargetIdleAsync(account.ID)
  360. <-dialer.firstStarted
  361. // Simulate a newer priority target arriving while the old model-only
  362. // prewarm dial is in flight. Routing hints are advisory, so this alone must
  363. // not discard an otherwise compatible connection.
  364. ap.mu.Lock()
  365. ap.lastAcquire = &openAIWSAcquireRequest{
  366. Account: account,
  367. WSURL: "wss://example.com/v1/responses",
  368. Headers: newHeaders,
  369. }
  370. ap.mu.Unlock()
  371. close(dialer.releaseFirst)
  372. require.Eventually(t, func() bool {
  373. ap.mu.Lock()
  374. defer ap.mu.Unlock()
  375. if ap.prewarmActive || len(ap.conns) != 1 {
  376. return false
  377. }
  378. for _, conn := range ap.conns {
  379. return conn != nil && conn.routingAffinity == "model=gpt-5.6-codex"
  380. }
  381. return false
  382. }, 2*time.Second, 10*time.Millisecond)
  383. require.Equal(t, 1, dialer.DialCount(), "routing-hint-only changes must not turn advisory metadata into hard reconnects")
  384. }
  385. func TestOpenAIWSConnPool_ClearAccountWakesIncompatibleTopologyWaiter(t *testing.T) {
  386. cfg := &config.Config{}
  387. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  388. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  389. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  390. pool := newOpenAIWSConnPool(cfg)
  391. dialer := &openAIWSCountingDialer{}
  392. pool.setClientDialerForTest(dialer)
  393. account := &Account{ID: 993, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  394. baseReq := openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"}
  395. betaAReq := baseReq
  396. betaAReq.Headers = http.Header{"X-Codex-Beta-Features": {"feature_a"}}
  397. betaBReq := baseReq
  398. betaBReq.Headers = http.Header{"X-Codex-Beta-Features": {"feature_b"}}
  399. busy, err := pool.Acquire(context.Background(), betaAReq)
  400. require.NoError(t, err)
  401. require.Equal(t, 1, dialer.DialCount())
  402. type result struct {
  403. lease *openAIWSConnLease
  404. err error
  405. }
  406. resultCh := make(chan result, 1)
  407. waitCtx, cancelWait := context.WithTimeout(context.Background(), time.Second)
  408. defer cancelWait()
  409. go func() {
  410. lease, acquireErr := pool.Acquire(waitCtx, betaBReq)
  411. resultCh <- result{lease: lease, err: acquireErr}
  412. }()
  413. require.Never(t, func() bool { return dialer.DialCount() > 1 }, 50*time.Millisecond, 5*time.Millisecond)
  414. pool.ClearAccount(account.ID)
  415. resultB := <-resultCh
  416. require.NoError(t, resultB.err)
  417. require.NotNil(t, resultB.lease)
  418. require.False(t, resultB.lease.Reused())
  419. require.Equal(t, 2, dialer.DialCount(), "ClearAccount must wake the waiter to redial immediately")
  420. resultB.lease.Release()
  421. busy.Release()
  422. }
  423. func TestOpenAIWSConnPool_ClearAccountDoesNotReviveInFlightDialGeneration(t *testing.T) {
  424. cfg := &config.Config{}
  425. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  426. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  427. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  428. pool := newOpenAIWSConnPool(cfg)
  429. dialer := newOpenAIWSFirstDialBlockingCaptureDialer()
  430. pool.setClientDialerForTest(dialer)
  431. account := &Account{ID: 994, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  432. req := openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"}
  433. type result struct {
  434. lease *openAIWSConnLease
  435. err error
  436. }
  437. resultCh := make(chan result, 1)
  438. go func() {
  439. lease, err := pool.Acquire(context.Background(), req)
  440. resultCh <- result{lease: lease, err: err}
  441. }()
  442. <-dialer.firstStarted
  443. pool.ClearAccount(account.ID)
  444. close(dialer.releaseFirst)
  445. got := <-resultCh
  446. require.NoError(t, got.err)
  447. require.NotNil(t, got.lease)
  448. require.Equal(t, 2, dialer.DialCount(), "the pre-clear dial must be discarded and retried in the new generation")
  449. require.True(t, strings.HasSuffix(got.lease.ConnID(), "_2"))
  450. ap, ok := pool.getAccountPool(account.ID)
  451. require.True(t, ok)
  452. ap.mu.Lock()
  453. require.Equal(t, uint64(1), ap.generation)
  454. require.Len(t, ap.conns, 1)
  455. require.NotNil(t, ap.lastAcquire, "only the post-clear successful acquire may restore the prewarm target")
  456. ap.mu.Unlock()
  457. got.lease.Release()
  458. }
  459. func TestOpenAIWSConnPool_ForceNewConnSkipsReuse(t *testing.T) {
  460. cfg := &config.Config{}
  461. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
  462. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  463. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
  464. pool := newOpenAIWSConnPool(cfg)
  465. dialer := &openAIWSCountingDialer{}
  466. pool.setClientDialerForTest(dialer)
  467. account := &Account{ID: 123, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  468. lease1, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
  469. Account: account,
  470. WSURL: "wss://example.com/v1/responses",
  471. })
  472. require.NoError(t, err)
  473. require.NotNil(t, lease1)
  474. lease1.Release()
  475. lease2, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
  476. Account: account,
  477. WSURL: "wss://example.com/v1/responses",
  478. ForceNewConn: true,
  479. })
  480. require.NoError(t, err)
  481. require.NotNil(t, lease2)
  482. lease2.Release()
  483. require.Equal(t, 2, dialer.DialCount(), "ForceNewConn=true 时应跳过空闲连接复用并新建连接")
  484. }
  485. func TestOpenAIWSConnPool_AcquireReusesOnlyMatchingBetaFeatures(t *testing.T) {
  486. cfg := &config.Config{}
  487. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
  488. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  489. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
  490. pool := newOpenAIWSConnPool(cfg)
  491. dialer := &openAIWSCountingDialer{}
  492. pool.setClientDialerForTest(dialer)
  493. account := &Account{ID: 128, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  494. baseReq := openAIWSAcquireRequest{
  495. Account: account,
  496. WSURL: "wss://example.com/v1/responses",
  497. }
  498. plainLease, err := pool.Acquire(context.Background(), baseReq)
  499. require.NoError(t, err)
  500. plainConnID := plainLease.ConnID()
  501. plainLease.Release()
  502. betaReq := baseReq
  503. betaReq.Headers = http.Header{"X-Codex-Beta-Features": {" remote_compaction_v2 ", " responses_websockets_v2 "}}
  504. betaLease, err := pool.Acquire(context.Background(), betaReq)
  505. require.NoError(t, err)
  506. require.False(t, betaLease.Reused())
  507. require.NotEqual(t, plainConnID, betaLease.ConnID())
  508. betaConnID := betaLease.ConnID()
  509. betaLease.Release()
  510. reorderedReq := baseReq
  511. reorderedReq.Headers = http.Header{"X-Codex-Beta-Features": {"responses_websockets_v2,remote_compaction_v2"}}
  512. reorderedLease, err := pool.Acquire(context.Background(), reorderedReq)
  513. require.NoError(t, err)
  514. require.True(t, reorderedLease.Reused())
  515. require.Equal(t, betaConnID, reorderedLease.ConnID())
  516. reorderedLease.Release()
  517. _, err = pool.Acquire(context.Background(), openAIWSAcquireRequest{
  518. Account: account,
  519. WSURL: baseReq.WSURL,
  520. Headers: betaReq.Headers,
  521. PreferredConnID: plainConnID,
  522. ForcePreferredConn: true,
  523. })
  524. require.ErrorIs(t, err, errOpenAIWSPreferredConnUnavailable)
  525. plainLease, err = pool.Acquire(context.Background(), baseReq)
  526. require.NoError(t, err)
  527. require.True(t, plainLease.Reused())
  528. require.Equal(t, plainConnID, plainLease.ConnID())
  529. plainLease.Release()
  530. require.Equal(t, 2, dialer.DialCount())
  531. }
  532. func activeCodexFingerprintPoolAccountForTest(id int64) *Account {
  533. return &Account{
  534. ID: id,
  535. Platform: PlatformOpenAI,
  536. Type: AccountTypeOAuth,
  537. Extra: map[string]any{
  538. codexFingerprintModeExtraKey: "session",
  539. codexFingerprintSeedExtraKey: "11111111-1111-4111-8111-111111111111",
  540. },
  541. }
  542. }
  543. func stableOpenAIWSIdentityHeadersForTest() http.Header {
  544. headers := make(http.Header)
  545. headers.Set("X-Codex-Beta-Features", "remote_compaction_v2,responses_websockets_v2")
  546. headers.Set("X-Codex-Installation-ID", "install-a")
  547. headers.Set("session-id", "session-hyphen-a")
  548. headers.Set("session_id", "session-underscore-a")
  549. headers.Set("thread-id", "thread-a")
  550. headers.Set("x-client-request-id", "client-request-a")
  551. headers.Set("x-codex-window-id", "window-a")
  552. return headers
  553. }
  554. func TestOpenAIWSConnPool_AcquireReusesSameStableIdentityWithDifferentTurnMetadata(t *testing.T) {
  555. cfg := &config.Config{}
  556. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  557. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  558. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  559. pool := newOpenAIWSConnPool(cfg)
  560. dialer := &openAIWSCountingDialer{}
  561. pool.setClientDialerForTest(dialer)
  562. account := activeCodexFingerprintPoolAccountForTest(132)
  563. headers := stableOpenAIWSIdentityHeadersForTest()
  564. headers.Set("Authorization", "Bearer token-a")
  565. headers.Set("x-codex-turn-metadata", `{"turn_id":"turn-a"}`)
  566. first, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
  567. Account: account,
  568. WSURL: "wss://example.com/v1/responses",
  569. Headers: headers,
  570. })
  571. require.NoError(t, err)
  572. firstConnID := first.ConnID()
  573. first.Release()
  574. nextHeaders := stableOpenAIWSIdentityHeadersForTest()
  575. nextHeaders.Set("Authorization", "Bearer token-b")
  576. nextHeaders.Set("x-codex-turn-metadata", `{"turn_id":"turn-b"}`)
  577. nextHeaders.Set(openAICodexRoutingHintHeader, "model=gpt-5.6-codex;tier=priority")
  578. second, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
  579. Account: account,
  580. WSURL: "wss://example.com/v1/responses",
  581. Headers: nextHeaders,
  582. })
  583. require.NoError(t, err)
  584. require.True(t, second.Reused())
  585. require.Equal(t, firstConnID, second.ConnID())
  586. second.Release()
  587. require.Equal(t, 1, dialer.DialCount(), "stable identity match should ignore auth, turn metadata, and soft routing hints")
  588. }
  589. func TestOpenAIWSConnPool_AcquireDoesNotReuseDifferentStableIdentity(t *testing.T) {
  590. for _, tt := range []struct {
  591. name string
  592. header string
  593. value string
  594. }{
  595. {name: "installation", header: "x-codex-installation-id", value: "install-b"},
  596. {name: "session hyphen", header: "session-id", value: "session-hyphen-b"},
  597. {name: "session underscore", header: "session_id", value: "session-underscore-b"},
  598. {name: "thread", header: "thread-id", value: "thread-b"},
  599. {name: "client request", header: "x-client-request-id", value: "client-request-b"},
  600. {name: "window", header: "x-codex-window-id", value: "window-b"},
  601. } {
  602. t.Run(tt.name, func(t *testing.T) {
  603. cfg := &config.Config{}
  604. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
  605. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  606. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
  607. pool := newOpenAIWSConnPool(cfg)
  608. dialer := &openAIWSCountingDialer{}
  609. pool.setClientDialerForTest(dialer)
  610. account := activeCodexFingerprintPoolAccountForTest(133)
  611. first, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
  612. Account: account,
  613. WSURL: "wss://example.com/v1/responses",
  614. Headers: stableOpenAIWSIdentityHeadersForTest(),
  615. })
  616. require.NoError(t, err)
  617. firstConnID := first.ConnID()
  618. first.Release()
  619. nextHeaders := stableOpenAIWSIdentityHeadersForTest()
  620. nextHeaders.Set(tt.header, tt.value)
  621. second, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
  622. Account: account,
  623. WSURL: "wss://example.com/v1/responses",
  624. Headers: nextHeaders,
  625. })
  626. require.NoError(t, err)
  627. require.False(t, second.Reused())
  628. require.NotEqual(t, firstConnID, second.ConnID())
  629. second.Release()
  630. require.Equal(t, 2, dialer.DialCount())
  631. })
  632. }
  633. }
  634. func TestOpenAIWSConnPool_AcquireRoutingHintRemainsSoftAffinity(t *testing.T) {
  635. cfg := &config.Config{}
  636. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  637. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  638. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  639. pool := newOpenAIWSConnPool(cfg)
  640. dialer := &openAIWSCountingDialer{}
  641. pool.setClientDialerForTest(dialer)
  642. account := activeCodexFingerprintPoolAccountForTest(134)
  643. firstHeaders := stableOpenAIWSIdentityHeadersForTest()
  644. firstHeaders.Set(openAICodexRoutingHintHeader, "model=gpt-5.6-codex")
  645. first, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
  646. Account: account,
  647. WSURL: "wss://example.com/v1/responses",
  648. Headers: firstHeaders,
  649. })
  650. require.NoError(t, err)
  651. firstConnID := first.ConnID()
  652. first.Release()
  653. secondHeaders := stableOpenAIWSIdentityHeadersForTest()
  654. secondHeaders.Set(openAICodexRoutingHintHeader, "model=gpt-5.6-codex;tier=priority")
  655. second, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
  656. Account: account,
  657. WSURL: "wss://example.com/v1/responses",
  658. Headers: secondHeaders,
  659. })
  660. require.NoError(t, err)
  661. require.True(t, second.Reused())
  662. require.Equal(t, firstConnID, second.ConnID())
  663. second.Release()
  664. require.Equal(t, 1, dialer.DialCount())
  665. }
  666. func TestOpenAIWSConnPool_DeviceModeKeysOnlyInstallationIdentity(t *testing.T) {
  667. cfg := &config.Config{}
  668. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
  669. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  670. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
  671. pool := newOpenAIWSConnPool(cfg)
  672. dialer := &openAIWSCountingDialer{}
  673. pool.setClientDialerForTest(dialer)
  674. account := activeCodexFingerprintPoolAccountForTest(135)
  675. account.Extra[codexFingerprintModeExtraKey] = "device"
  676. firstHeaders := stableOpenAIWSIdentityHeadersForTest()
  677. first, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
  678. Account: account,
  679. WSURL: "wss://example.com/v1/responses",
  680. Headers: firstHeaders,
  681. })
  682. require.NoError(t, err)
  683. firstConnID := first.ConnID()
  684. first.Release()
  685. sessionChanged := stableOpenAIWSIdentityHeadersForTest()
  686. sessionChanged.Set("session-id", "session-hyphen-b")
  687. sessionChanged.Set("session_id", "session-underscore-b")
  688. sessionChanged.Set("thread-id", "thread-b")
  689. sessionChanged.Set("x-client-request-id", "client-request-b")
  690. sessionChanged.Set("x-codex-window-id", "window-b")
  691. second, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
  692. Account: account,
  693. WSURL: "wss://example.com/v1/responses",
  694. Headers: sessionChanged,
  695. })
  696. require.NoError(t, err)
  697. require.True(t, second.Reused())
  698. require.Equal(t, firstConnID, second.ConnID())
  699. second.Release()
  700. installationChanged := sessionChanged.Clone()
  701. installationChanged.Set("x-codex-installation-id", "install-b")
  702. third, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
  703. Account: account,
  704. WSURL: "wss://example.com/v1/responses",
  705. Headers: installationChanged,
  706. })
  707. require.NoError(t, err)
  708. require.False(t, third.Reused())
  709. require.NotEqual(t, firstConnID, third.ConnID())
  710. third.Release()
  711. require.Equal(t, 2, dialer.DialCount())
  712. }
  713. func TestOpenAIWSConnPool_AcquireReplacesIdleConnWithDifferentBetaFeatures(t *testing.T) {
  714. cfg := &config.Config{}
  715. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  716. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  717. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  718. pool := newOpenAIWSConnPool(cfg)
  719. dialer := &openAIWSCountingDialer{}
  720. pool.setClientDialerForTest(dialer)
  721. account := &Account{ID: 129, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  722. plainLease, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
  723. Account: account,
  724. WSURL: "wss://example.com/v1/responses",
  725. })
  726. require.NoError(t, err)
  727. plainConnID := plainLease.ConnID()
  728. plainLease.Release()
  729. betaLease, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
  730. Account: account,
  731. WSURL: "wss://example.com/v1/responses",
  732. Headers: http.Header{"X-Codex-Beta-Features": {"remote_compaction_v2"}},
  733. })
  734. require.NoError(t, err)
  735. require.False(t, betaLease.Reused())
  736. require.NotEqual(t, plainConnID, betaLease.ConnID())
  737. betaLease.Release()
  738. require.Equal(t, 2, dialer.DialCount())
  739. }
  740. func TestOpenAIWSConnPool_AcquireWaitsForBusyIncompatibleConnection(t *testing.T) {
  741. cfg := &config.Config{}
  742. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  743. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  744. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
  745. pool := newOpenAIWSConnPool(cfg)
  746. dialer := &openAIWSCountingDialer{}
  747. pool.setClientDialerForTest(dialer)
  748. account := &Account{ID: 130, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  749. baseReq := openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"}
  750. plainLease, err := pool.Acquire(context.Background(), baseReq)
  751. require.NoError(t, err)
  752. plainConnID := plainLease.ConnID()
  753. type acquireResult struct {
  754. lease *openAIWSConnLease
  755. err error
  756. }
  757. resultCh := make(chan acquireResult, 1)
  758. var done atomic.Bool
  759. ctx, cancel := context.WithTimeout(context.Background(), time.Second)
  760. defer cancel()
  761. go func() {
  762. betaReq := baseReq
  763. betaReq.Headers = http.Header{"X-Codex-Beta-Features": {"remote_compaction_v2"}}
  764. lease, acquireErr := pool.Acquire(ctx, betaReq)
  765. resultCh <- acquireResult{lease: lease, err: acquireErr}
  766. done.Store(true)
  767. }()
  768. require.Never(t, done.Load, 50*time.Millisecond, 5*time.Millisecond)
  769. plainLease.Release()
  770. result := <-resultCh
  771. require.NoError(t, result.err)
  772. require.NotNil(t, result.lease)
  773. require.False(t, result.lease.Reused())
  774. require.NotEqual(t, plainConnID, result.lease.ConnID())
  775. result.lease.Release()
  776. require.Equal(t, 2, dialer.DialCount())
  777. }
  778. func TestOpenAIWSConnPool_AcquireReplacesIncompatibleIdleWhenMatchingBusy(t *testing.T) {
  779. cfg := &config.Config{}
  780. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
  781. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  782. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
  783. pool := newOpenAIWSConnPool(cfg)
  784. dialer := &openAIWSCountingDialer{}
  785. pool.setClientDialerForTest(dialer)
  786. account := &Account{ID: 131, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  787. baseReq := openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"}
  788. plainLease, err := pool.Acquire(context.Background(), baseReq)
  789. require.NoError(t, err)
  790. plainConnID := plainLease.ConnID()
  791. plainLease.Release()
  792. betaReq := baseReq
  793. betaReq.Headers = http.Header{"X-Codex-Beta-Features": {"remote_compaction_v2"}}
  794. busyBetaLease, err := pool.Acquire(context.Background(), betaReq)
  795. require.NoError(t, err)
  796. secondBetaLease, err := pool.Acquire(context.Background(), betaReq)
  797. require.NoError(t, err)
  798. require.False(t, secondBetaLease.Reused())
  799. require.NotEqual(t, plainConnID, secondBetaLease.ConnID())
  800. require.NotEqual(t, busyBetaLease.ConnID(), secondBetaLease.ConnID())
  801. secondBetaLease.Release()
  802. busyBetaLease.Release()
  803. require.Equal(t, 3, dialer.DialCount())
  804. }
  805. func TestOpenAIWSConnPool_AcquireForcePreferredConnUnavailable(t *testing.T) {
  806. cfg := &config.Config{}
  807. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
  808. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  809. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
  810. pool := newOpenAIWSConnPool(cfg)
  811. account := &Account{ID: 124, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  812. ap := pool.getOrCreateAccountPool(account.ID)
  813. otherConn := newOpenAIWSConn("other_conn", account.ID, &openAIWSFakeConn{}, nil)
  814. ap.mu.Lock()
  815. ap.conns[otherConn.id] = otherConn
  816. ap.mu.Unlock()
  817. _, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
  818. Account: account,
  819. WSURL: "wss://example.com/v1/responses",
  820. ForcePreferredConn: true,
  821. })
  822. require.ErrorIs(t, err, errOpenAIWSPreferredConnUnavailable)
  823. _, err = pool.Acquire(context.Background(), openAIWSAcquireRequest{
  824. Account: account,
  825. WSURL: "wss://example.com/v1/responses",
  826. PreferredConnID: "missing_conn",
  827. ForcePreferredConn: true,
  828. })
  829. require.ErrorIs(t, err, errOpenAIWSPreferredConnUnavailable)
  830. }
  831. func TestOpenAIWSConnPool_AcquireForcePreferredConnQueuesOnPreferredOnly(t *testing.T) {
  832. cfg := &config.Config{}
  833. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
  834. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  835. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
  836. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 4
  837. pool := newOpenAIWSConnPool(cfg)
  838. account := &Account{ID: 125, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  839. ap := pool.getOrCreateAccountPool(account.ID)
  840. preferredConn := newOpenAIWSConn("preferred_conn", account.ID, &openAIWSFakeConn{}, nil)
  841. otherConn := newOpenAIWSConn("other_conn_idle", account.ID, &openAIWSFakeConn{}, nil)
  842. require.True(t, preferredConn.tryAcquire(), "先占用 preferred 连接,触发排队获取")
  843. ap.mu.Lock()
  844. ap.conns[preferredConn.id] = preferredConn
  845. ap.conns[otherConn.id] = otherConn
  846. ap.lastCleanupAt = time.Now()
  847. ap.mu.Unlock()
  848. go func() {
  849. time.Sleep(60 * time.Millisecond)
  850. preferredConn.release()
  851. }()
  852. ctx, cancel := context.WithTimeout(context.Background(), time.Second)
  853. defer cancel()
  854. lease, err := pool.Acquire(ctx, openAIWSAcquireRequest{
  855. Account: account,
  856. WSURL: "wss://example.com/v1/responses",
  857. PreferredConnID: preferredConn.id,
  858. ForcePreferredConn: true,
  859. })
  860. require.NoError(t, err)
  861. require.NotNil(t, lease)
  862. require.Equal(t, preferredConn.id, lease.ConnID(), "严格模式应只等待并复用 preferred 连接,不可漂移")
  863. require.GreaterOrEqual(t, lease.QueueWaitDuration(), 40*time.Millisecond)
  864. lease.Release()
  865. require.True(t, otherConn.tryAcquire(), "other 连接不应被严格模式抢占")
  866. otherConn.release()
  867. }
  868. func TestOpenAIWSConnPool_AcquireForcePreferredConnDirectAndQueueFull(t *testing.T) {
  869. cfg := &config.Config{}
  870. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
  871. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  872. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
  873. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 1
  874. pool := newOpenAIWSConnPool(cfg)
  875. account := &Account{ID: 127, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  876. ap := pool.getOrCreateAccountPool(account.ID)
  877. preferredConn := newOpenAIWSConn("preferred_conn_direct", account.ID, &openAIWSFakeConn{}, nil)
  878. otherConn := newOpenAIWSConn("other_conn_direct", account.ID, &openAIWSFakeConn{}, nil)
  879. ap.mu.Lock()
  880. ap.conns[preferredConn.id] = preferredConn
  881. ap.conns[otherConn.id] = otherConn
  882. ap.lastCleanupAt = time.Now()
  883. ap.mu.Unlock()
  884. lease, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
  885. Account: account,
  886. WSURL: "wss://example.com/v1/responses",
  887. PreferredConnID: preferredConn.id,
  888. ForcePreferredConn: true,
  889. })
  890. require.NoError(t, err)
  891. require.Equal(t, preferredConn.id, lease.ConnID(), "preferred 空闲时应直接命中")
  892. lease.Release()
  893. require.True(t, preferredConn.tryAcquire())
  894. preferredConn.waiters.Store(1)
  895. _, err = pool.Acquire(context.Background(), openAIWSAcquireRequest{
  896. Account: account,
  897. WSURL: "wss://example.com/v1/responses",
  898. PreferredConnID: preferredConn.id,
  899. ForcePreferredConn: true,
  900. })
  901. require.ErrorIs(t, err, errOpenAIWSConnQueueFull, "严格模式下队列满应直接失败,不得漂移")
  902. preferredConn.waiters.Store(0)
  903. preferredConn.release()
  904. }
  905. func TestOpenAIWSConnPool_CleanupSkipsPinnedConn(t *testing.T) {
  906. cfg := &config.Config{}
  907. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
  908. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 0
  909. pool := newOpenAIWSConnPool(cfg)
  910. accountID := int64(126)
  911. ap := pool.getOrCreateAccountPool(accountID)
  912. pinnedConn := newOpenAIWSConn("pinned_conn", accountID, &openAIWSFakeConn{}, nil)
  913. idleConn := newOpenAIWSConn("idle_conn", accountID, &openAIWSFakeConn{}, nil)
  914. ap.mu.Lock()
  915. ap.conns[pinnedConn.id] = pinnedConn
  916. ap.conns[idleConn.id] = idleConn
  917. ap.mu.Unlock()
  918. require.True(t, pool.PinConn(accountID, pinnedConn.id))
  919. evicted := pool.cleanupAccountLocked(ap, time.Now(), pool.maxConnsHardCap())
  920. closeOpenAIWSConns(evicted)
  921. ap.mu.Lock()
  922. _, pinnedExists := ap.conns[pinnedConn.id]
  923. _, idleExists := ap.conns[idleConn.id]
  924. ap.mu.Unlock()
  925. require.True(t, pinnedExists, "被 active ingress 绑定的连接不应被 cleanup 回收")
  926. require.False(t, idleExists, "非绑定的空闲连接应被回收")
  927. pool.UnpinConn(accountID, pinnedConn.id)
  928. evicted = pool.cleanupAccountLocked(ap, time.Now(), pool.maxConnsHardCap())
  929. closeOpenAIWSConns(evicted)
  930. ap.mu.Lock()
  931. _, pinnedExists = ap.conns[pinnedConn.id]
  932. ap.mu.Unlock()
  933. require.False(t, pinnedExists, "解绑后连接应可被正常回收")
  934. }
  935. func TestOpenAIWSConnPool_PinUnpinConnBranches(t *testing.T) {
  936. var nilPool *openAIWSConnPool
  937. require.False(t, nilPool.PinConn(1, "x"))
  938. nilPool.UnpinConn(1, "x")
  939. cfg := &config.Config{}
  940. pool := newOpenAIWSConnPool(cfg)
  941. accountID := int64(128)
  942. ap := &openAIWSAccountPool{
  943. conns: map[string]*openAIWSConn{},
  944. }
  945. pool.accounts.Store(accountID, ap)
  946. require.False(t, pool.PinConn(0, "x"))
  947. require.False(t, pool.PinConn(999, "x"))
  948. require.False(t, pool.PinConn(accountID, ""))
  949. require.False(t, pool.PinConn(accountID, "missing"))
  950. conn := newOpenAIWSConn("pin_refcount", accountID, &openAIWSFakeConn{}, nil)
  951. ap.mu.Lock()
  952. ap.conns[conn.id] = conn
  953. ap.mu.Unlock()
  954. require.True(t, pool.PinConn(accountID, conn.id))
  955. require.True(t, pool.PinConn(accountID, conn.id))
  956. ap.mu.Lock()
  957. require.Equal(t, 2, ap.pinnedConns[conn.id])
  958. ap.mu.Unlock()
  959. pool.UnpinConn(accountID, conn.id)
  960. ap.mu.Lock()
  961. require.Equal(t, 1, ap.pinnedConns[conn.id])
  962. ap.mu.Unlock()
  963. pool.UnpinConn(accountID, conn.id)
  964. ap.mu.Lock()
  965. _, exists := ap.pinnedConns[conn.id]
  966. ap.mu.Unlock()
  967. require.False(t, exists)
  968. pool.UnpinConn(accountID, conn.id)
  969. pool.UnpinConn(accountID, "")
  970. pool.UnpinConn(0, conn.id)
  971. pool.UnpinConn(999, conn.id)
  972. }
  973. func TestOpenAIWSConnPool_EffectiveMaxConnsByAccount(t *testing.T) {
  974. cfg := &config.Config{}
  975. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 8
  976. cfg.Gateway.OpenAIWS.DynamicMaxConnsByAccountConcurrencyEnabled = true
  977. cfg.Gateway.OpenAIWS.OAuthMaxConnsFactor = 1.0
  978. cfg.Gateway.OpenAIWS.APIKeyMaxConnsFactor = 0.6
  979. pool := newOpenAIWSConnPool(cfg)
  980. oauthHigh := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 10}
  981. require.Equal(t, 8, pool.effectiveMaxConnsByAccount(oauthHigh), "应受全局硬上限约束")
  982. oauthLow := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 3}
  983. require.Equal(t, 3, pool.effectiveMaxConnsByAccount(oauthLow))
  984. apiKeyHigh := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 10}
  985. require.Equal(t, 6, pool.effectiveMaxConnsByAccount(apiKeyHigh), "API Key 应按系数缩放")
  986. apiKeyLow := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1}
  987. require.Equal(t, 1, pool.effectiveMaxConnsByAccount(apiKeyLow), "最小值应保持为 1")
  988. unlimited := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 0}
  989. require.Equal(t, 8, pool.effectiveMaxConnsByAccount(unlimited), "无限并发应回退到全局硬上限")
  990. require.Equal(t, 8, pool.effectiveMaxConnsByAccount(nil), "缺少账号上下文应回退到全局硬上限")
  991. }
  992. func TestOpenAIWSConnPool_EffectiveMaxConnsDisabledFallbackHardCap(t *testing.T) {
  993. cfg := &config.Config{}
  994. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 8
  995. cfg.Gateway.OpenAIWS.DynamicMaxConnsByAccountConcurrencyEnabled = false
  996. cfg.Gateway.OpenAIWS.OAuthMaxConnsFactor = 1.0
  997. cfg.Gateway.OpenAIWS.APIKeyMaxConnsFactor = 1.0
  998. pool := newOpenAIWSConnPool(cfg)
  999. account := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 2}
  1000. require.Equal(t, 8, pool.effectiveMaxConnsByAccount(account), "关闭动态模式后应保持旧行为")
  1001. }
  1002. func TestOpenAIWSConnPool_EffectiveMaxConnsByAccount_ModeRouterV2RespectsHardCap(t *testing.T) {
  1003. cfg := &config.Config{}
  1004. cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
  1005. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 8
  1006. cfg.Gateway.OpenAIWS.DynamicMaxConnsByAccountConcurrencyEnabled = true
  1007. cfg.Gateway.OpenAIWS.OAuthMaxConnsFactor = 0.3
  1008. cfg.Gateway.OpenAIWS.APIKeyMaxConnsFactor = 0.6
  1009. pool := newOpenAIWSConnPool(cfg)
  1010. high := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 20}
  1011. require.Equal(t, 8, pool.effectiveMaxConnsByAccount(high), "v2 路径也必须受连接池硬上限约束")
  1012. nonPositive := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 0}
  1013. require.Equal(t, 0, pool.effectiveMaxConnsByAccount(nonPositive), "并发数<=0 时应不可调度")
  1014. }
  1015. func TestOpenAIWSConnPool_AcquireRejectsWhenEffectiveMaxConnsIsZero(t *testing.T) {
  1016. cfg := &config.Config{}
  1017. cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
  1018. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 8
  1019. pool := newOpenAIWSConnPool(cfg)
  1020. account := &Account{ID: 901, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 0}
  1021. _, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
  1022. Account: account,
  1023. WSURL: "wss://example.com/v1/responses",
  1024. })
  1025. require.ErrorIs(t, err, errOpenAIWSConnQueueFull)
  1026. }
  1027. func TestOpenAIWSConnLease_ReadMessageWithContextTimeout_PerRead(t *testing.T) {
  1028. conn := newOpenAIWSConn("timeout", 1, &openAIWSBlockingConn{readDelay: 80 * time.Millisecond}, nil)
  1029. lease := &openAIWSConnLease{conn: conn}
  1030. _, err := lease.ReadMessageWithContextTimeout(context.Background(), 20*time.Millisecond)
  1031. require.Error(t, err)
  1032. require.ErrorIs(t, err, context.DeadlineExceeded)
  1033. payload, err := lease.ReadMessageWithContextTimeout(context.Background(), 150*time.Millisecond)
  1034. require.NoError(t, err)
  1035. require.Contains(t, string(payload), "response.completed")
  1036. parentCtx, cancel := context.WithCancel(context.Background())
  1037. cancel()
  1038. _, err = lease.ReadMessageWithContextTimeout(parentCtx, 150*time.Millisecond)
  1039. require.Error(t, err)
  1040. require.ErrorIs(t, err, context.Canceled)
  1041. }
  1042. func TestOpenAIWSConnLease_WriteJSONWithContextTimeout_RespectsParentContext(t *testing.T) {
  1043. conn := newOpenAIWSConn("write_timeout_ctx", 1, &openAIWSWriteBlockingConn{}, nil)
  1044. lease := &openAIWSConnLease{conn: conn}
  1045. parentCtx, cancel := context.WithCancel(context.Background())
  1046. go func() {
  1047. time.Sleep(20 * time.Millisecond)
  1048. cancel()
  1049. }()
  1050. start := time.Now()
  1051. err := lease.WriteJSONWithContextTimeout(parentCtx, map[string]any{"type": "response.create"}, 2*time.Minute)
  1052. elapsed := time.Since(start)
  1053. require.Error(t, err)
  1054. require.ErrorIs(t, err, context.Canceled)
  1055. require.Less(t, elapsed, 200*time.Millisecond)
  1056. }
  1057. func TestOpenAIWSConnLease_PingWithTimeout(t *testing.T) {
  1058. conn := newOpenAIWSConn("ping_ok", 1, &openAIWSFakeConn{}, nil)
  1059. lease := &openAIWSConnLease{conn: conn}
  1060. require.NoError(t, lease.PingWithTimeout(50*time.Millisecond))
  1061. var nilLease *openAIWSConnLease
  1062. err := nilLease.PingWithTimeout(50 * time.Millisecond)
  1063. require.ErrorIs(t, err, errOpenAIWSConnClosed)
  1064. }
  1065. func TestOpenAIWSConn_ReadAndWriteCanProceedConcurrently(t *testing.T) {
  1066. conn := newOpenAIWSConn("full_duplex", 1, &openAIWSBlockingConn{readDelay: 120 * time.Millisecond}, nil)
  1067. readDone := make(chan error, 1)
  1068. go func() {
  1069. _, err := conn.readMessageWithContextTimeout(context.Background(), 200*time.Millisecond)
  1070. readDone <- err
  1071. }()
  1072. // 让读取先占用 readMu。
  1073. time.Sleep(20 * time.Millisecond)
  1074. start := time.Now()
  1075. err := conn.pingWithTimeout(50 * time.Millisecond)
  1076. elapsed := time.Since(start)
  1077. require.NoError(t, err)
  1078. require.Less(t, elapsed, 80*time.Millisecond, "写路径不应被读锁长期阻塞")
  1079. require.NoError(t, <-readDone)
  1080. }
  1081. func TestOpenAIWSConnPool_BackgroundPingSweep_EvictsDeadIdleConn(t *testing.T) {
  1082. cfg := &config.Config{}
  1083. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
  1084. pool := newOpenAIWSConnPool(cfg)
  1085. accountID := int64(301)
  1086. ap := pool.getOrCreateAccountPool(accountID)
  1087. conn := newOpenAIWSConn("dead_idle", accountID, &openAIWSPingFailConn{}, nil)
  1088. ap.mu.Lock()
  1089. ap.conns[conn.id] = conn
  1090. ap.mu.Unlock()
  1091. pool.runBackgroundPingSweep()
  1092. ap.mu.Lock()
  1093. _, exists := ap.conns[conn.id]
  1094. ap.mu.Unlock()
  1095. require.False(t, exists, "后台 ping 失败的空闲连接应被回收")
  1096. }
  1097. func TestOpenAIWSConnPool_BackgroundCleanupSweep_WithoutAcquire(t *testing.T) {
  1098. cfg := &config.Config{}
  1099. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
  1100. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
  1101. pool := newOpenAIWSConnPool(cfg)
  1102. accountID := int64(302)
  1103. ap := pool.getOrCreateAccountPool(accountID)
  1104. stale := newOpenAIWSConn("stale_bg", accountID, &openAIWSFakeConn{}, nil)
  1105. stale.createdAtNano.Store(time.Now().Add(-2 * time.Hour).UnixNano())
  1106. stale.lastUsedNano.Store(time.Now().Add(-2 * time.Hour).UnixNano())
  1107. ap.mu.Lock()
  1108. ap.conns[stale.id] = stale
  1109. ap.mu.Unlock()
  1110. pool.runBackgroundCleanupSweep(time.Now())
  1111. ap.mu.Lock()
  1112. _, exists := ap.conns[stale.id]
  1113. ap.mu.Unlock()
  1114. require.False(t, exists, "后台清理应在无新 acquire 时也回收过期连接")
  1115. }
  1116. func TestOpenAIWSConnPool_RecyclesUnsupportedIdlePingConnection(t *testing.T) {
  1117. cfg := &config.Config{}
  1118. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
  1119. cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
  1120. pool := &openAIWSConnPool{cfg: cfg}
  1121. accountID := int64(303)
  1122. ap := &openAIWSAccountPool{conns: make(map[string]*openAIWSConn)}
  1123. conn := newOpenAIWSConn("stale_unsupported_idle_ping", accountID, &openAIWSIdlePingUnsupportedConn{}, nil)
  1124. conn.lastUsedNano.Store(time.Now().Add(-openAIWSConnIdleRecycleAfter - time.Second).UnixNano())
  1125. ap.conns[conn.id] = conn
  1126. pool.accounts.Store(accountID, ap)
  1127. pool.runBackgroundCleanupSweep(time.Now())
  1128. ap.mu.Lock()
  1129. _, exists := ap.conns[conn.id]
  1130. ap.mu.Unlock()
  1131. require.False(t, exists, "不支持无 reader idle ping 的陈旧连接应被主动回收")
  1132. select {
  1133. case <-conn.closedCh:
  1134. default:
  1135. t.Fatal("被回收的陈旧连接应已关闭")
  1136. }
  1137. }
  1138. func TestOpenAIWSConnPool_BackgroundWorkerGuardBranches(t *testing.T) {
  1139. var nilPool *openAIWSConnPool
  1140. require.NotPanics(t, func() {
  1141. nilPool.startBackgroundWorkers()
  1142. nilPool.runBackgroundPingWorker()
  1143. nilPool.runBackgroundPingSweep()
  1144. _ = nilPool.snapshotIdleConnsForPing()
  1145. nilPool.runBackgroundCleanupWorker()
  1146. nilPool.runBackgroundCleanupSweep(time.Now())
  1147. })
  1148. poolNoStop := &openAIWSConnPool{}
  1149. require.NotPanics(t, func() {
  1150. poolNoStop.startBackgroundWorkers()
  1151. })
  1152. poolStopPing := &openAIWSConnPool{workerStopCh: make(chan struct{})}
  1153. pingDone := make(chan struct{})
  1154. go func() {
  1155. poolStopPing.runBackgroundPingWorker()
  1156. close(pingDone)
  1157. }()
  1158. close(poolStopPing.workerStopCh)
  1159. select {
  1160. case <-pingDone:
  1161. case <-time.After(500 * time.Millisecond):
  1162. t.Fatal("runBackgroundPingWorker 未在 stop 信号后退出")
  1163. }
  1164. poolStopCleanup := &openAIWSConnPool{workerStopCh: make(chan struct{})}
  1165. cleanupDone := make(chan struct{})
  1166. go func() {
  1167. poolStopCleanup.runBackgroundCleanupWorker()
  1168. close(cleanupDone)
  1169. }()
  1170. close(poolStopCleanup.workerStopCh)
  1171. select {
  1172. case <-cleanupDone:
  1173. case <-time.After(500 * time.Millisecond):
  1174. t.Fatal("runBackgroundCleanupWorker 未在 stop 信号后退出")
  1175. }
  1176. }
  1177. func TestOpenAIWSConnPool_SnapshotIdleConnsForPing_SkipsInvalidEntries(t *testing.T) {
  1178. pool := &openAIWSConnPool{}
  1179. pool.accounts.Store("invalid-key", &openAIWSAccountPool{})
  1180. pool.accounts.Store(int64(123), "invalid-value")
  1181. accountID := int64(123)
  1182. ap := &openAIWSAccountPool{
  1183. conns: make(map[string]*openAIWSConn),
  1184. }
  1185. ap.conns["nil_conn"] = nil
  1186. leased := newOpenAIWSConn("leased", accountID, &openAIWSFakeConn{}, nil)
  1187. require.True(t, leased.tryAcquire())
  1188. ap.conns[leased.id] = leased
  1189. waiting := newOpenAIWSConn("waiting", accountID, &openAIWSFakeConn{}, nil)
  1190. waiting.waiters.Store(1)
  1191. ap.conns[waiting.id] = waiting
  1192. idle := newOpenAIWSConn("idle", accountID, &openAIWSFakeConn{}, nil)
  1193. ap.conns[idle.id] = idle
  1194. pool.accounts.Store(accountID, ap)
  1195. candidates := pool.snapshotIdleConnsForPing()
  1196. require.Len(t, candidates, 1)
  1197. require.Equal(t, idle.id, candidates[0].conn.id)
  1198. }
  1199. func TestOpenAIWSConnPool_RunBackgroundCleanupSweep_SkipsInvalidAndUsesAccountCap(t *testing.T) {
  1200. cfg := &config.Config{}
  1201. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 4
  1202. cfg.Gateway.OpenAIWS.DynamicMaxConnsByAccountConcurrencyEnabled = true
  1203. pool := &openAIWSConnPool{cfg: cfg}
  1204. pool.accounts.Store("bad-key", "bad-value")
  1205. accountID := int64(2026)
  1206. ap := &openAIWSAccountPool{
  1207. conns: make(map[string]*openAIWSConn),
  1208. }
  1209. ap.conns["nil_conn"] = nil
  1210. stale := newOpenAIWSConn("stale_bg_cleanup", accountID, &openAIWSFakeConn{}, nil)
  1211. stale.createdAtNano.Store(time.Now().Add(-2 * time.Hour).UnixNano())
  1212. stale.lastUsedNano.Store(time.Now().Add(-2 * time.Hour).UnixNano())
  1213. ap.conns[stale.id] = stale
  1214. ap.lastAcquire = &openAIWSAcquireRequest{
  1215. Account: &Account{
  1216. ID: accountID,
  1217. Platform: PlatformOpenAI,
  1218. Type: AccountTypeAPIKey,
  1219. Concurrency: 1,
  1220. },
  1221. }
  1222. pool.accounts.Store(accountID, ap)
  1223. now := time.Now()
  1224. require.NotPanics(t, func() {
  1225. pool.runBackgroundCleanupSweep(now)
  1226. })
  1227. ap.mu.Lock()
  1228. _, nilConnExists := ap.conns["nil_conn"]
  1229. _, exists := ap.conns[stale.id]
  1230. lastCleanupAt := ap.lastCleanupAt
  1231. ap.mu.Unlock()
  1232. require.False(t, nilConnExists, "后台清理应移除无效 nil 连接条目")
  1233. require.False(t, exists, "后台清理应清理过期连接")
  1234. require.Equal(t, now, lastCleanupAt)
  1235. }
  1236. func TestOpenAIWSConnPool_QueueLimitPerConn_DefaultAndConfigured(t *testing.T) {
  1237. var nilPool *openAIWSConnPool
  1238. require.Equal(t, 256, nilPool.queueLimitPerConn())
  1239. pool := &openAIWSConnPool{cfg: &config.Config{}}
  1240. require.Equal(t, 256, pool.queueLimitPerConn())
  1241. pool.cfg.Gateway.OpenAIWS.QueueLimitPerConn = 9
  1242. require.Equal(t, 9, pool.queueLimitPerConn())
  1243. }
  1244. func TestOpenAIWSConnPool_Close(t *testing.T) {
  1245. cfg := &config.Config{}
  1246. pool := newOpenAIWSConnPool(cfg)
  1247. // Close 应该可以安全调用
  1248. pool.Close()
  1249. // workerStopCh 应已关闭
  1250. select {
  1251. case <-pool.workerStopCh:
  1252. // 预期:channel 已关闭
  1253. default:
  1254. t.Fatal("Close 后 workerStopCh 应已关闭")
  1255. }
  1256. // 多次调用 Close 不应 panic
  1257. pool.Close()
  1258. // nil pool 调用 Close 不应 panic
  1259. var nilPool *openAIWSConnPool
  1260. nilPool.Close()
  1261. }
  1262. func TestOpenAIWSDialError_ErrorAndUnwrap(t *testing.T) {
  1263. baseErr := errors.New("boom")
  1264. dialErr := &openAIWSDialError{StatusCode: 502, Err: baseErr}
  1265. require.Contains(t, dialErr.Error(), "status=502")
  1266. require.ErrorIs(t, dialErr.Unwrap(), baseErr)
  1267. noStatus := &openAIWSDialError{Err: baseErr}
  1268. require.Contains(t, noStatus.Error(), "boom")
  1269. var nilDialErr *openAIWSDialError
  1270. require.Equal(t, "", nilDialErr.Error())
  1271. require.NoError(t, nilDialErr.Unwrap())
  1272. }
  1273. func TestOpenAIWSConnLease_ReadWriteHelpersAndConnStats(t *testing.T) {
  1274. conn := newOpenAIWSConn("helper_conn", 1, &openAIWSFakeConn{}, http.Header{
  1275. "X-Test": []string{" value "},
  1276. })
  1277. lease := &openAIWSConnLease{conn: conn}
  1278. require.NoError(t, lease.WriteJSONContext(context.Background(), map[string]any{"type": "response.create"}))
  1279. payload, err := lease.ReadMessage(100 * time.Millisecond)
  1280. require.NoError(t, err)
  1281. require.Contains(t, string(payload), "response.completed")
  1282. payload, err = lease.ReadMessageContext(context.Background())
  1283. require.NoError(t, err)
  1284. require.Contains(t, string(payload), "response.completed")
  1285. payload, err = conn.readMessageWithTimeout(100 * time.Millisecond)
  1286. require.NoError(t, err)
  1287. require.Contains(t, string(payload), "response.completed")
  1288. require.Equal(t, "value", conn.handshakeHeader(" X-Test "))
  1289. require.NotZero(t, conn.createdAt())
  1290. require.NotZero(t, conn.lastUsedAt())
  1291. require.GreaterOrEqual(t, conn.age(time.Now()), time.Duration(0))
  1292. require.GreaterOrEqual(t, conn.idleDuration(time.Now()), time.Duration(0))
  1293. require.False(t, conn.isLeased())
  1294. // 覆盖空上下文路径
  1295. _, err = conn.readMessage(context.Background())
  1296. require.NoError(t, err)
  1297. // 覆盖 nil 保护分支
  1298. var nilConn *openAIWSConn
  1299. require.ErrorIs(t, nilConn.writeJSONWithTimeout(context.Background(), map[string]any{}, time.Second), errOpenAIWSConnClosed)
  1300. _, err = nilConn.readMessageWithTimeout(10 * time.Millisecond)
  1301. require.ErrorIs(t, err, errOpenAIWSConnClosed)
  1302. _, err = nilConn.readMessageWithContextTimeout(context.Background(), 10*time.Millisecond)
  1303. require.ErrorIs(t, err, errOpenAIWSConnClosed)
  1304. }
  1305. func TestOpenAIWSConnPool_PickOldestIdleAndAccountPoolLoad(t *testing.T) {
  1306. pool := &openAIWSConnPool{}
  1307. accountID := int64(404)
  1308. ap := &openAIWSAccountPool{conns: map[string]*openAIWSConn{}}
  1309. idleOld := newOpenAIWSConn("idle_old", accountID, &openAIWSFakeConn{}, nil)
  1310. idleOld.lastUsedNano.Store(time.Now().Add(-10 * time.Minute).UnixNano())
  1311. idleNew := newOpenAIWSConn("idle_new", accountID, &openAIWSFakeConn{}, nil)
  1312. idleNew.lastUsedNano.Store(time.Now().Add(-1 * time.Minute).UnixNano())
  1313. leased := newOpenAIWSConn("leased", accountID, &openAIWSFakeConn{}, nil)
  1314. require.True(t, leased.tryAcquire())
  1315. leased.waiters.Store(2)
  1316. ap.conns[idleOld.id] = idleOld
  1317. ap.conns[idleNew.id] = idleNew
  1318. ap.conns[leased.id] = leased
  1319. oldest := pool.pickOldestIdleConnLocked(ap)
  1320. require.NotNil(t, oldest)
  1321. require.Equal(t, idleOld.id, oldest.id)
  1322. inflight, waiters := accountPoolLoadLocked(ap)
  1323. require.Equal(t, 1, inflight)
  1324. require.Equal(t, 2, waiters)
  1325. pool.accounts.Store(accountID, ap)
  1326. loadInflight, loadWaiters, conns := pool.AccountPoolLoad(accountID)
  1327. require.Equal(t, 1, loadInflight)
  1328. require.Equal(t, 2, loadWaiters)
  1329. require.Equal(t, 3, conns)
  1330. zeroInflight, zeroWaiters, zeroConns := pool.AccountPoolLoad(0)
  1331. require.Equal(t, 0, zeroInflight)
  1332. require.Equal(t, 0, zeroWaiters)
  1333. require.Equal(t, 0, zeroConns)
  1334. }
  1335. func TestOpenAIWSConnPool_Close_WaitsWorkerGroupAndNilStopChannel(t *testing.T) {
  1336. pool := &openAIWSConnPool{}
  1337. release := make(chan struct{})
  1338. pool.workerWg.Add(1)
  1339. go func() {
  1340. defer pool.workerWg.Done()
  1341. <-release
  1342. }()
  1343. closed := make(chan struct{})
  1344. go func() {
  1345. pool.Close()
  1346. close(closed)
  1347. }()
  1348. select {
  1349. case <-closed:
  1350. t.Fatal("Close 不应在 WaitGroup 未完成时提前返回")
  1351. case <-time.After(30 * time.Millisecond):
  1352. }
  1353. close(release)
  1354. select {
  1355. case <-closed:
  1356. case <-time.After(time.Second):
  1357. t.Fatal("Close 未等待 workerWg 完成")
  1358. }
  1359. }
  1360. func TestOpenAIWSConnPool_Close_ClosesOnlyIdleConnections(t *testing.T) {
  1361. pool := &openAIWSConnPool{
  1362. workerStopCh: make(chan struct{}),
  1363. }
  1364. accountID := int64(606)
  1365. ap := &openAIWSAccountPool{
  1366. conns: map[string]*openAIWSConn{},
  1367. }
  1368. idle := newOpenAIWSConn("idle_conn", accountID, &openAIWSFakeConn{}, nil)
  1369. leased := newOpenAIWSConn("leased_conn", accountID, &openAIWSFakeConn{}, nil)
  1370. require.True(t, leased.tryAcquire())
  1371. ap.conns[idle.id] = idle
  1372. ap.conns[leased.id] = leased
  1373. pool.accounts.Store(accountID, ap)
  1374. pool.accounts.Store("invalid-key", "invalid-value")
  1375. pool.Close()
  1376. select {
  1377. case <-idle.closedCh:
  1378. // idle should be closed
  1379. default:
  1380. t.Fatal("空闲连接应在 Close 时被关闭")
  1381. }
  1382. select {
  1383. case <-leased.closedCh:
  1384. t.Fatal("已租赁连接不应在 Close 时被关闭")
  1385. default:
  1386. }
  1387. leased.release()
  1388. pool.Close()
  1389. }
  1390. func TestOpenAIWSConnPool_RunBackgroundPingSweep_ConcurrencyLimit(t *testing.T) {
  1391. cfg := &config.Config{}
  1392. pool := newOpenAIWSConnPool(cfg)
  1393. accountID := int64(505)
  1394. ap := pool.getOrCreateAccountPool(accountID)
  1395. var current atomic.Int32
  1396. var maxConcurrent atomic.Int32
  1397. release := make(chan struct{})
  1398. for i := 0; i < 25; i++ {
  1399. conn := newOpenAIWSConn(pool.nextConnID(accountID), accountID, &openAIWSPingBlockingConn{
  1400. current: &current,
  1401. maxConcurrent: &maxConcurrent,
  1402. release: release,
  1403. }, nil)
  1404. ap.mu.Lock()
  1405. ap.conns[conn.id] = conn
  1406. ap.mu.Unlock()
  1407. }
  1408. done := make(chan struct{})
  1409. go func() {
  1410. pool.runBackgroundPingSweep()
  1411. close(done)
  1412. }()
  1413. require.Eventually(t, func() bool {
  1414. return maxConcurrent.Load() >= 10
  1415. }, time.Second, 10*time.Millisecond)
  1416. close(release)
  1417. select {
  1418. case <-done:
  1419. case <-time.After(2 * time.Second):
  1420. t.Fatal("runBackgroundPingSweep 未在释放后完成")
  1421. }
  1422. require.LessOrEqual(t, maxConcurrent.Load(), int32(10))
  1423. }
  1424. func TestOpenAIWSConnLease_BasicGetterBranches(t *testing.T) {
  1425. var nilLease *openAIWSConnLease
  1426. require.Equal(t, "", nilLease.ConnID())
  1427. require.Equal(t, time.Duration(0), nilLease.QueueWaitDuration())
  1428. require.Equal(t, time.Duration(0), nilLease.ConnPickDuration())
  1429. require.False(t, nilLease.Reused())
  1430. require.Equal(t, "", nilLease.HandshakeHeader("x-test"))
  1431. require.False(t, nilLease.IsPrewarmed())
  1432. nilLease.MarkPrewarmed()
  1433. nilLease.Release()
  1434. conn := newOpenAIWSConn("getter_conn", 1, &openAIWSFakeConn{}, http.Header{"X-Test": []string{"ok"}})
  1435. lease := &openAIWSConnLease{
  1436. conn: conn,
  1437. queueWait: 3 * time.Millisecond,
  1438. connPick: 4 * time.Millisecond,
  1439. reused: true,
  1440. }
  1441. require.Equal(t, "getter_conn", lease.ConnID())
  1442. require.Equal(t, 3*time.Millisecond, lease.QueueWaitDuration())
  1443. require.Equal(t, 4*time.Millisecond, lease.ConnPickDuration())
  1444. require.True(t, lease.Reused())
  1445. require.Equal(t, "ok", lease.HandshakeHeader("x-test"))
  1446. require.False(t, lease.IsPrewarmed())
  1447. lease.MarkPrewarmed()
  1448. require.True(t, lease.IsPrewarmed())
  1449. lease.Release()
  1450. }
  1451. func TestOpenAIWSConnPool_UtilityBranches(t *testing.T) {
  1452. var nilPool *openAIWSConnPool
  1453. require.Equal(t, OpenAIWSPoolMetricsSnapshot{}, nilPool.SnapshotMetrics())
  1454. require.Equal(t, OpenAIWSTransportMetricsSnapshot{}, nilPool.SnapshotTransportMetrics())
  1455. pool := &openAIWSConnPool{cfg: &config.Config{}}
  1456. pool.metrics.acquireTotal.Store(7)
  1457. pool.metrics.acquireReuseTotal.Store(3)
  1458. metrics := pool.SnapshotMetrics()
  1459. require.Equal(t, int64(7), metrics.AcquireTotal)
  1460. require.Equal(t, int64(3), metrics.AcquireReuseTotal)
  1461. // 非 transport metrics dialer 路径
  1462. pool.clientDialer = &openAIWSFakeDialer{}
  1463. require.Equal(t, OpenAIWSTransportMetricsSnapshot{}, pool.SnapshotTransportMetrics())
  1464. pool.setClientDialerForTest(nil)
  1465. require.NotNil(t, pool.clientDialer)
  1466. require.Equal(t, 8, nilPool.maxConnsHardCap())
  1467. require.False(t, nilPool.dynamicMaxConnsEnabled())
  1468. require.Equal(t, 1.0, nilPool.maxConnsFactorByAccount(nil))
  1469. require.Equal(t, 0, nilPool.minIdlePerAccount())
  1470. require.Equal(t, 4, nilPool.maxIdlePerAccount())
  1471. require.Equal(t, 256, nilPool.queueLimitPerConn())
  1472. require.Equal(t, 0.7, nilPool.targetUtilization())
  1473. require.Equal(t, time.Duration(0), nilPool.prewarmCooldown())
  1474. require.Equal(t, 10*time.Second, nilPool.dialTimeout())
  1475. // shouldSuppressPrewarmLocked 覆盖 3 条分支
  1476. now := time.Now()
  1477. apNilFail := &openAIWSAccountPool{prewarmFails: 1}
  1478. require.False(t, pool.shouldSuppressPrewarmLocked(apNilFail, now))
  1479. apZeroTime := &openAIWSAccountPool{prewarmFails: 2}
  1480. require.False(t, pool.shouldSuppressPrewarmLocked(apZeroTime, now))
  1481. require.Equal(t, 0, apZeroTime.prewarmFails)
  1482. apOldFail := &openAIWSAccountPool{prewarmFails: 2, prewarmFailAt: now.Add(-openAIWSPrewarmFailureWindow - time.Second)}
  1483. require.False(t, pool.shouldSuppressPrewarmLocked(apOldFail, now))
  1484. apRecentFail := &openAIWSAccountPool{prewarmFails: openAIWSPrewarmFailureSuppress, prewarmFailAt: now}
  1485. require.True(t, pool.shouldSuppressPrewarmLocked(apRecentFail, now))
  1486. // recordConnPickDuration 的保护分支
  1487. nilPool.recordConnPickDuration(10 * time.Millisecond)
  1488. pool.recordConnPickDuration(-10 * time.Millisecond)
  1489. require.Equal(t, int64(1), pool.metrics.connPickTotal.Load())
  1490. // account pool 读写分支
  1491. require.Nil(t, nilPool.getOrCreateAccountPool(1))
  1492. require.Nil(t, pool.getOrCreateAccountPool(0))
  1493. pool.accounts.Store(int64(7), "invalid")
  1494. ap := pool.getOrCreateAccountPool(7)
  1495. require.NotNil(t, ap)
  1496. _, ok := pool.getAccountPool(0)
  1497. require.False(t, ok)
  1498. _, ok = pool.getAccountPool(12345)
  1499. require.False(t, ok)
  1500. pool.accounts.Store(int64(8), "bad-type")
  1501. _, ok = pool.getAccountPool(8)
  1502. require.False(t, ok)
  1503. // health check 条件
  1504. require.False(t, pool.shouldHealthCheckConn(nil))
  1505. conn := newOpenAIWSConn("health", 1, &openAIWSFakeConn{}, nil)
  1506. conn.lastUsedNano.Store(time.Now().Add(-openAIWSConnHealthCheckIdle - time.Second).UnixNano())
  1507. require.True(t, pool.shouldHealthCheckConn(conn))
  1508. unsafeConn := newOpenAIWSConn("unsafe_health", 1, &openAIWSIdlePingUnsupportedConn{}, nil)
  1509. unsafeConn.lastUsedNano.Store(time.Now().Add(-openAIWSConnHealthCheckIdle - time.Second).UnixNano())
  1510. require.False(t, pool.shouldHealthCheckConn(unsafeConn))
  1511. }
  1512. func TestOpenAIWSConn_LeaseAndTimeHelpers_NilAndClosedBranches(t *testing.T) {
  1513. var nilConn *openAIWSConn
  1514. nilConn.touch()
  1515. require.Equal(t, time.Time{}, nilConn.createdAt())
  1516. require.Equal(t, time.Time{}, nilConn.lastUsedAt())
  1517. require.Equal(t, time.Duration(0), nilConn.idleDuration(time.Now()))
  1518. require.Equal(t, time.Duration(0), nilConn.age(time.Now()))
  1519. require.False(t, nilConn.isLeased())
  1520. require.False(t, nilConn.isPrewarmed())
  1521. nilConn.markPrewarmed()
  1522. conn := newOpenAIWSConn("lease_state", 1, &openAIWSFakeConn{}, nil)
  1523. require.True(t, conn.tryAcquire())
  1524. require.True(t, conn.isLeased())
  1525. conn.release()
  1526. require.False(t, conn.isLeased())
  1527. conn.close()
  1528. require.False(t, conn.tryAcquire())
  1529. ctx, cancel := context.WithCancel(context.Background())
  1530. cancel()
  1531. err := conn.acquire(ctx)
  1532. require.Error(t, err)
  1533. }
  1534. func TestOpenAIWSConnLease_ReadWriteNilConnBranches(t *testing.T) {
  1535. lease := &openAIWSConnLease{}
  1536. require.ErrorIs(t, lease.WriteJSON(map[string]any{"k": "v"}, time.Second), errOpenAIWSConnClosed)
  1537. require.ErrorIs(t, lease.WriteJSONContext(context.Background(), map[string]any{"k": "v"}), errOpenAIWSConnClosed)
  1538. _, err := lease.ReadMessage(10 * time.Millisecond)
  1539. require.ErrorIs(t, err, errOpenAIWSConnClosed)
  1540. _, err = lease.ReadMessageContext(context.Background())
  1541. require.ErrorIs(t, err, errOpenAIWSConnClosed)
  1542. _, err = lease.ReadMessageWithContextTimeout(context.Background(), 10*time.Millisecond)
  1543. require.ErrorIs(t, err, errOpenAIWSConnClosed)
  1544. }
  1545. func TestOpenAIWSConnLease_ReleasedLeaseGuards(t *testing.T) {
  1546. conn := newOpenAIWSConn("released_guard", 1, &openAIWSFakeConn{}, nil)
  1547. lease := &openAIWSConnLease{conn: conn}
  1548. require.NoError(t, lease.PingWithTimeout(50*time.Millisecond))
  1549. lease.Release()
  1550. lease.Release() // idempotent
  1551. require.ErrorIs(t, lease.WriteJSON(map[string]any{"k": "v"}, time.Second), errOpenAIWSConnClosed)
  1552. require.ErrorIs(t, lease.WriteJSONContext(context.Background(), map[string]any{"k": "v"}), errOpenAIWSConnClosed)
  1553. require.ErrorIs(t, lease.WriteJSONWithContextTimeout(context.Background(), map[string]any{"k": "v"}, time.Second), errOpenAIWSConnClosed)
  1554. _, err := lease.ReadMessage(10 * time.Millisecond)
  1555. require.ErrorIs(t, err, errOpenAIWSConnClosed)
  1556. _, err = lease.ReadMessageContext(context.Background())
  1557. require.ErrorIs(t, err, errOpenAIWSConnClosed)
  1558. _, err = lease.ReadMessageWithContextTimeout(context.Background(), 10*time.Millisecond)
  1559. require.ErrorIs(t, err, errOpenAIWSConnClosed)
  1560. require.ErrorIs(t, lease.PingWithTimeout(50*time.Millisecond), errOpenAIWSConnClosed)
  1561. }
  1562. func TestOpenAIWSConnLease_MarkBrokenAfterRelease_NoEviction(t *testing.T) {
  1563. conn := newOpenAIWSConn("released_markbroken", 7, &openAIWSFakeConn{}, nil)
  1564. ap := &openAIWSAccountPool{
  1565. conns: map[string]*openAIWSConn{
  1566. conn.id: conn,
  1567. },
  1568. }
  1569. pool := &openAIWSConnPool{}
  1570. pool.accounts.Store(int64(7), ap)
  1571. lease := &openAIWSConnLease{
  1572. pool: pool,
  1573. accountID: 7,
  1574. conn: conn,
  1575. }
  1576. lease.Release()
  1577. lease.MarkBroken()
  1578. ap.mu.Lock()
  1579. _, exists := ap.conns[conn.id]
  1580. ap.mu.Unlock()
  1581. require.True(t, exists, "released lease should not evict active pool connection")
  1582. }
  1583. func TestOpenAIWSConn_AdditionalGuardBranches(t *testing.T) {
  1584. var nilConn *openAIWSConn
  1585. require.False(t, nilConn.tryAcquire())
  1586. require.ErrorIs(t, nilConn.acquire(context.Background()), errOpenAIWSConnClosed)
  1587. nilConn.release()
  1588. nilConn.close()
  1589. require.Equal(t, "", nilConn.handshakeHeader("x-test"))
  1590. connBusy := newOpenAIWSConn("busy_ctx", 1, &openAIWSFakeConn{}, nil)
  1591. require.True(t, connBusy.tryAcquire())
  1592. ctx, cancel := context.WithCancel(context.Background())
  1593. cancel()
  1594. require.ErrorIs(t, connBusy.acquire(ctx), context.Canceled)
  1595. connBusy.release()
  1596. connClosed := newOpenAIWSConn("closed_guard", 1, &openAIWSFakeConn{}, nil)
  1597. connClosed.close()
  1598. require.ErrorIs(
  1599. t,
  1600. connClosed.writeJSONWithTimeout(context.Background(), map[string]any{"k": "v"}, time.Second),
  1601. errOpenAIWSConnClosed,
  1602. )
  1603. _, err := connClosed.readMessageWithContextTimeout(context.Background(), time.Second)
  1604. require.ErrorIs(t, err, errOpenAIWSConnClosed)
  1605. require.ErrorIs(t, connClosed.pingWithTimeout(time.Second), errOpenAIWSConnClosed)
  1606. connNoWS := newOpenAIWSConn("no_ws", 1, nil, nil)
  1607. require.ErrorIs(t, connNoWS.writeJSON(map[string]any{"k": "v"}, context.Background()), errOpenAIWSConnClosed)
  1608. _, err = connNoWS.readMessage(context.Background())
  1609. require.ErrorIs(t, err, errOpenAIWSConnClosed)
  1610. require.ErrorIs(t, connNoWS.pingWithTimeout(time.Second), errOpenAIWSConnClosed)
  1611. require.Equal(t, "", connNoWS.handshakeHeader("x-test"))
  1612. connOK := newOpenAIWSConn("ok", 1, &openAIWSFakeConn{}, nil)
  1613. require.NoError(t, connOK.writeJSON(map[string]any{"k": "v"}, nil))
  1614. _, err = connOK.readMessageWithContextTimeout(context.Background(), 0)
  1615. require.NoError(t, err)
  1616. require.NoError(t, connOK.pingWithTimeout(0))
  1617. connZero := newOpenAIWSConn("zero_ts", 1, &openAIWSFakeConn{}, nil)
  1618. connZero.createdAtNano.Store(0)
  1619. connZero.lastUsedNano.Store(0)
  1620. require.True(t, connZero.createdAt().IsZero())
  1621. require.True(t, connZero.lastUsedAt().IsZero())
  1622. require.Equal(t, time.Duration(0), connZero.idleDuration(time.Now()))
  1623. require.Equal(t, time.Duration(0), connZero.age(time.Now()))
  1624. require.Nil(t, cloneOpenAIWSAcquireRequestPtr(nil))
  1625. copied := cloneHeader(http.Header{
  1626. "X-Empty": []string{},
  1627. "X-Test": []string{"v1"},
  1628. })
  1629. require.Contains(t, copied, "X-Empty")
  1630. require.Nil(t, copied["X-Empty"])
  1631. require.Equal(t, "v1", copied.Get("X-Test"))
  1632. closeOpenAIWSConns([]*openAIWSConn{nil, connOK})
  1633. }
  1634. func TestOpenAIWSConnPool_CanceledWaiterReturnsDeliveredLease(t *testing.T) {
  1635. conn := newOpenAIWSConn("cancelled_delivery", 1, &openAIWSFakeConn{}, nil)
  1636. // Both branches of acquire's select are ready. Before the post-delivery
  1637. // cancellation check this intermittently returned nil after consuming the
  1638. // only lease token, which made the next pool acquire block forever.
  1639. for range 64 {
  1640. ctx, cancel := context.WithCancel(context.Background())
  1641. cancel()
  1642. require.ErrorIs(t, conn.acquire(ctx), context.Canceled)
  1643. require.True(t, conn.tryAcquire(), "a canceled waiter must return a delivered lease token")
  1644. conn.release()
  1645. }
  1646. }
  1647. func TestOpenAIWSConnLease_MarkBrokenEvictsConn(t *testing.T) {
  1648. pool := newOpenAIWSConnPool(&config.Config{})
  1649. accountID := int64(5001)
  1650. conn := newOpenAIWSConn("broken_me", accountID, &openAIWSFakeConn{}, nil)
  1651. ap := pool.getOrCreateAccountPool(accountID)
  1652. ap.mu.Lock()
  1653. ap.conns[conn.id] = conn
  1654. ap.mu.Unlock()
  1655. lease := &openAIWSConnLease{
  1656. pool: pool,
  1657. accountID: accountID,
  1658. conn: conn,
  1659. }
  1660. lease.MarkBroken()
  1661. ap.mu.Lock()
  1662. _, exists := ap.conns[conn.id]
  1663. ap.mu.Unlock()
  1664. require.False(t, exists)
  1665. require.False(t, conn.tryAcquire(), "被标记为 broken 的连接应被关闭")
  1666. }
  1667. func TestOpenAIWSConnPool_TargetConnCountAndPrewarmBranches(t *testing.T) {
  1668. cfg := &config.Config{}
  1669. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  1670. pool := newOpenAIWSConnPool(cfg)
  1671. require.Equal(t, 0, pool.targetConnCountLocked(nil, 1))
  1672. ap := &openAIWSAccountPool{conns: map[string]*openAIWSConn{}}
  1673. require.Equal(t, 0, pool.targetConnCountLocked(ap, 0))
  1674. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 3
  1675. require.Equal(t, 1, pool.targetConnCountLocked(ap, 1), "minIdle 应被 maxConns 截断")
  1676. // 覆盖 waiters>0 且 target 需要至少 len(conns)+1 的分支
  1677. cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
  1678. cfg.Gateway.OpenAIWS.PoolTargetUtilization = 0.9
  1679. busy := newOpenAIWSConn("busy_target", 2, &openAIWSFakeConn{}, nil)
  1680. require.True(t, busy.tryAcquire())
  1681. busy.waiters.Store(1)
  1682. ap.conns[busy.id] = busy
  1683. target := pool.targetConnCountLocked(ap, 4)
  1684. require.GreaterOrEqual(t, target, len(ap.conns)+1)
  1685. // prewarm: account pool 缺失时,拨号后的连接应被关闭并提前返回
  1686. req := openAIWSAcquireRequest{
  1687. Account: &Account{ID: 999, Platform: PlatformOpenAI, Type: AccountTypeAPIKey},
  1688. WSURL: "wss://example.com/v1/responses",
  1689. }
  1690. pool.prewarmConns(999, req, 1)
  1691. // prewarm: 拨号失败分支(prewarmFails 累加)
  1692. accountID := int64(1000)
  1693. failPool := newOpenAIWSConnPool(cfg)
  1694. failPool.setClientDialerForTest(&openAIWSAlwaysFailDialer{})
  1695. apFail := failPool.getOrCreateAccountPool(accountID)
  1696. apFail.mu.Lock()
  1697. apFail.creating = 1
  1698. apFail.mu.Unlock()
  1699. req.Account.ID = accountID
  1700. failPool.prewarmConns(accountID, req, 1)
  1701. apFail.mu.Lock()
  1702. require.GreaterOrEqual(t, apFail.prewarmFails, 1)
  1703. apFail.mu.Unlock()
  1704. }
  1705. func TestOpenAIWSConnPool_Acquire_ErrorBranches(t *testing.T) {
  1706. var nilPool *openAIWSConnPool
  1707. _, err := nilPool.Acquire(context.Background(), openAIWSAcquireRequest{})
  1708. require.Error(t, err)
  1709. pool := newOpenAIWSConnPool(&config.Config{})
  1710. _, err = pool.Acquire(context.Background(), openAIWSAcquireRequest{
  1711. Account: &Account{ID: 1},
  1712. WSURL: " ",
  1713. })
  1714. require.Error(t, err)
  1715. require.Contains(t, err.Error(), "ws url is empty")
  1716. // target=nil 分支:池满且仅有 nil 连接
  1717. cfg := &config.Config{}
  1718. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
  1719. cfg.Gateway.OpenAIWS.QueueLimitPerConn = 1
  1720. fullPool := newOpenAIWSConnPool(cfg)
  1721. account := &Account{ID: 2001, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  1722. ap := fullPool.getOrCreateAccountPool(account.ID)
  1723. ap.mu.Lock()
  1724. ap.conns["nil"] = nil
  1725. ap.lastCleanupAt = time.Now()
  1726. ap.mu.Unlock()
  1727. _, err = fullPool.Acquire(context.Background(), openAIWSAcquireRequest{
  1728. Account: account,
  1729. WSURL: "wss://example.com/v1/responses",
  1730. })
  1731. require.ErrorIs(t, err, errOpenAIWSConnClosed)
  1732. // queue full 分支:waiters 达上限
  1733. account2 := &Account{ID: 2002, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  1734. ap2 := fullPool.getOrCreateAccountPool(account2.ID)
  1735. conn := newOpenAIWSConn("queue_full", account2.ID, &openAIWSFakeConn{}, nil)
  1736. require.True(t, conn.tryAcquire())
  1737. conn.waiters.Store(1)
  1738. ap2.mu.Lock()
  1739. ap2.conns[conn.id] = conn
  1740. ap2.lastCleanupAt = time.Now()
  1741. ap2.mu.Unlock()
  1742. _, err = fullPool.Acquire(context.Background(), openAIWSAcquireRequest{
  1743. Account: account2,
  1744. WSURL: "wss://example.com/v1/responses",
  1745. })
  1746. require.ErrorIs(t, err, errOpenAIWSConnQueueFull)
  1747. }
  1748. type openAIWSFakeDialer struct{}
  1749. func (d *openAIWSFakeDialer) Dial(
  1750. ctx context.Context,
  1751. wsURL string,
  1752. headers http.Header,
  1753. proxyURL string,
  1754. ) (openAIWSClientConn, int, http.Header, error) {
  1755. _ = ctx
  1756. _ = wsURL
  1757. _ = headers
  1758. _ = proxyURL
  1759. return &openAIWSFakeConn{}, 0, nil, nil
  1760. }
  1761. type openAIWSCountingDialer struct {
  1762. mu sync.Mutex
  1763. dialCount int
  1764. }
  1765. type openAIWSFirstDialBlockingCaptureDialer struct {
  1766. mu sync.Mutex
  1767. dialCount int
  1768. headers []http.Header
  1769. firstStarted chan struct{}
  1770. releaseFirst chan struct{}
  1771. }
  1772. func newOpenAIWSFirstDialBlockingCaptureDialer() *openAIWSFirstDialBlockingCaptureDialer {
  1773. return &openAIWSFirstDialBlockingCaptureDialer{
  1774. firstStarted: make(chan struct{}),
  1775. releaseFirst: make(chan struct{}),
  1776. }
  1777. }
  1778. type openAIWSAlwaysFailDialer struct {
  1779. mu sync.Mutex
  1780. dialCount int
  1781. }
  1782. type openAIWSPingBlockingConn struct {
  1783. current *atomic.Int32
  1784. maxConcurrent *atomic.Int32
  1785. release <-chan struct{}
  1786. }
  1787. type openAIWSIdlePingUnsupportedConn struct {
  1788. openAIWSFakeConn
  1789. }
  1790. func (c *openAIWSIdlePingUnsupportedConn) SupportsIdlePingWithoutReader() bool {
  1791. return false
  1792. }
  1793. func (c *openAIWSPingBlockingConn) WriteJSON(context.Context, any) error {
  1794. return nil
  1795. }
  1796. func (c *openAIWSPingBlockingConn) ReadMessage(context.Context) ([]byte, error) {
  1797. return []byte(`{"type":"response.completed","response":{"id":"resp_blocking_ping"}}`), nil
  1798. }
  1799. func (c *openAIWSPingBlockingConn) Ping(ctx context.Context) error {
  1800. if c.current == nil || c.maxConcurrent == nil {
  1801. return nil
  1802. }
  1803. now := c.current.Add(1)
  1804. for {
  1805. prev := c.maxConcurrent.Load()
  1806. if now <= prev || c.maxConcurrent.CompareAndSwap(prev, now) {
  1807. break
  1808. }
  1809. }
  1810. defer c.current.Add(-1)
  1811. select {
  1812. case <-ctx.Done():
  1813. return ctx.Err()
  1814. case <-c.release:
  1815. return nil
  1816. }
  1817. }
  1818. func (c *openAIWSPingBlockingConn) Close() error {
  1819. return nil
  1820. }
  1821. func (d *openAIWSCountingDialer) Dial(
  1822. ctx context.Context,
  1823. wsURL string,
  1824. headers http.Header,
  1825. proxyURL string,
  1826. ) (openAIWSClientConn, int, http.Header, error) {
  1827. _ = ctx
  1828. _ = wsURL
  1829. _ = headers
  1830. _ = proxyURL
  1831. d.mu.Lock()
  1832. d.dialCount++
  1833. d.mu.Unlock()
  1834. return &openAIWSFakeConn{}, 0, nil, nil
  1835. }
  1836. func (d *openAIWSFirstDialBlockingCaptureDialer) Dial(
  1837. ctx context.Context,
  1838. wsURL string,
  1839. headers http.Header,
  1840. proxyURL string,
  1841. ) (openAIWSClientConn, int, http.Header, error) {
  1842. _ = wsURL
  1843. _ = proxyURL
  1844. d.mu.Lock()
  1845. d.dialCount++
  1846. dialNumber := d.dialCount
  1847. d.headers = append(d.headers, cloneHeader(headers))
  1848. d.mu.Unlock()
  1849. if dialNumber == 1 {
  1850. close(d.firstStarted)
  1851. select {
  1852. case <-ctx.Done():
  1853. return nil, 0, nil, ctx.Err()
  1854. case <-d.releaseFirst:
  1855. }
  1856. }
  1857. return &openAIWSFakeConn{}, 0, nil, nil
  1858. }
  1859. func (d *openAIWSFirstDialBlockingCaptureDialer) DialCount() int {
  1860. d.mu.Lock()
  1861. defer d.mu.Unlock()
  1862. return d.dialCount
  1863. }
  1864. func (d *openAIWSCountingDialer) DialCount() int {
  1865. d.mu.Lock()
  1866. defer d.mu.Unlock()
  1867. return d.dialCount
  1868. }
  1869. func (d *openAIWSAlwaysFailDialer) Dial(
  1870. ctx context.Context,
  1871. wsURL string,
  1872. headers http.Header,
  1873. proxyURL string,
  1874. ) (openAIWSClientConn, int, http.Header, error) {
  1875. _ = ctx
  1876. _ = wsURL
  1877. _ = headers
  1878. _ = proxyURL
  1879. d.mu.Lock()
  1880. d.dialCount++
  1881. d.mu.Unlock()
  1882. return nil, 503, nil, errors.New("dial failed")
  1883. }
  1884. func (d *openAIWSAlwaysFailDialer) DialCount() int {
  1885. d.mu.Lock()
  1886. defer d.mu.Unlock()
  1887. return d.dialCount
  1888. }
  1889. type openAIWSFakeConn struct {
  1890. mu sync.Mutex
  1891. closed bool
  1892. payload [][]byte
  1893. }
  1894. func (c *openAIWSFakeConn) WriteJSON(ctx context.Context, value any) error {
  1895. _ = ctx
  1896. c.mu.Lock()
  1897. defer c.mu.Unlock()
  1898. if c.closed {
  1899. return errors.New("closed")
  1900. }
  1901. c.payload = append(c.payload, []byte("ok"))
  1902. _ = value
  1903. return nil
  1904. }
  1905. func (c *openAIWSFakeConn) ReadMessage(ctx context.Context) ([]byte, error) {
  1906. _ = ctx
  1907. c.mu.Lock()
  1908. defer c.mu.Unlock()
  1909. if c.closed {
  1910. return nil, errors.New("closed")
  1911. }
  1912. return []byte(`{"type":"response.completed","response":{"id":"resp_fake"}}`), nil
  1913. }
  1914. func (c *openAIWSFakeConn) Ping(ctx context.Context) error {
  1915. _ = ctx
  1916. return nil
  1917. }
  1918. func (c *openAIWSFakeConn) Close() error {
  1919. c.mu.Lock()
  1920. defer c.mu.Unlock()
  1921. c.closed = true
  1922. return nil
  1923. }
  1924. type openAIWSBlockingConn struct {
  1925. readDelay time.Duration
  1926. }
  1927. func (c *openAIWSBlockingConn) WriteJSON(ctx context.Context, value any) error {
  1928. _ = ctx
  1929. _ = value
  1930. return nil
  1931. }
  1932. func (c *openAIWSBlockingConn) ReadMessage(ctx context.Context) ([]byte, error) {
  1933. delay := c.readDelay
  1934. if delay <= 0 {
  1935. delay = 10 * time.Millisecond
  1936. }
  1937. timer := time.NewTimer(delay)
  1938. defer timer.Stop()
  1939. select {
  1940. case <-ctx.Done():
  1941. return nil, ctx.Err()
  1942. case <-timer.C:
  1943. return []byte(`{"type":"response.completed","response":{"id":"resp_blocking"}}`), nil
  1944. }
  1945. }
  1946. func (c *openAIWSBlockingConn) Ping(ctx context.Context) error {
  1947. _ = ctx
  1948. return nil
  1949. }
  1950. func (c *openAIWSBlockingConn) Close() error {
  1951. return nil
  1952. }
  1953. type openAIWSWriteBlockingConn struct{}
  1954. func (c *openAIWSWriteBlockingConn) WriteJSON(ctx context.Context, _ any) error {
  1955. <-ctx.Done()
  1956. return ctx.Err()
  1957. }
  1958. func (c *openAIWSWriteBlockingConn) ReadMessage(context.Context) ([]byte, error) {
  1959. return []byte(`{"type":"response.completed","response":{"id":"resp_write_block"}}`), nil
  1960. }
  1961. func (c *openAIWSWriteBlockingConn) Ping(context.Context) error {
  1962. return nil
  1963. }
  1964. func (c *openAIWSWriteBlockingConn) Close() error {
  1965. return nil
  1966. }
  1967. type openAIWSPingFailConn struct{}
  1968. func (c *openAIWSPingFailConn) WriteJSON(context.Context, any) error {
  1969. return nil
  1970. }
  1971. func (c *openAIWSPingFailConn) ReadMessage(context.Context) ([]byte, error) {
  1972. return []byte(`{"type":"response.completed","response":{"id":"resp_ping_fail"}}`), nil
  1973. }
  1974. func (c *openAIWSPingFailConn) Ping(context.Context) error {
  1975. return errors.New("ping failed")
  1976. }
  1977. func (c *openAIWSPingFailConn) Close() error {
  1978. return nil
  1979. }
  1980. type openAIWSContextProbeConn struct {
  1981. lastWriteCtx context.Context
  1982. }
  1983. func (c *openAIWSContextProbeConn) WriteJSON(ctx context.Context, _ any) error {
  1984. c.lastWriteCtx = ctx
  1985. return nil
  1986. }
  1987. func (c *openAIWSContextProbeConn) ReadMessage(context.Context) ([]byte, error) {
  1988. return []byte(`{"type":"response.completed","response":{"id":"resp_ctx_probe"}}`), nil
  1989. }
  1990. func (c *openAIWSContextProbeConn) Ping(context.Context) error {
  1991. return nil
  1992. }
  1993. func (c *openAIWSContextProbeConn) Close() error {
  1994. return nil
  1995. }
  1996. type openAIWSNilConnDialer struct{}
  1997. func (d *openAIWSNilConnDialer) Dial(
  1998. ctx context.Context,
  1999. wsURL string,
  2000. headers http.Header,
  2001. proxyURL string,
  2002. ) (openAIWSClientConn, int, http.Header, error) {
  2003. _ = ctx
  2004. _ = wsURL
  2005. _ = headers
  2006. _ = proxyURL
  2007. return nil, 200, nil, nil
  2008. }
  2009. func TestOpenAIWSConnPool_DialConnNilConnection(t *testing.T) {
  2010. cfg := &config.Config{}
  2011. cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
  2012. cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 1
  2013. pool := newOpenAIWSConnPool(cfg)
  2014. pool.setClientDialerForTest(&openAIWSNilConnDialer{})
  2015. account := &Account{ID: 91, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  2016. _, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
  2017. Account: account,
  2018. WSURL: "wss://example.com/v1/responses",
  2019. })
  2020. require.Error(t, err)
  2021. require.Contains(t, err.Error(), "nil connection")
  2022. }
  2023. func TestOpenAIWSConnPool_SnapshotTransportMetrics(t *testing.T) {
  2024. cfg := &config.Config{}
  2025. pool := newOpenAIWSConnPool(cfg)
  2026. dialer, ok := pool.clientDialer.(*coderOpenAIWSClientDialer)
  2027. require.True(t, ok)
  2028. _, err := dialer.proxyHTTPClient("http://127.0.0.1:28080")
  2029. require.NoError(t, err)
  2030. _, err = dialer.proxyHTTPClient("http://127.0.0.1:28080")
  2031. require.NoError(t, err)
  2032. _, err = dialer.proxyHTTPClient("http://127.0.0.1:28081")
  2033. require.NoError(t, err)
  2034. snapshot := pool.SnapshotTransportMetrics()
  2035. require.Equal(t, int64(1), snapshot.ProxyClientCacheHits)
  2036. require.Equal(t, int64(2), snapshot.ProxyClientCacheMisses)
  2037. require.InDelta(t, 1.0/3.0, snapshot.TransportReuseRatio, 0.0001)
  2038. }