content_moderation_runtime_cache_test.go 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456
  1. package service
  2. import (
  3. "context"
  4. "encoding/json"
  5. "errors"
  6. "sync"
  7. "testing"
  8. "time"
  9. "github.com/stretchr/testify/require"
  10. )
  11. type contentModerationRuntimeSettingRepo struct {
  12. mu sync.Mutex
  13. values map[string]string
  14. getValueCalls int
  15. getMultipleCalls int
  16. getMultipleErr error
  17. getMultipleStart chan<- struct{}
  18. getMultipleWait <-chan struct{}
  19. }
  20. func (r *contentModerationRuntimeSettingRepo) Get(_ context.Context, key string) (*Setting, error) {
  21. r.mu.Lock()
  22. defer r.mu.Unlock()
  23. value, ok := r.values[key]
  24. if !ok {
  25. return nil, ErrSettingNotFound
  26. }
  27. return &Setting{Key: key, Value: value}, nil
  28. }
  29. func (r *contentModerationRuntimeSettingRepo) GetValue(_ context.Context, key string) (string, error) {
  30. r.mu.Lock()
  31. defer r.mu.Unlock()
  32. r.getValueCalls++
  33. value, ok := r.values[key]
  34. if !ok {
  35. return "", ErrSettingNotFound
  36. }
  37. return value, nil
  38. }
  39. func (r *contentModerationRuntimeSettingRepo) Set(_ context.Context, key, value string) error {
  40. r.mu.Lock()
  41. defer r.mu.Unlock()
  42. if r.values == nil {
  43. r.values = make(map[string]string)
  44. }
  45. r.values[key] = value
  46. return nil
  47. }
  48. func (r *contentModerationRuntimeSettingRepo) GetMultiple(_ context.Context, keys []string) (map[string]string, error) {
  49. r.mu.Lock()
  50. r.getMultipleCalls++
  51. if err := r.getMultipleErr; err != nil {
  52. r.mu.Unlock()
  53. return nil, err
  54. }
  55. out := make(map[string]string, len(keys))
  56. for _, key := range keys {
  57. if value, ok := r.values[key]; ok {
  58. out[key] = value
  59. }
  60. }
  61. start := r.getMultipleStart
  62. wait := r.getMultipleWait
  63. r.getMultipleStart = nil
  64. r.getMultipleWait = nil
  65. r.mu.Unlock()
  66. if start != nil {
  67. start <- struct{}{}
  68. }
  69. if wait != nil {
  70. <-wait
  71. }
  72. return out, nil
  73. }
  74. func (r *contentModerationRuntimeSettingRepo) SetMultiple(_ context.Context, values map[string]string) error {
  75. r.mu.Lock()
  76. defer r.mu.Unlock()
  77. if r.values == nil {
  78. r.values = make(map[string]string)
  79. }
  80. for key, value := range values {
  81. r.values[key] = value
  82. }
  83. return nil
  84. }
  85. func (r *contentModerationRuntimeSettingRepo) GetAll(_ context.Context) (map[string]string, error) {
  86. r.mu.Lock()
  87. defer r.mu.Unlock()
  88. out := make(map[string]string, len(r.values))
  89. for key, value := range r.values {
  90. out[key] = value
  91. }
  92. return out, nil
  93. }
  94. func (r *contentModerationRuntimeSettingRepo) Delete(_ context.Context, key string) error {
  95. r.mu.Lock()
  96. defer r.mu.Unlock()
  97. delete(r.values, key)
  98. return nil
  99. }
  100. func (r *contentModerationRuntimeSettingRepo) calls() (getValue, getMultiple int) {
  101. r.mu.Lock()
  102. defer r.mu.Unlock()
  103. return r.getValueCalls, r.getMultipleCalls
  104. }
  105. func (r *contentModerationRuntimeSettingRepo) failMultiple(err error) {
  106. r.mu.Lock()
  107. defer r.mu.Unlock()
  108. r.getMultipleErr = err
  109. }
  110. func (r *contentModerationRuntimeSettingRepo) blockNextMultiple(start chan<- struct{}, wait <-chan struct{}) {
  111. r.mu.Lock()
  112. defer r.mu.Unlock()
  113. r.getMultipleStart = start
  114. r.getMultipleWait = wait
  115. }
  116. func runtimeCacheTestConfig(t *testing.T, keywords ...string) string {
  117. t.Helper()
  118. cfg := defaultContentModerationConfig()
  119. cfg.Enabled = true
  120. cfg.Mode = ContentModerationModePreBlock
  121. cfg.KeywordBlockingMode = ContentModerationKeywordModeKeywordOnly
  122. cfg.BlockedKeywords = keywords
  123. raw, err := json.Marshal(cfg)
  124. require.NoError(t, err)
  125. return string(raw)
  126. }
  127. func runtimeCacheTestService(repo *contentModerationRuntimeSettingRepo, ttl time.Duration) *ContentModerationService {
  128. return &ContentModerationService{
  129. settingRepo: repo,
  130. repo: &contentModerationTestRepo{},
  131. runtimeCacheTTL: ttl,
  132. }
  133. }
  134. func runtimeCacheTestInput(text string) ContentModerationCheckInput {
  135. return ContentModerationCheckInput{
  136. Protocol: ContentModerationProtocolOpenAIChat,
  137. Model: "risk-cache-test",
  138. Body: []byte(`{"messages":[{"role":"user","content":"` + text + `"}]}`),
  139. }
  140. }
  141. func TestContentModerationRuntimeSnapshotCachesSettings(t *testing.T) {
  142. repo := &contentModerationRuntimeSettingRepo{values: map[string]string{
  143. SettingKeyRiskControlEnabled: "true",
  144. SettingKeyContentModerationConfig: runtimeCacheTestConfig(t, "blocked"),
  145. }}
  146. svc := runtimeCacheTestService(repo, time.Hour)
  147. for range 20 {
  148. decision, err := svc.Check(context.Background(), runtimeCacheTestInput("clean prompt"))
  149. require.NoError(t, err)
  150. require.True(t, decision.Allowed)
  151. }
  152. getValue, getMultiple := repo.calls()
  153. require.Zero(t, getValue)
  154. require.Equal(t, 1, getMultiple)
  155. }
  156. func TestContentModerationRuntimeSnapshotUpdateConfigIsImmediate(t *testing.T) {
  157. repo := &contentModerationRuntimeSettingRepo{values: map[string]string{
  158. SettingKeyRiskControlEnabled: "true",
  159. SettingKeyContentModerationConfig: runtimeCacheTestConfig(t, "old-keyword"),
  160. }}
  161. svc := runtimeCacheTestService(repo, time.Hour)
  162. decision, err := svc.Check(context.Background(), runtimeCacheTestInput("new-keyword"))
  163. require.NoError(t, err)
  164. require.True(t, decision.Allowed)
  165. keywords := []string{"new-keyword"}
  166. _, err = svc.UpdateConfig(context.Background(), UpdateContentModerationConfigInput{
  167. BlockedKeywords: &keywords,
  168. })
  169. require.NoError(t, err)
  170. decision, err = svc.Check(context.Background(), runtimeCacheTestInput("new-keyword"))
  171. require.NoError(t, err)
  172. require.True(t, decision.Blocked)
  173. _, getMultiple := repo.calls()
  174. require.Equal(t, 1, getMultiple)
  175. }
  176. func TestContentModerationRuntimeSnapshotUpdateWinsOverInitialLoad(t *testing.T) {
  177. repo := &contentModerationRuntimeSettingRepo{values: map[string]string{
  178. SettingKeyRiskControlEnabled: "true",
  179. SettingKeyContentModerationConfig: runtimeCacheTestConfig(t, "old-keyword"),
  180. }}
  181. svc := runtimeCacheTestService(repo, time.Hour)
  182. refreshStarted := make(chan struct{}, 1)
  183. releaseRefresh := make(chan struct{})
  184. released := false
  185. defer func() {
  186. if !released {
  187. close(releaseRefresh)
  188. }
  189. }()
  190. repo.blockNextMultiple(refreshStarted, releaseRefresh)
  191. initialCheckDone := make(chan error, 1)
  192. go func() {
  193. decision, err := svc.Check(context.Background(), runtimeCacheTestInput("clean prompt"))
  194. if err == nil && (decision == nil || !decision.Allowed) {
  195. err = errors.New("unexpected initial moderation decision")
  196. }
  197. initialCheckDone <- err
  198. }()
  199. require.Eventually(t, func() bool {
  200. select {
  201. case <-refreshStarted:
  202. return true
  203. default:
  204. return false
  205. }
  206. }, time.Second, time.Millisecond)
  207. updateDone := make(chan error, 1)
  208. go func() {
  209. keywords := []string{"new-keyword"}
  210. _, updateErr := svc.UpdateConfig(context.Background(), UpdateContentModerationConfigInput{
  211. BlockedKeywords: &keywords,
  212. })
  213. updateDone <- updateErr
  214. }()
  215. select {
  216. case updateErr := <-updateDone:
  217. require.NoError(t, updateErr)
  218. t.Fatal("configuration update completed before the initial load released its lock")
  219. case <-time.After(10 * time.Millisecond):
  220. }
  221. close(releaseRefresh)
  222. released = true
  223. require.NoError(t, <-initialCheckDone)
  224. require.NoError(t, <-updateDone)
  225. decision, err := svc.Check(context.Background(), runtimeCacheTestInput("new-keyword"))
  226. require.NoError(t, err)
  227. require.True(t, decision.Blocked)
  228. decision, err = svc.Check(context.Background(), runtimeCacheTestInput("old-keyword"))
  229. require.NoError(t, err)
  230. require.True(t, decision.Allowed)
  231. }
  232. func TestContentModerationRuntimeSnapshotRefreshFailureKeepsStaleConfig(t *testing.T) {
  233. repo := &contentModerationRuntimeSettingRepo{values: map[string]string{
  234. SettingKeyRiskControlEnabled: "true",
  235. SettingKeyContentModerationConfig: runtimeCacheTestConfig(t, "blocked"),
  236. }}
  237. svc := runtimeCacheTestService(repo, time.Nanosecond)
  238. input := runtimeCacheTestInput("blocked")
  239. decision, err := svc.Check(context.Background(), input)
  240. require.NoError(t, err)
  241. require.True(t, decision.Blocked)
  242. repo.failMultiple(errors.New("database unavailable"))
  243. decision, err = svc.Check(context.Background(), input)
  244. require.NoError(t, err)
  245. require.True(t, decision.Blocked)
  246. // Windows 单调时钟步进约 0.5ms,两次相邻调用可能读到同一时刻,需反复调用才能观察到 TTL 过期。
  247. require.Eventually(t, func() bool {
  248. _, _ = svc.Check(context.Background(), input)
  249. _, calls := repo.calls()
  250. return calls >= 2
  251. }, time.Second, time.Millisecond)
  252. decision, err = svc.Check(context.Background(), input)
  253. require.NoError(t, err)
  254. require.True(t, decision.Blocked)
  255. }
  256. func TestContentModerationRuntimeSnapshotRefreshFailureBacksOff(t *testing.T) {
  257. repo := &contentModerationRuntimeSettingRepo{values: map[string]string{
  258. SettingKeyRiskControlEnabled: "true",
  259. SettingKeyContentModerationConfig: runtimeCacheTestConfig(t, "blocked"),
  260. }}
  261. svc := runtimeCacheTestService(repo, time.Minute)
  262. input := runtimeCacheTestInput("blocked")
  263. decision, err := svc.Check(context.Background(), input)
  264. require.NoError(t, err)
  265. require.True(t, decision.Blocked)
  266. current := svc.runtimeSnapshot.Load()
  267. require.NotNil(t, current)
  268. expired := *current
  269. expired.loadedAt = time.Now().Add(-2 * time.Minute)
  270. svc.runtimeSnapshot.Store(&expired)
  271. repo.failMultiple(errors.New("database unavailable"))
  272. decision, err = svc.Check(context.Background(), input)
  273. require.NoError(t, err)
  274. require.True(t, decision.Blocked)
  275. require.Eventually(t, func() bool {
  276. _, calls := repo.calls()
  277. return calls == 2
  278. }, time.Second, time.Millisecond)
  279. for range 100 {
  280. decision, err = svc.Check(context.Background(), input)
  281. require.NoError(t, err)
  282. require.True(t, decision.Blocked)
  283. }
  284. _, calls := repo.calls()
  285. require.Equal(t, 2, calls)
  286. }
  287. func TestContentModerationRuntimeSnapshotRefreshReusesUnchangedMatcher(t *testing.T) {
  288. repo := &contentModerationRuntimeSettingRepo{values: map[string]string{
  289. SettingKeyRiskControlEnabled: "true",
  290. SettingKeyContentModerationConfig: runtimeCacheTestConfig(t, "blocked"),
  291. }}
  292. svc := runtimeCacheTestService(repo, time.Minute)
  293. input := runtimeCacheTestInput("blocked")
  294. decision, err := svc.Check(context.Background(), input)
  295. require.NoError(t, err)
  296. require.True(t, decision.Blocked)
  297. current := svc.runtimeSnapshot.Load()
  298. require.NotNil(t, current)
  299. expired := *current
  300. expired.loadedAt = time.Now().Add(-2 * time.Minute)
  301. svc.runtimeSnapshot.Store(&expired)
  302. decision, err = svc.Check(context.Background(), input)
  303. require.NoError(t, err)
  304. require.True(t, decision.Blocked)
  305. require.Eventually(t, func() bool {
  306. refreshed := svc.runtimeSnapshot.Load()
  307. return refreshed != nil && refreshed.loadedAt.After(expired.loadedAt)
  308. }, time.Second, time.Millisecond)
  309. refreshed := svc.runtimeSnapshot.Load()
  310. require.Same(t, current.config, refreshed.config)
  311. require.Same(t, current.keywordMatcher, refreshed.keywordMatcher)
  312. _, calls := repo.calls()
  313. require.Equal(t, 2, calls)
  314. }
  315. func TestContentModerationRuntimeSnapshotUpdateWinsOverInFlightRefresh(t *testing.T) {
  316. repo := &contentModerationRuntimeSettingRepo{values: map[string]string{
  317. SettingKeyRiskControlEnabled: "true",
  318. SettingKeyContentModerationConfig: runtimeCacheTestConfig(t, "old-keyword"),
  319. }}
  320. svc := runtimeCacheTestService(repo, time.Minute)
  321. decision, err := svc.Check(context.Background(), runtimeCacheTestInput("old-keyword"))
  322. require.NoError(t, err)
  323. require.True(t, decision.Blocked)
  324. current := svc.runtimeSnapshot.Load()
  325. require.NotNil(t, current)
  326. expired := *current
  327. expired.loadedAt = time.Now().Add(-2 * time.Minute)
  328. svc.runtimeSnapshot.Store(&expired)
  329. refreshStarted := make(chan struct{}, 1)
  330. releaseRefresh := make(chan struct{})
  331. repo.blockNextMultiple(refreshStarted, releaseRefresh)
  332. decision, err = svc.Check(context.Background(), runtimeCacheTestInput("clean prompt"))
  333. require.NoError(t, err)
  334. require.True(t, decision.Allowed)
  335. require.Eventually(t, func() bool {
  336. select {
  337. case <-refreshStarted:
  338. return true
  339. default:
  340. return false
  341. }
  342. }, time.Second, time.Millisecond)
  343. updateDone := make(chan error, 1)
  344. go func() {
  345. keywords := []string{"new-keyword"}
  346. _, updateErr := svc.UpdateConfig(context.Background(), UpdateContentModerationConfigInput{
  347. BlockedKeywords: &keywords,
  348. })
  349. updateDone <- updateErr
  350. }()
  351. select {
  352. case updateErr := <-updateDone:
  353. require.NoError(t, updateErr)
  354. t.Fatal("configuration update completed before the in-flight refresh released its lock")
  355. case <-time.After(10 * time.Millisecond):
  356. }
  357. close(releaseRefresh)
  358. require.NoError(t, <-updateDone)
  359. decision, err = svc.Check(context.Background(), runtimeCacheTestInput("new-keyword"))
  360. require.NoError(t, err)
  361. require.True(t, decision.Blocked)
  362. decision, err = svc.Check(context.Background(), runtimeCacheTestInput("old-keyword"))
  363. require.NoError(t, err)
  364. require.True(t, decision.Allowed)
  365. }
  366. func TestContentModerationRuntimeSnapshotConcurrentReadAndReplace(t *testing.T) {
  367. repo := &contentModerationRuntimeSettingRepo{values: map[string]string{
  368. SettingKeyRiskControlEnabled: "true",
  369. SettingKeyContentModerationConfig: runtimeCacheTestConfig(t, "blocked-0"),
  370. }}
  371. svc := runtimeCacheTestService(repo, time.Hour)
  372. _, err := svc.Check(context.Background(), runtimeCacheTestInput("clean prompt"))
  373. require.NoError(t, err)
  374. var wg sync.WaitGroup
  375. errs := make(chan error, 8)
  376. for range 8 {
  377. wg.Add(1)
  378. go func() {
  379. defer wg.Done()
  380. for range 100 {
  381. decision, checkErr := svc.Check(context.Background(), runtimeCacheTestInput("clean prompt"))
  382. if checkErr != nil {
  383. errs <- checkErr
  384. return
  385. }
  386. if decision == nil || !decision.Allowed {
  387. errs <- errors.New("unexpected moderation decision")
  388. return
  389. }
  390. }
  391. }()
  392. }
  393. for i := 1; i <= 20; i++ {
  394. keywords := []string{"blocked-" + time.Duration(i).String()}
  395. _, err := svc.UpdateConfig(context.Background(), UpdateContentModerationConfigInput{
  396. BlockedKeywords: &keywords,
  397. })
  398. require.NoError(t, err)
  399. }
  400. wg.Wait()
  401. close(errs)
  402. for err := range errs {
  403. require.NoError(t, err)
  404. }
  405. }