Rogeeandrogee
f36606a4f2
* fix(HH-564): harden durable realtime enqueue * fix(HH-564): wire production SSE stream --------- Co-authored-by: Rogee <rogee@ipao.vip>
98 lines
2.9 KiB
Go
98 lines
2.9 KiB
Go
package ws
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"os"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/driver/postgres"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/logger"
|
|
|
|
"github.com/gochat/gochat/internal/model"
|
|
"github.com/gochat/gochat/internal/worker"
|
|
)
|
|
|
|
func TestEventPublisherConcurrentTransactionalEnqueueIsIdempotentPostgres(t *testing.T) {
|
|
if os.Getenv("GOCHAT_TEST_DB") == "sqlite" {
|
|
t.Skip("PostgreSQL-only concurrency regression")
|
|
}
|
|
|
|
dsn := os.Getenv("GOCHAT_TEST_DB_URL")
|
|
if dsn == "" {
|
|
dsn = "host=localhost port=5432 user=postgres password=postgres dbname=gochat_test sslmode=disable"
|
|
}
|
|
db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
|
|
require.NoError(t, err)
|
|
sqlDB, err := db.DB()
|
|
require.NoError(t, err)
|
|
t.Cleanup(func() { require.NoError(t, sqlDB.Close()) })
|
|
|
|
schema := fmt.Sprintf("realtime_enqueue_%d", time.Now().UnixNano())
|
|
require.NoError(t, db.Exec(`CREATE SCHEMA "`+schema+`"`).Error)
|
|
t.Cleanup(func() { require.NoError(t, db.Exec(`DROP SCHEMA "`+schema+`" CASCADE`).Error) })
|
|
inSchema := func(tx *gorm.DB) error {
|
|
return tx.Exec(`SET LOCAL search_path TO "` + schema + `", public`).Error
|
|
}
|
|
require.NoError(t, db.Transaction(func(tx *gorm.DB) error {
|
|
if err := inSchema(tx); err != nil {
|
|
return err
|
|
}
|
|
if err := tx.AutoMigrate(&model.BackgroundJob{}); err != nil {
|
|
return err
|
|
}
|
|
return tx.Exec(`CREATE UNIQUE INDEX idx_background_jobs_idempotency_key_unique ON background_jobs(idempotency_key) WHERE idempotency_key IS NOT NULL AND idempotency_key <> ''`).Error
|
|
}))
|
|
|
|
pool := worker.NewWorkerPool(db)
|
|
publisher := NewEventPublisherLocal(nil, nil)
|
|
publisher.SetWorkerPool(pool)
|
|
|
|
const writers = 8
|
|
barrier := make(chan struct{})
|
|
var arrived atomic.Int32
|
|
require.NoError(t, db.Callback().Create().Before("gorm:create").Register("test:concurrent_realtime_enqueue", func(tx *gorm.DB) {
|
|
if tx.Statement.Schema != nil && tx.Statement.Schema.Table == "background_jobs" && arrived.Add(1) == writers {
|
|
close(barrier)
|
|
}
|
|
<-barrier
|
|
}))
|
|
start := make(chan struct{})
|
|
errs := make(chan error, writers)
|
|
var wg sync.WaitGroup
|
|
for range writers {
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
<-start
|
|
errs <- db.Transaction(func(tx *gorm.DB) error {
|
|
if err := inSchema(tx); err != nil {
|
|
return err
|
|
}
|
|
_, err := publisher.EnqueueInTransaction(context.Background(), tx, 1, "visitor", EventMessageCreated, map[string]any{"id": 7})
|
|
return err
|
|
})
|
|
}()
|
|
}
|
|
close(start)
|
|
wg.Wait()
|
|
close(errs)
|
|
for err := range errs {
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
var count int64
|
|
require.NoError(t, db.Transaction(func(tx *gorm.DB) error {
|
|
if err := inSchema(tx); err != nil {
|
|
return err
|
|
}
|
|
return tx.Model(&model.BackgroundJob{}).Where("job_type = ?", taskTypeRealtimeEventPublish).Count(&count).Error
|
|
}))
|
|
require.Equal(t, int64(2), count, "one account and one token job must survive concurrent duplicate enqueue")
|
|
}
|