openai_ws_state_store.go 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666
  1. package service
  2. import (
  3. "context"
  4. "crypto/sha256"
  5. "encoding/hex"
  6. "fmt"
  7. "strings"
  8. "sync"
  9. "sync/atomic"
  10. "time"
  11. )
  12. const (
  13. openAIWSResponseAccountCachePrefix = "openai:response:"
  14. openAIHTTPResponseOwnerUserPrefix = "openai:http-response-owner:user:"
  15. openAIHTTPResponseOwnerKeyPrefix = "openai:http-response-owner:key:"
  16. openAIWSStateStoreCleanupInterval = time.Minute
  17. openAIWSStateStoreCleanupMaxPerMap = 512
  18. openAIWSStateStoreMaxEntriesPerMap = 65536
  19. openAIWSStateStoreRedisTimeout = 3 * time.Second
  20. )
  21. type openAIWSAccountBinding struct {
  22. accountID int64
  23. expiresAt time.Time
  24. }
  25. type openAIHTTPResponseOwnerBinding struct {
  26. userID int64
  27. apiKeyID int64
  28. expiresAt time.Time
  29. }
  30. type openAIWSConnBinding struct {
  31. connID string
  32. expiresAt time.Time
  33. }
  34. type openAIWSTurnStateBinding struct {
  35. turnState string
  36. expiresAt time.Time
  37. }
  38. type openAIWSSessionConnBinding struct {
  39. connID string
  40. expiresAt time.Time
  41. }
  42. // openAIWSInvalidEncryptedBinding 记录一个会话中已被上游判定失效的
  43. // encrypted_content 摘要(invalid_encrypted_content lineage)。
  44. type openAIWSInvalidEncryptedBinding struct {
  45. digests map[string]struct{}
  46. expiresAt time.Time
  47. }
  48. // openAIWSInvalidEncryptedDigestsPerSession 是单会话摘要集合的存储自保护上限。
  49. // 超出后新增摘要被忽略,仅退化为"该项下次仍触发一次上游拒绝后的常规 recovery",
  50. // 不影响正确性。
  51. const openAIWSInvalidEncryptedDigestsPerSession = 512
  52. // OpenAIWSStateStore 管理 WSv2 的粘连状态。
  53. // - response_id -> account_id 用于续链路由
  54. // - response_id -> conn_id 用于连接内上下文复用
  55. //
  56. // response_id -> account_id 优先走 GatewayCache(Redis),同时维护本地热缓存。
  57. // response_id -> conn_id 仅在本进程内有效。
  58. type OpenAIWSStateStore interface {
  59. BindResponseAccount(ctx context.Context, groupID int64, responseID string, accountID int64, ttl time.Duration) error
  60. GetResponseAccount(ctx context.Context, groupID int64, responseID string) (int64, error)
  61. DeleteResponseAccount(ctx context.Context, groupID int64, responseID string) error
  62. BindHTTPResponseOwner(ctx context.Context, groupID int64, responseID string, userID, apiKeyID int64, ttl time.Duration) error
  63. GetHTTPResponseOwner(ctx context.Context, groupID int64, responseID string) (userID, apiKeyID int64, found bool, err error)
  64. BindResponseConn(responseID, connID string, ttl time.Duration)
  65. GetResponseConn(responseID string) (string, bool)
  66. DeleteResponseConn(responseID string)
  67. BindSessionTurnState(groupID int64, sessionHash, turnState string, ttl time.Duration)
  68. GetSessionTurnState(groupID int64, sessionHash string) (string, bool)
  69. DeleteSessionTurnState(groupID int64, sessionHash string)
  70. BindSessionConn(groupID int64, sessionHash, connID string, ttl time.Duration)
  71. GetSessionConn(groupID int64, sessionHash string) (string, bool)
  72. DeleteSessionConn(groupID int64, sessionHash string)
  73. // invalid_encrypted_content lineage:按会话记录已被上游拒绝的
  74. // encrypted_content 摘要,后续 turn 进场时仅剥离命中项,避免同一失效
  75. // 密文随客户端历史反复触发"整包被拒→剥离→重试/重连"。仅进程内有效。
  76. MarkSessionInvalidEncryptedContent(groupID int64, sessionHash string, digests []string, ttl time.Duration)
  77. GetSessionInvalidEncryptedContentDigests(groupID int64, sessionHash string) map[string]struct{}
  78. // HasAnySessionInvalidEncryptedContent 是热路径快速探测:全局无记录时
  79. // 调用方可跳过会话哈希计算与摘要匹配。
  80. HasAnySessionInvalidEncryptedContent() bool
  81. }
  82. type defaultOpenAIWSStateStore struct {
  83. cache GatewayCache
  84. responseToAccountMu sync.RWMutex
  85. responseToAccount map[string]openAIWSAccountBinding
  86. responseOwnerMu sync.RWMutex
  87. responseOwners map[string]openAIHTTPResponseOwnerBinding
  88. responseToConnMu sync.RWMutex
  89. responseToConn map[string]openAIWSConnBinding
  90. sessionToTurnStateMu sync.RWMutex
  91. sessionToTurnState map[string]openAIWSTurnStateBinding
  92. sessionToConnMu sync.RWMutex
  93. sessionToConn map[string]openAIWSSessionConnBinding
  94. sessionInvalidEncryptedMu sync.RWMutex
  95. sessionInvalidEncrypted map[string]openAIWSInvalidEncryptedBinding
  96. lastCleanupUnixNano atomic.Int64
  97. }
  98. // NewOpenAIWSStateStore 创建默认 WS 状态存储。
  99. func NewOpenAIWSStateStore(cache GatewayCache) OpenAIWSStateStore {
  100. store := &defaultOpenAIWSStateStore{
  101. cache: cache,
  102. responseToAccount: make(map[string]openAIWSAccountBinding, 256),
  103. responseOwners: make(map[string]openAIHTTPResponseOwnerBinding, 256),
  104. responseToConn: make(map[string]openAIWSConnBinding, 256),
  105. sessionToTurnState: make(map[string]openAIWSTurnStateBinding, 256),
  106. sessionToConn: make(map[string]openAIWSSessionConnBinding, 256),
  107. sessionInvalidEncrypted: make(map[string]openAIWSInvalidEncryptedBinding),
  108. }
  109. store.lastCleanupUnixNano.Store(time.Now().UnixNano())
  110. return store
  111. }
  112. func (s *defaultOpenAIWSStateStore) BindHTTPResponseOwner(ctx context.Context, groupID int64, responseID string, userID, apiKeyID int64, ttl time.Duration) error {
  113. id := normalizeOpenAIWSResponseID(responseID)
  114. if id == "" || userID <= 0 || apiKeyID <= 0 {
  115. return nil
  116. }
  117. ttl = normalizeOpenAIWSTTL(ttl)
  118. s.maybeCleanup()
  119. mapKey := openAIWSResponseAccountMapKey(groupID, id)
  120. s.responseOwnerMu.Lock()
  121. ensureBindingCapacity(s.responseOwners, mapKey, openAIWSStateStoreMaxEntriesPerMap)
  122. s.responseOwners[mapKey] = openAIHTTPResponseOwnerBinding{
  123. userID: userID, apiKeyID: apiKeyID, expiresAt: time.Now().Add(ttl),
  124. }
  125. s.responseOwnerMu.Unlock()
  126. if s.cache == nil {
  127. return nil
  128. }
  129. cacheCtx, cancel := withOpenAIWSStateStoreRedisTimeout(ctx)
  130. defer cancel()
  131. if err := s.cache.SetSessionAccountID(cacheCtx, groupID, openAIHTTPResponseOwnerCacheKey(openAIHTTPResponseOwnerUserPrefix, id), userID, ttl); err != nil {
  132. return err
  133. }
  134. return s.cache.SetSessionAccountID(cacheCtx, groupID, openAIHTTPResponseOwnerCacheKey(openAIHTTPResponseOwnerKeyPrefix, id), apiKeyID, ttl)
  135. }
  136. func (s *defaultOpenAIWSStateStore) GetHTTPResponseOwner(ctx context.Context, groupID int64, responseID string) (int64, int64, bool, error) {
  137. id := normalizeOpenAIWSResponseID(responseID)
  138. if id == "" {
  139. return 0, 0, false, nil
  140. }
  141. s.maybeCleanup()
  142. now := time.Now()
  143. mapKey := openAIWSResponseAccountMapKey(groupID, id)
  144. s.responseOwnerMu.RLock()
  145. if binding, ok := s.responseOwners[mapKey]; ok && now.Before(binding.expiresAt) {
  146. s.responseOwnerMu.RUnlock()
  147. return binding.userID, binding.apiKeyID, true, nil
  148. }
  149. s.responseOwnerMu.RUnlock()
  150. if s.cache == nil {
  151. return 0, 0, false, nil
  152. }
  153. cacheCtx, cancel := withOpenAIWSStateStoreRedisTimeout(ctx)
  154. defer cancel()
  155. userID, err := s.cache.GetSessionAccountID(cacheCtx, groupID, openAIHTTPResponseOwnerCacheKey(openAIHTTPResponseOwnerUserPrefix, id))
  156. if err != nil || userID <= 0 {
  157. return 0, 0, false, err
  158. }
  159. apiKeyID, err := s.cache.GetSessionAccountID(cacheCtx, groupID, openAIHTTPResponseOwnerCacheKey(openAIHTTPResponseOwnerKeyPrefix, id))
  160. if err != nil || apiKeyID <= 0 {
  161. return 0, 0, false, err
  162. }
  163. s.responseOwnerMu.Lock()
  164. ensureBindingCapacity(s.responseOwners, mapKey, openAIWSStateStoreMaxEntriesPerMap)
  165. s.responseOwners[mapKey] = openAIHTTPResponseOwnerBinding{
  166. userID: userID, apiKeyID: apiKeyID, expiresAt: now.Add(time.Minute),
  167. }
  168. s.responseOwnerMu.Unlock()
  169. return userID, apiKeyID, true, nil
  170. }
  171. func (s *defaultOpenAIWSStateStore) BindResponseAccount(ctx context.Context, groupID int64, responseID string, accountID int64, ttl time.Duration) error {
  172. id := normalizeOpenAIWSResponseID(responseID)
  173. if id == "" || accountID <= 0 {
  174. return nil
  175. }
  176. ttl = normalizeOpenAIWSTTL(ttl)
  177. s.maybeCleanup()
  178. expiresAt := time.Now().Add(ttl)
  179. mapKey := openAIWSResponseAccountMapKey(groupID, id)
  180. s.responseToAccountMu.Lock()
  181. ensureBindingCapacity(s.responseToAccount, mapKey, openAIWSStateStoreMaxEntriesPerMap)
  182. s.responseToAccount[mapKey] = openAIWSAccountBinding{accountID: accountID, expiresAt: expiresAt}
  183. s.responseToAccountMu.Unlock()
  184. if s.cache == nil {
  185. return nil
  186. }
  187. cacheKey := openAIWSResponseAccountCacheKey(id)
  188. cacheCtx, cancel := withOpenAIWSStateStoreRedisTimeout(ctx)
  189. defer cancel()
  190. return s.cache.SetSessionAccountID(cacheCtx, groupID, cacheKey, accountID, ttl)
  191. }
  192. func cleanupExpiredHTTPResponseOwnerBindings(bindings map[string]openAIHTTPResponseOwnerBinding, now time.Time, maxScan int) {
  193. if len(bindings) == 0 || maxScan <= 0 {
  194. return
  195. }
  196. scanned := 0
  197. for key, binding := range bindings {
  198. if now.After(binding.expiresAt) {
  199. delete(bindings, key)
  200. }
  201. scanned++
  202. if scanned >= maxScan {
  203. break
  204. }
  205. }
  206. }
  207. func (s *defaultOpenAIWSStateStore) GetResponseAccount(ctx context.Context, groupID int64, responseID string) (int64, error) {
  208. id := normalizeOpenAIWSResponseID(responseID)
  209. if id == "" {
  210. return 0, nil
  211. }
  212. s.maybeCleanup()
  213. now := time.Now()
  214. mapKey := openAIWSResponseAccountMapKey(groupID, id)
  215. s.responseToAccountMu.RLock()
  216. if binding, ok := s.responseToAccount[mapKey]; ok {
  217. if now.Before(binding.expiresAt) {
  218. accountID := binding.accountID
  219. s.responseToAccountMu.RUnlock()
  220. return accountID, nil
  221. }
  222. }
  223. s.responseToAccountMu.RUnlock()
  224. if s.cache == nil {
  225. return 0, nil
  226. }
  227. cacheKey := openAIWSResponseAccountCacheKey(id)
  228. cacheCtx, cancel := withOpenAIWSStateStoreRedisTimeout(ctx)
  229. defer cancel()
  230. accountID, err := s.cache.GetSessionAccountID(cacheCtx, groupID, cacheKey)
  231. if err != nil || accountID <= 0 {
  232. // 缓存读取失败不阻断主流程,按未命中降级。
  233. return 0, nil
  234. }
  235. return accountID, nil
  236. }
  237. func (s *defaultOpenAIWSStateStore) DeleteResponseAccount(ctx context.Context, groupID int64, responseID string) error {
  238. id := normalizeOpenAIWSResponseID(responseID)
  239. if id == "" {
  240. return nil
  241. }
  242. s.responseToAccountMu.Lock()
  243. delete(s.responseToAccount, openAIWSResponseAccountMapKey(groupID, id))
  244. s.responseToAccountMu.Unlock()
  245. if s.cache == nil {
  246. return nil
  247. }
  248. cacheCtx, cancel := withOpenAIWSStateStoreRedisTimeout(ctx)
  249. defer cancel()
  250. return s.cache.DeleteSessionAccountID(cacheCtx, groupID, openAIWSResponseAccountCacheKey(id))
  251. }
  252. func (s *defaultOpenAIWSStateStore) BindResponseConn(responseID, connID string, ttl time.Duration) {
  253. id := normalizeOpenAIWSResponseID(responseID)
  254. conn := strings.TrimSpace(connID)
  255. if id == "" || conn == "" {
  256. return
  257. }
  258. ttl = normalizeOpenAIWSTTL(ttl)
  259. s.maybeCleanup()
  260. s.responseToConnMu.Lock()
  261. ensureBindingCapacity(s.responseToConn, id, openAIWSStateStoreMaxEntriesPerMap)
  262. s.responseToConn[id] = openAIWSConnBinding{
  263. connID: conn,
  264. expiresAt: time.Now().Add(ttl),
  265. }
  266. s.responseToConnMu.Unlock()
  267. }
  268. func (s *defaultOpenAIWSStateStore) GetResponseConn(responseID string) (string, bool) {
  269. id := normalizeOpenAIWSResponseID(responseID)
  270. if id == "" {
  271. return "", false
  272. }
  273. s.maybeCleanup()
  274. now := time.Now()
  275. s.responseToConnMu.RLock()
  276. binding, ok := s.responseToConn[id]
  277. s.responseToConnMu.RUnlock()
  278. if !ok || now.After(binding.expiresAt) || strings.TrimSpace(binding.connID) == "" {
  279. return "", false
  280. }
  281. return binding.connID, true
  282. }
  283. func (s *defaultOpenAIWSStateStore) DeleteResponseConn(responseID string) {
  284. id := normalizeOpenAIWSResponseID(responseID)
  285. if id == "" {
  286. return
  287. }
  288. s.responseToConnMu.Lock()
  289. delete(s.responseToConn, id)
  290. s.responseToConnMu.Unlock()
  291. }
  292. func (s *defaultOpenAIWSStateStore) BindSessionTurnState(groupID int64, sessionHash, turnState string, ttl time.Duration) {
  293. key := openAIWSSessionTurnStateKey(groupID, sessionHash)
  294. state := strings.TrimSpace(turnState)
  295. if key == "" || state == "" {
  296. return
  297. }
  298. ttl = normalizeOpenAIWSTTL(ttl)
  299. s.maybeCleanup()
  300. s.sessionToTurnStateMu.Lock()
  301. ensureBindingCapacity(s.sessionToTurnState, key, openAIWSStateStoreMaxEntriesPerMap)
  302. s.sessionToTurnState[key] = openAIWSTurnStateBinding{
  303. turnState: state,
  304. expiresAt: time.Now().Add(ttl),
  305. }
  306. s.sessionToTurnStateMu.Unlock()
  307. }
  308. func (s *defaultOpenAIWSStateStore) GetSessionTurnState(groupID int64, sessionHash string) (string, bool) {
  309. key := openAIWSSessionTurnStateKey(groupID, sessionHash)
  310. if key == "" {
  311. return "", false
  312. }
  313. s.maybeCleanup()
  314. now := time.Now()
  315. s.sessionToTurnStateMu.RLock()
  316. binding, ok := s.sessionToTurnState[key]
  317. s.sessionToTurnStateMu.RUnlock()
  318. if !ok || now.After(binding.expiresAt) || strings.TrimSpace(binding.turnState) == "" {
  319. return "", false
  320. }
  321. return binding.turnState, true
  322. }
  323. func (s *defaultOpenAIWSStateStore) DeleteSessionTurnState(groupID int64, sessionHash string) {
  324. key := openAIWSSessionTurnStateKey(groupID, sessionHash)
  325. if key == "" {
  326. return
  327. }
  328. s.sessionToTurnStateMu.Lock()
  329. delete(s.sessionToTurnState, key)
  330. s.sessionToTurnStateMu.Unlock()
  331. }
  332. func (s *defaultOpenAIWSStateStore) BindSessionConn(groupID int64, sessionHash, connID string, ttl time.Duration) {
  333. key := openAIWSSessionTurnStateKey(groupID, sessionHash)
  334. conn := strings.TrimSpace(connID)
  335. if key == "" || conn == "" {
  336. return
  337. }
  338. ttl = normalizeOpenAIWSTTL(ttl)
  339. s.maybeCleanup()
  340. s.sessionToConnMu.Lock()
  341. ensureBindingCapacity(s.sessionToConn, key, openAIWSStateStoreMaxEntriesPerMap)
  342. s.sessionToConn[key] = openAIWSSessionConnBinding{
  343. connID: conn,
  344. expiresAt: time.Now().Add(ttl),
  345. }
  346. s.sessionToConnMu.Unlock()
  347. }
  348. func (s *defaultOpenAIWSStateStore) GetSessionConn(groupID int64, sessionHash string) (string, bool) {
  349. key := openAIWSSessionTurnStateKey(groupID, sessionHash)
  350. if key == "" {
  351. return "", false
  352. }
  353. s.maybeCleanup()
  354. now := time.Now()
  355. s.sessionToConnMu.RLock()
  356. binding, ok := s.sessionToConn[key]
  357. s.sessionToConnMu.RUnlock()
  358. if !ok || now.After(binding.expiresAt) || strings.TrimSpace(binding.connID) == "" {
  359. return "", false
  360. }
  361. return binding.connID, true
  362. }
  363. func (s *defaultOpenAIWSStateStore) DeleteSessionConn(groupID int64, sessionHash string) {
  364. key := openAIWSSessionTurnStateKey(groupID, sessionHash)
  365. if key == "" {
  366. return
  367. }
  368. s.sessionToConnMu.Lock()
  369. delete(s.sessionToConn, key)
  370. s.sessionToConnMu.Unlock()
  371. }
  372. func (s *defaultOpenAIWSStateStore) MarkSessionInvalidEncryptedContent(groupID int64, sessionHash string, digests []string, ttl time.Duration) {
  373. key := openAIWSSessionTurnStateKey(groupID, sessionHash)
  374. if key == "" || len(digests) == 0 {
  375. return
  376. }
  377. ttl = normalizeOpenAIWSTTL(ttl)
  378. s.maybeCleanup()
  379. now := time.Now()
  380. overflowed := 0
  381. s.sessionInvalidEncryptedMu.Lock()
  382. ensureBindingCapacity(s.sessionInvalidEncrypted, key, openAIWSStateStoreMaxEntriesPerMap)
  383. binding, ok := s.sessionInvalidEncrypted[key]
  384. if !ok || now.After(binding.expiresAt) || binding.digests == nil {
  385. binding = openAIWSInvalidEncryptedBinding{digests: make(map[string]struct{}, len(digests))}
  386. }
  387. for _, digest := range digests {
  388. digest = strings.TrimSpace(digest)
  389. if digest == "" {
  390. continue
  391. }
  392. if len(binding.digests) >= openAIWSInvalidEncryptedDigestsPerSession {
  393. if _, exists := binding.digests[digest]; !exists {
  394. overflowed++
  395. }
  396. continue
  397. }
  398. binding.digests[digest] = struct{}{}
  399. }
  400. binding.expiresAt = now.Add(ttl)
  401. s.sessionInvalidEncrypted[key] = binding
  402. s.sessionInvalidEncryptedMu.Unlock()
  403. // HasAny/Get 在每个 OpenAI 请求热路径上抢同一把锁,日志 IO 必须在锁外。
  404. if overflowed > 0 {
  405. logOpenAIWSModeInfo(
  406. "invalid_encrypted_lineage_capacity_overflow dropped_digests=%d capacity=%d",
  407. overflowed,
  408. openAIWSInvalidEncryptedDigestsPerSession,
  409. )
  410. }
  411. }
  412. func (s *defaultOpenAIWSStateStore) GetSessionInvalidEncryptedContentDigests(groupID int64, sessionHash string) map[string]struct{} {
  413. key := openAIWSSessionTurnStateKey(groupID, sessionHash)
  414. if key == "" {
  415. return nil
  416. }
  417. s.maybeCleanup()
  418. now := time.Now()
  419. s.sessionInvalidEncryptedMu.RLock()
  420. defer s.sessionInvalidEncryptedMu.RUnlock()
  421. binding, ok := s.sessionInvalidEncrypted[key]
  422. if !ok || now.After(binding.expiresAt) || len(binding.digests) == 0 {
  423. return nil
  424. }
  425. digests := make(map[string]struct{}, len(binding.digests))
  426. for digest := range binding.digests {
  427. digests[digest] = struct{}{}
  428. }
  429. return digests
  430. }
  431. func (s *defaultOpenAIWSStateStore) HasAnySessionInvalidEncryptedContent() bool {
  432. s.sessionInvalidEncryptedMu.RLock()
  433. defer s.sessionInvalidEncryptedMu.RUnlock()
  434. return len(s.sessionInvalidEncrypted) > 0
  435. }
  436. func (s *defaultOpenAIWSStateStore) maybeCleanup() {
  437. if s == nil {
  438. return
  439. }
  440. now := time.Now()
  441. last := time.Unix(0, s.lastCleanupUnixNano.Load())
  442. if now.Sub(last) < openAIWSStateStoreCleanupInterval {
  443. return
  444. }
  445. if !s.lastCleanupUnixNano.CompareAndSwap(last.UnixNano(), now.UnixNano()) {
  446. return
  447. }
  448. // 增量限额清理,避免高规模下一次性全量扫描导致长时间阻塞。
  449. s.responseToAccountMu.Lock()
  450. cleanupExpiredAccountBindings(s.responseToAccount, now, openAIWSStateStoreCleanupMaxPerMap)
  451. s.responseToAccountMu.Unlock()
  452. s.responseOwnerMu.Lock()
  453. cleanupExpiredHTTPResponseOwnerBindings(s.responseOwners, now, openAIWSStateStoreCleanupMaxPerMap)
  454. s.responseOwnerMu.Unlock()
  455. s.responseToConnMu.Lock()
  456. cleanupExpiredConnBindings(s.responseToConn, now, openAIWSStateStoreCleanupMaxPerMap)
  457. s.responseToConnMu.Unlock()
  458. s.sessionToTurnStateMu.Lock()
  459. cleanupExpiredTurnStateBindings(s.sessionToTurnState, now, openAIWSStateStoreCleanupMaxPerMap)
  460. s.sessionToTurnStateMu.Unlock()
  461. s.sessionToConnMu.Lock()
  462. cleanupExpiredSessionConnBindings(s.sessionToConn, now, openAIWSStateStoreCleanupMaxPerMap)
  463. s.sessionToConnMu.Unlock()
  464. s.sessionInvalidEncryptedMu.Lock()
  465. cleanupExpiredInvalidEncryptedBindings(s.sessionInvalidEncrypted, now, openAIWSStateStoreCleanupMaxPerMap)
  466. s.sessionInvalidEncryptedMu.Unlock()
  467. }
  468. func cleanupExpiredInvalidEncryptedBindings(bindings map[string]openAIWSInvalidEncryptedBinding, now time.Time, maxScan int) {
  469. if len(bindings) == 0 || maxScan <= 0 {
  470. return
  471. }
  472. scanned := 0
  473. for key, binding := range bindings {
  474. if now.After(binding.expiresAt) {
  475. delete(bindings, key)
  476. }
  477. scanned++
  478. if scanned >= maxScan {
  479. break
  480. }
  481. }
  482. }
  483. func cleanupExpiredAccountBindings(bindings map[string]openAIWSAccountBinding, now time.Time, maxScan int) {
  484. if len(bindings) == 0 || maxScan <= 0 {
  485. return
  486. }
  487. scanned := 0
  488. for key, binding := range bindings {
  489. if now.After(binding.expiresAt) {
  490. delete(bindings, key)
  491. }
  492. scanned++
  493. if scanned >= maxScan {
  494. break
  495. }
  496. }
  497. }
  498. func cleanupExpiredConnBindings(bindings map[string]openAIWSConnBinding, now time.Time, maxScan int) {
  499. if len(bindings) == 0 || maxScan <= 0 {
  500. return
  501. }
  502. scanned := 0
  503. for key, binding := range bindings {
  504. if now.After(binding.expiresAt) {
  505. delete(bindings, key)
  506. }
  507. scanned++
  508. if scanned >= maxScan {
  509. break
  510. }
  511. }
  512. }
  513. func cleanupExpiredTurnStateBindings(bindings map[string]openAIWSTurnStateBinding, now time.Time, maxScan int) {
  514. if len(bindings) == 0 || maxScan <= 0 {
  515. return
  516. }
  517. scanned := 0
  518. for key, binding := range bindings {
  519. if now.After(binding.expiresAt) {
  520. delete(bindings, key)
  521. }
  522. scanned++
  523. if scanned >= maxScan {
  524. break
  525. }
  526. }
  527. }
  528. func cleanupExpiredSessionConnBindings(bindings map[string]openAIWSSessionConnBinding, now time.Time, maxScan int) {
  529. if len(bindings) == 0 || maxScan <= 0 {
  530. return
  531. }
  532. scanned := 0
  533. for key, binding := range bindings {
  534. if now.After(binding.expiresAt) {
  535. delete(bindings, key)
  536. }
  537. scanned++
  538. if scanned >= maxScan {
  539. break
  540. }
  541. }
  542. }
  543. func ensureBindingCapacity[T any](bindings map[string]T, incomingKey string, maxEntries int) {
  544. if len(bindings) < maxEntries || maxEntries <= 0 {
  545. return
  546. }
  547. if _, exists := bindings[incomingKey]; exists {
  548. return
  549. }
  550. // 固定上限保护:淘汰任意一项,优先保证内存有界。
  551. for key := range bindings {
  552. delete(bindings, key)
  553. return
  554. }
  555. }
  556. func normalizeOpenAIWSResponseID(responseID string) string {
  557. return strings.TrimSpace(responseID)
  558. }
  559. func openAIWSResponseAccountCacheKey(responseID string) string {
  560. sum := sha256.Sum256([]byte(responseID))
  561. return openAIWSResponseAccountCachePrefix + hex.EncodeToString(sum[:])
  562. }
  563. func openAIHTTPResponseOwnerCacheKey(prefix, responseID string) string {
  564. sum := sha256.Sum256([]byte(responseID))
  565. return prefix + hex.EncodeToString(sum[:])
  566. }
  567. // openAIWSResponseAccountMapKey 本地热缓存按分组隔离的 key,与 Redis 层保持一致,避免跨组命中。
  568. func openAIWSResponseAccountMapKey(groupID int64, responseID string) string {
  569. return fmt.Sprintf("%d:%s", groupID, responseID)
  570. }
  571. func normalizeOpenAIWSTTL(ttl time.Duration) time.Duration {
  572. if ttl <= 0 {
  573. return time.Hour
  574. }
  575. return ttl
  576. }
  577. func openAIWSSessionTurnStateKey(groupID int64, sessionHash string) string {
  578. hash := strings.TrimSpace(sessionHash)
  579. if hash == "" {
  580. return ""
  581. }
  582. return fmt.Sprintf("%d:%s", groupID, hash)
  583. }
  584. func withOpenAIWSStateStoreRedisTimeout(ctx context.Context) (context.Context, context.CancelFunc) {
  585. if ctx == nil {
  586. ctx = context.Background()
  587. }
  588. return context.WithTimeout(ctx, openAIWSStateStoreRedisTimeout)
  589. }