test: strengthen concurrent seed coverage
This commit is contained in:
@@ -5,7 +5,6 @@ import (
|
||||
"fmt"
|
||||
"os"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -17,6 +16,15 @@ import (
|
||||
"senlinai-agent/backend/internal/models"
|
||||
)
|
||||
|
||||
type demoSeedRunContextKey struct{}
|
||||
|
||||
type demoSeedRunMarker string
|
||||
|
||||
const (
|
||||
firstDemoSeedRun demoSeedRunMarker = "first"
|
||||
secondDemoSeedRun demoSeedRunMarker = "second"
|
||||
)
|
||||
|
||||
func TestPostgresDemoSeedSerializesConcurrentRunsForSameOwner(t *testing.T) {
|
||||
dsn := os.Getenv("DATABASE_URL")
|
||||
if dsn == "" {
|
||||
@@ -50,71 +58,83 @@ func TestPostgresDemoSeedSerializesConcurrentRunsForSameOwner(t *testing.T) {
|
||||
ProjectID: projects[0].ID, CreatedBy: user.ID, TagID: &uiTag.ID,
|
||||
Title: "智能报表导出功能", Status: "open",
|
||||
}).Error)
|
||||
t.Cleanup(func() { cleanupPostgresDemoSeed(database, user.ID) })
|
||||
|
||||
var ownerLockObserved atomic.Bool
|
||||
var tagMisses atomic.Int32
|
||||
releaseTagMisses := make(chan struct{})
|
||||
firstAtChildQuery := make(chan struct{}, 1)
|
||||
secondOwnerQueryAttempted := make(chan struct{}, 1)
|
||||
secondEnteredChildQuery := make(chan struct{}, 1)
|
||||
releaseFirst := make(chan struct{})
|
||||
var firstChildOnce sync.Once
|
||||
var releaseOnce sync.Once
|
||||
releaseFirstRun := func() { releaseOnce.Do(func() { close(releaseFirst) }) }
|
||||
callbackName := "test:coordinate_concurrent_seed_" + suffix
|
||||
require.NoError(t, database.Callback().Query().After("gorm:query").Register(callbackName, func(tx *gorm.DB) {
|
||||
require.NoError(t, database.Callback().Query().Before("gorm:query").Register(callbackName, func(tx *gorm.DB) {
|
||||
if tx.Statement.Schema == nil {
|
||||
return
|
||||
}
|
||||
marker, ok := tx.Statement.Context.Value(demoSeedRunContextKey{}).(demoSeedRunMarker)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if tx.Statement.Schema.Table == (models.SenlinAgentUser{}).TableName() {
|
||||
if _, ok := tx.Statement.Clauses["FOR"]; ok {
|
||||
ownerLockObserved.Store(true)
|
||||
if _, locking := tx.Statement.Clauses["FOR"]; locking && marker == secondDemoSeedRun {
|
||||
signalDemoSeedStage(secondOwnerQueryAttempted)
|
||||
}
|
||||
return
|
||||
}
|
||||
if ownerLockObserved.Load() || tx.Statement.Schema.Table != (models.SenlinAgentTag{}).TableName() || tx.RowsAffected != 0 {
|
||||
return
|
||||
}
|
||||
if tagMisses.Add(1) >= 2 {
|
||||
releaseOnce.Do(func() { close(releaseTagMisses) })
|
||||
}
|
||||
select {
|
||||
case <-releaseTagMisses:
|
||||
case <-time.After(5 * time.Second):
|
||||
tx.AddError(fmt.Errorf("timed out coordinating concurrent tag misses"))
|
||||
switch marker {
|
||||
case firstDemoSeedRun:
|
||||
firstChildOnce.Do(func() {
|
||||
signalDemoSeedStage(firstAtChildQuery)
|
||||
<-releaseFirst
|
||||
})
|
||||
case secondDemoSeedRun:
|
||||
signalDemoSeedStage(secondEnteredChildQuery)
|
||||
}
|
||||
}))
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
var wait sync.WaitGroup
|
||||
firstResult := make(chan error, 1)
|
||||
secondResult := make(chan error, 1)
|
||||
t.Cleanup(func() {
|
||||
cancel()
|
||||
releaseFirstRun()
|
||||
wait.Wait()
|
||||
database.Callback().Query().Remove(callbackName)
|
||||
cleanupPostgresDemoSeed(database, user.ID)
|
||||
})
|
||||
|
||||
start := make(chan struct{})
|
||||
results := make(chan error, 2)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
var wait sync.WaitGroup
|
||||
for range 2 {
|
||||
wait.Add(1)
|
||||
go func() {
|
||||
defer wait.Done()
|
||||
<-start
|
||||
_, err := Demo(database.WithContext(ctx), DemoOptions{Email: email, DisplayName: "不应覆盖", Password: "new-password"})
|
||||
results <- err
|
||||
}()
|
||||
}
|
||||
close(start)
|
||||
done := make(chan struct{})
|
||||
wait.Add(1)
|
||||
go func() {
|
||||
wait.Wait()
|
||||
close(done)
|
||||
defer wait.Done()
|
||||
firstContext := context.WithValue(ctx, demoSeedRunContextKey{}, firstDemoSeedRun)
|
||||
_, err := Demo(database.WithContext(firstContext), DemoOptions{Email: email, DisplayName: "不应覆盖", Password: "new-password"})
|
||||
firstResult <- err
|
||||
}()
|
||||
select {
|
||||
case <-done:
|
||||
case <-ctx.Done():
|
||||
<-done
|
||||
t.Fatal("timed out waiting for concurrent Demo calls")
|
||||
}
|
||||
close(results)
|
||||
for err := range results {
|
||||
require.NoError(t, err)
|
||||
}
|
||||
waitForDemoSeedStage(t, ctx, firstAtChildQuery, "first Demo to hold the owner lock and reach its first child query")
|
||||
|
||||
wait.Add(1)
|
||||
go func() {
|
||||
defer wait.Done()
|
||||
secondContext := context.WithValue(ctx, demoSeedRunContextKey{}, secondDemoSeedRun)
|
||||
_, err := Demo(database.WithContext(secondContext), DemoOptions{Email: email, DisplayName: "不应覆盖", Password: "new-password"})
|
||||
secondResult <- err
|
||||
}()
|
||||
waitForDemoSeedStage(t, ctx, secondOwnerQueryAttempted, "second Demo to attempt the owner FOR UPDATE query")
|
||||
|
||||
select {
|
||||
case <-secondEnteredChildQuery:
|
||||
t.Fatal("second Demo entered a child query before the first released the owner lock")
|
||||
case err := <-secondResult:
|
||||
t.Fatalf("second Demo returned before the first released the owner lock: %v", err)
|
||||
case <-time.After(200 * time.Millisecond):
|
||||
case <-ctx.Done():
|
||||
t.Fatal("timed out while proving the second Demo is blocked on the owner lock")
|
||||
}
|
||||
releaseFirstRun()
|
||||
require.NoError(t, waitForDemoSeedResult(t, ctx, firstResult, "first Demo result"))
|
||||
require.NoError(t, waitForDemoSeedResult(t, ctx, secondResult, "second Demo result"))
|
||||
|
||||
require.True(t, ownerLockObserved.Load(), "seed must lock the owner user row before creating child objects")
|
||||
var storedUser models.SenlinAgentUser
|
||||
require.NoError(t, database.Where("email = ?", email).First(&storedUser).Error)
|
||||
require.Equal(t, "保留用户名称", storedUser.DisplayName)
|
||||
@@ -133,6 +153,33 @@ func TestPostgresDemoSeedSerializesConcurrentRunsForSameOwner(t *testing.T) {
|
||||
assertPostgresDemoCount(t, database, &models.SenlinAgentProjectChannel{}, "project_id = ?", []any{projects[0].ID}, 2)
|
||||
}
|
||||
|
||||
func signalDemoSeedStage(stage chan<- struct{}) {
|
||||
select {
|
||||
case stage <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func waitForDemoSeedStage(t *testing.T, ctx context.Context, stage <-chan struct{}, description string) {
|
||||
t.Helper()
|
||||
select {
|
||||
case <-stage:
|
||||
case <-ctx.Done():
|
||||
t.Fatalf("timed out waiting for %s", description)
|
||||
}
|
||||
}
|
||||
|
||||
func waitForDemoSeedResult(t *testing.T, ctx context.Context, result <-chan error, description string) error {
|
||||
t.Helper()
|
||||
select {
|
||||
case err := <-result:
|
||||
return err
|
||||
case <-ctx.Done():
|
||||
t.Fatalf("timed out waiting for %s", description)
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
func assertPostgresDemoCount(t *testing.T, database *gorm.DB, model any, where string, args []any, expected int64) {
|
||||
t.Helper()
|
||||
var count int64
|
||||
|
||||
Reference in New Issue
Block a user