wonder
0027704169
- 修复多个 goroutine 并发 append 同一 slice 导致的 DATA RACE - 使用预分配 slice + 索引赋值替代并发 append - 通过 go test -race 验证
441 lines
10 KiB
Go
441 lines
10 KiB
Go
package service
|
||
|
||
import (
|
||
"context"
|
||
"sync"
|
||
"testing"
|
||
"time"
|
||
|
||
"github.com/1024XEngineer/xinfra/server/internal/model"
|
||
)
|
||
|
||
// newTestDeliveryService 创建一个用于测试的 DeliveryService，不需要真实的 DB 和配置。
|
||
// 仅适用于测试 pub/sub、缓存等内存逻辑。
|
||
func newTestDeliveryService() *DeliveryService {
|
||
return &DeliveryService{
|
||
streams: make(map[string]map[chan DeliveryTaskSnapshot]struct{}),
|
||
stdoutCache: make(map[string]*stdoutCacheItem),
|
||
}
|
||
}
|
||
|
||
func TestSubscribeTask(t *testing.T) {
|
||
svc := newTestDeliveryService()
|
||
ctx := context.Background()
|
||
_ = ctx
|
||
|
||
taskID := "test-task-1"
|
||
|
||
// 订阅任务
|
||
ch, cancel := svc.SubscribeTask(taskID)
|
||
defer cancel()
|
||
|
||
// 验证 channel 已注册
|
||
svc.streamMu.Lock()
|
||
if _, ok := svc.streams[taskID]; !ok {
|
||
svc.streamMu.Unlock()
|
||
t.Fatal("SubscribeTask did not register channel in streams map")
|
||
}
|
||
svc.streamMu.Unlock()
|
||
|
||
// 模拟 broadcastTask 推送 snapshot
|
||
snapshot := DeliveryTaskSnapshot{
|
||
Task: &model.DeliveryTask{ID: taskID, Status: model.TaskRunning},
|
||
}
|
||
svc.streamMu.Lock()
|
||
for ch := range svc.streams[taskID] {
|
||
select {
|
||
case ch <- snapshot:
|
||
default:
|
||
}
|
||
}
|
||
svc.streamMu.Unlock()
|
||
|
||
// 接收推送
|
||
select {
|
||
case received := <-ch:
|
||
if received.Task == nil || received.Task.ID != taskID {
|
||
t.Fatalf("received snapshot task ID = %v, want %q", received.Task, taskID)
|
||
}
|
||
if received.Task.Status != model.TaskRunning {
|
||
t.Fatalf("received snapshot status = %q, want %q", received.Task.Status, model.TaskRunning)
|
||
}
|
||
case <-time.After(time.Second):
|
||
t.Fatal("timeout waiting for snapshot from SubscribeTask")
|
||
}
|
||
}
|
||
|
||
func TestSubscribeTask_Cancel(t *testing.T) {
|
||
svc := newTestDeliveryService()
|
||
taskID := "test-task-cancel"
|
||
|
||
ch, cancel := svc.SubscribeTask(taskID)
|
||
|
||
// 调用 cancel
|
||
cancel()
|
||
|
||
// 验证 channel 已关闭
|
||
select {
|
||
case _, ok := <-ch:
|
||
if ok {
|
||
t.Fatal("channel should be closed after cancel, but got a value")
|
||
}
|
||
case <-time.After(time.Second):
|
||
t.Fatal("timeout waiting for channel close")
|
||
}
|
||
|
||
// 验证已从 streams 中移除
|
||
svc.streamMu.Lock()
|
||
if subs := svc.streams[taskID]; subs != nil && len(subs) > 0 {
|
||
svc.streamMu.Unlock()
|
||
t.Fatal("cancel did not remove channel from streams map")
|
||
}
|
||
svc.streamMu.Unlock()
|
||
}
|
||
|
||
func TestSubscribeTask_MultipleSubscribers(t *testing.T) {
|
||
svc := newTestDeliveryService()
|
||
taskID := "test-task-multi"
|
||
|
||
ch1, cancel1 := svc.SubscribeTask(taskID)
|
||
defer cancel1()
|
||
ch2, cancel2 := svc.SubscribeTask(taskID)
|
||
defer cancel2()
|
||
ch3, cancel3 := svc.SubscribeTask(taskID)
|
||
defer cancel3()
|
||
|
||
// 验证三个订阅者都已注册
|
||
svc.streamMu.Lock()
|
||
subs := svc.streams[taskID]
|
||
if len(subs) != 3 {
|
||
svc.streamMu.Unlock()
|
||
t.Fatalf("expected 3 subscribers, got %d", len(subs))
|
||
}
|
||
svc.streamMu.Unlock()
|
||
|
||
// 推送 snapshot
|
||
snapshot := DeliveryTaskSnapshot{
|
||
Task: &model.DeliveryTask{ID: taskID, Status: model.TaskFinished},
|
||
}
|
||
svc.streamMu.Lock()
|
||
for ch := range svc.streams[taskID] {
|
||
select {
|
||
case ch <- snapshot:
|
||
default:
|
||
}
|
||
}
|
||
svc.streamMu.Unlock()
|
||
|
||
// 验证三个订阅者都收到
|
||
for i, ch := range []<-chan DeliveryTaskSnapshot{ch1, ch2, ch3} {
|
||
select {
|
||
case received := <-ch:
|
||
if received.Task == nil || received.Task.Status != model.TaskFinished {
|
||
t.Fatalf("subscriber %d: expected TaskFinished, got %+v", i+1, received)
|
||
}
|
||
case <-time.After(time.Second):
|
||
t.Fatalf("subscriber %d: timeout waiting for snapshot", i+1)
|
||
}
|
||
}
|
||
}
|
||
|
||
func TestSubscribeTask_CancelOneDoesNotAffectOthers(t *testing.T) {
|
||
svc := newTestDeliveryService()
|
||
taskID := "test-task-cancel-one"
|
||
|
||
ch1, cancel1 := svc.SubscribeTask(taskID)
|
||
_, cancel2 := svc.SubscribeTask(taskID)
|
||
_ = cancel2
|
||
|
||
// 取消第一个订阅者
|
||
cancel1()
|
||
|
||
// 验证还剩一个订阅者
|
||
svc.streamMu.Lock()
|
||
subs := svc.streams[taskID]
|
||
if len(subs) != 1 {
|
||
svc.streamMu.Unlock()
|
||
t.Fatalf("expected 1 subscriber after cancel, got %d", len(subs))
|
||
}
|
||
svc.streamMu.Unlock()
|
||
|
||
// 推送 snapshot
|
||
snapshot := DeliveryTaskSnapshot{
|
||
Task: &model.DeliveryTask{ID: taskID, Status: model.TaskRunning},
|
||
}
|
||
svc.streamMu.Lock()
|
||
for ch := range svc.streams[taskID] {
|
||
select {
|
||
case ch <- snapshot:
|
||
default:
|
||
}
|
||
}
|
||
svc.streamMu.Unlock()
|
||
|
||
// ch1 已关闭，不应收到消息
|
||
select {
|
||
case _, ok := <-ch1:
|
||
if ok {
|
||
t.Fatal("ch1 should be closed after cancel")
|
||
}
|
||
case <-time.After(100 * time.Millisecond):
|
||
// OK: channel is closed, no value received
|
||
}
|
||
}
|
||
|
||
func TestBroadcastTask_ClosedChannelCleanup(t *testing.T) {
|
||
svc := newTestDeliveryService()
|
||
taskID := "test-task-cleanup"
|
||
|
||
// 创建一个订阅者并立即关闭
|
||
_, cancel := svc.SubscribeTask(taskID)
|
||
cancel()
|
||
// 等待 cancel 完成
|
||
time.Sleep(10 * time.Millisecond)
|
||
|
||
// 创建一个新的正常订阅者
|
||
ch2, cancel2 := svc.SubscribeTask(taskID)
|
||
defer cancel2()
|
||
|
||
// 模拟 broadcastTask 行为（带 recover）
|
||
snapshot := DeliveryTaskSnapshot{
|
||
Task: &model.DeliveryTask{ID: taskID, Status: model.TaskRunning},
|
||
}
|
||
|
||
svc.streamMu.Lock()
|
||
var closedChannels []chan DeliveryTaskSnapshot
|
||
for ch := range svc.streams[taskID] {
|
||
func() {
|
||
defer func() {
|
||
if r := recover(); r != nil {
|
||
closedChannels = append(closedChannels, ch)
|
||
}
|
||
}()
|
||
select {
|
||
case ch <- snapshot:
|
||
default:
|
||
}
|
||
}()
|
||
}
|
||
for _, ch := range closedChannels {
|
||
delete(svc.streams[taskID], ch)
|
||
}
|
||
if len(svc.streams[taskID]) == 0 {
|
||
delete(svc.streams, taskID)
|
||
}
|
||
svc.streamMu.Unlock()
|
||
|
||
// 验证 closed channel 被清理
|
||
svc.streamMu.Lock()
|
||
if subs := svc.streams[taskID]; subs != nil && len(subs) != 1 {
|
||
svc.streamMu.Unlock()
|
||
t.Fatalf("expected 1 subscriber after cleanup, got %d", len(subs))
|
||
}
|
||
svc.streamMu.Unlock()
|
||
|
||
// 正常订阅者应该收到消息
|
||
select {
|
||
case received := <-ch2:
|
||
if received.Task == nil || received.Task.ID != taskID {
|
||
t.Fatalf("expected task ID %q, got %+v", taskID, received)
|
||
}
|
||
case <-time.After(time.Second):
|
||
t.Fatal("timeout waiting for snapshot from normal subscriber")
|
||
}
|
||
}
|
||
|
||
func TestInvalidateStdoutCache(t *testing.T) {
|
||
svc := newTestDeliveryService()
|
||
|
||
// 填充缓存
|
||
svc.cacheMu.Lock()
|
||
svc.stdoutCache["job-1"] = &stdoutCacheItem{stdout: "cached output", createdAt: time.Now()}
|
||
svc.cacheMu.Unlock()
|
||
|
||
// 验证缓存命中
|
||
svc.cacheMu.RLock()
|
||
item, ok := svc.stdoutCache["job-1"]
|
||
svc.cacheMu.RUnlock()
|
||
if !ok {
|
||
t.Fatal("cache entry not found before invalidation")
|
||
}
|
||
if item.stdout != "cached output" {
|
||
t.Fatalf("cache stdout = %q, want %q", item.stdout, "cached output")
|
||
}
|
||
|
||
// 失效缓存
|
||
svc.InvalidateStdoutCache("job-1")
|
||
|
||
// 验证缓存已失效
|
||
svc.cacheMu.RLock()
|
||
_, ok = svc.stdoutCache["job-1"]
|
||
svc.cacheMu.RUnlock()
|
||
if ok {
|
||
t.Fatal("cache entry still exists after invalidation")
|
||
}
|
||
}
|
||
|
||
func TestInvalidateStdoutCache_NonExistent(t *testing.T) {
|
||
svc := newTestDeliveryService()
|
||
|
||
// 对不存在的 key 调用 invalidate 不应 panic
|
||
svc.InvalidateStdoutCache("non-existent-job")
|
||
|
||
// 验证缓存为空
|
||
svc.cacheMu.RLock()
|
||
size := len(svc.stdoutCache)
|
||
svc.cacheMu.RUnlock()
|
||
if size != 0 {
|
||
t.Fatalf("cache size = %d, want 0", size)
|
||
}
|
||
}
|
||
|
||
func TestStdoutCache_TTLExpiry(t *testing.T) {
|
||
svc := newTestDeliveryService()
|
||
|
||
// 填充一个已过期的缓存条目
|
||
svc.cacheMu.Lock()
|
||
svc.stdoutCache["job-expired"] = &stdoutCacheItem{
|
||
stdout: "old output",
|
||
createdAt: time.Now().Add(-stdoutCacheTTL - time.Second),
|
||
}
|
||
svc.cacheMu.Unlock()
|
||
|
||
// 模拟 AWXJobStdout 的缓存检查逻辑
|
||
svc.cacheMu.RLock()
|
||
cacheHit := false
|
||
if item, ok := svc.stdoutCache["job-expired"]; ok {
|
||
if time.Since(item.createdAt) < stdoutCacheTTL {
|
||
cacheHit = true
|
||
}
|
||
}
|
||
svc.cacheMu.RUnlock()
|
||
|
||
if cacheHit {
|
||
t.Fatal("expired cache entry should not be a hit")
|
||
}
|
||
}
|
||
|
||
func TestStdoutCache_FreshEntry(t *testing.T) {
|
||
svc := newTestDeliveryService()
|
||
|
||
// 填充一个新鲜的缓存条目
|
||
svc.cacheMu.Lock()
|
||
svc.stdoutCache["job-fresh"] = &stdoutCacheItem{
|
||
stdout: "fresh output",
|
||
createdAt: time.Now(),
|
||
}
|
||
svc.cacheMu.Unlock()
|
||
|
||
// 模拟 AWXJobStdout 的缓存检查逻辑
|
||
svc.cacheMu.RLock()
|
||
cacheHit := false
|
||
var cachedStdout string
|
||
if item, ok := svc.stdoutCache["job-fresh"]; ok {
|
||
if time.Since(item.createdAt) < stdoutCacheTTL {
|
||
cacheHit = true
|
||
cachedStdout = item.stdout
|
||
}
|
||
}
|
||
svc.cacheMu.RUnlock()
|
||
|
||
if !cacheHit {
|
||
t.Fatal("fresh cache entry should be a hit")
|
||
}
|
||
if cachedStdout != "fresh output" {
|
||
t.Fatalf("cached stdout = %q, want %q", cachedStdout, "fresh output")
|
||
}
|
||
}
|
||
|
||
func TestSubscribeTask_ConcurrentSafety(t *testing.T) {
|
||
svc := newTestDeliveryService()
|
||
taskID := "test-concurrent"
|
||
|
||
var wg sync.WaitGroup
|
||
const goroutines = 50
|
||
|
||
// 并发订阅
|
||
cancels := make([]func(), goroutines)
|
||
for i := 0; i < goroutines; i++ {
|
||
wg.Add(1)
|
||
go func(idx int) {
|
||
defer wg.Done()
|
||
_, cancel := svc.SubscribeTask(taskID)
|
||
cancels[idx] = cancel
|
||
}(i)
|
||
}
|
||
wg.Wait()
|
||
|
||
// 验证所有订阅者都已注册
|
||
svc.streamMu.Lock()
|
||
subs := svc.streams[taskID]
|
||
if len(subs) != goroutines {
|
||
svc.streamMu.Unlock()
|
||
t.Fatalf("expected %d subscribers, got %d", goroutines, len(subs))
|
||
}
|
||
svc.streamMu.Unlock()
|
||
|
||
// 并发取消
|
||
for _, cancel := range cancels {
|
||
wg.Add(1)
|
||
go func(c func()) {
|
||
defer wg.Done()
|
||
c()
|
||
}(cancel)
|
||
}
|
||
wg.Wait()
|
||
|
||
// 验证所有订阅者都已移除
|
||
svc.streamMu.Lock()
|
||
if subs := svc.streams[taskID]; subs != nil && len(subs) > 0 {
|
||
svc.streamMu.Unlock()
|
||
t.Fatalf("expected 0 subscribers after concurrent cancel, got %d", len(subs))
|
||
}
|
||
svc.streamMu.Unlock()
|
||
}
|
||
|
||
func TestStdoutCache_ConcurrentAccess(t *testing.T) {
|
||
svc := newTestDeliveryService()
|
||
|
||
var wg sync.WaitGroup
|
||
const goroutines = 50
|
||
|
||
// 并发写入缓存
|
||
for i := 0; i < goroutines; i++ {
|
||
wg.Add(1)
|
||
go func(idx int) {
|
||
defer wg.Done()
|
||
jobID := "job-" + string(rune('A'+idx%26))
|
||
svc.cacheMu.Lock()
|
||
svc.stdoutCache[jobID] = &stdoutCacheItem{
|
||
stdout: "output",
|
||
createdAt: time.Now(),
|
||
}
|
||
svc.cacheMu.Unlock()
|
||
}(i)
|
||
}
|
||
|
||
// 并发读取缓存
|
||
for i := 0; i < goroutines; i++ {
|
||
wg.Add(1)
|
||
go func(idx int) {
|
||
defer wg.Done()
|
||
jobID := "job-" + string(rune('A'+idx%26))
|
||
svc.cacheMu.RLock()
|
||
_ = svc.stdoutCache[jobID]
|
||
svc.cacheMu.RUnlock()
|
||
}(i)
|
||
}
|
||
|
||
// 并发失效缓存
|
||
for i := 0; i < goroutines; i++ {
|
||
wg.Add(1)
|
||
go func(idx int) {
|
||
defer wg.Done()
|
||
jobID := "job-" + string(rune('A'+idx%26))
|
||
svc.InvalidateStdoutCache(jobID)
|
||
}(i)
|
||
}
|
||
|
||
wg.Wait()
|
||
}
|