package repository import ( "context" "errors" "fmt" "sync" "testing" "time" "evobgp/internal/store" "github.com/google/uuid" "github.com/jackc/pgx/v5/pgxpool" ) func snapshotIntegrationModule(t *testing.T, p *Postgres, tenant string) string { t.Helper() id := uuid.NewString() if _, err := p.pool.Exec(context.Background(), `INSERT INTO module(id,tenant_id,type,name) VALUES($1,$2,'CDN_CIDRS',$3)`, id, tenant, id); err != nil { t.Fatal(err) } return id } func TestPostgresModuleSnapshotAtomicRMWMaxPoolOneIntegration(t *testing.T) { p, tenant := queueIntegration(t) module := snapshotIntegrationModule(t, p, tenant) cfg := p.pool.Config() cfg.MaxConns = 1 smallPool, err := pgxpool.NewWithConfig(context.Background(), cfg) if err != nil { t.Fatal(err) } defer smallPool.Close() small := &Postgres{pool: smallPool} var wg sync.WaitGroup errs := make(chan error, 2) for i := 0; i < 2; i++ { wg.Add(1) go func(i int) { defer wg.Done() errs <- small.UpdateModulePrefixSnapshot(context.Background(), tenant, module, func(cur *store.ModulePrefixSnapshot) (*store.ModulePrefixSnapshot, error) { var rows []store.PrefixRow if cur != nil { rows = append(rows, cur.Prefixes...) } rows = append(rows, store.PrefixRow{Prefix: fmt.Sprintf("192.0.%d.0/24", i), Source: fmt.Sprintf("cdn:%d", i)}) return &store.ModulePrefixSnapshot{InputHash: "atomic", Prefixes: rows}, nil }) }(i) } wg.Wait() close(errs) for err := range errs { if err != nil { t.Fatal(err) } } snap, ok, err := small.GetModulePrefixSnapshot(context.Background(), tenant, module) if err != nil || !ok || len(snap.Prefixes) != 2 { t.Fatalf("lost concurrent update: %+v %v %v", snap, ok, err) } } func TestPostgresModuleSnapshotAdvisoryWaitCancellationIntegration(t *testing.T) { p, tenant := queueIntegration(t) module := snapshotIntegrationModule(t, p, tenant) other := snapshotIntegrationModule(t, p, tenant) cfg := p.pool.Config() cfg.MaxConns = 3 testPool, err := pgxpool.NewWithConfig(context.Background(), cfg) if err != nil { t.Fatal(err) } defer testPool.Close() locked := &Postgres{pool: testPool} started := make(chan struct{}) release := make(chan struct{}) firstDone := make(chan error, 1) go func() { firstDone <- locked.UpdateModulePrefixSnapshot(context.Background(), tenant, module, func(*store.ModulePrefixSnapshot) (*store.ModulePrefixSnapshot, error) { close(started) <-release return &store.ModulePrefixSnapshot{InputHash: "first"}, nil }) }() select { case <-started: case <-time.After(5 * time.Second): close(release) t.Fatal("first transaction did not acquire lock") } defer func() { select { case <-release: default: close(release) } }() waitCtx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) defer cancel() err = locked.UpdateModulePrefixSnapshot(waitCtx, tenant, module, func(*store.ModulePrefixSnapshot) (*store.ModulePrefixSnapshot, error) { return &store.ModulePrefixSnapshot{InputHash: "second"}, nil }) if !errors.Is(err, context.DeadlineExceeded) { t.Fatalf("same-key wait did not cancel: %v", err) } otherCtx, otherCancel := context.WithTimeout(context.Background(), 3*time.Second) defer otherCancel() if err := locked.SetModulePrefixSnapshot(otherCtx, tenant, other, "other", nil); err != nil { t.Fatalf("unrelated module blocked: %v", err) } close(release) if err := <-firstDone; err != nil { t.Fatal(err) } } func TestPostgresModuleSnapshotPartialWriteRollbackIntegration(t *testing.T) { p, tenant := queueIntegration(t) module := snapshotIntegrationModule(t, p, tenant) ctx := context.Background() if err := p.SetModulePrefixSnapshot(ctx, tenant, module, "before", []store.PrefixRow{{Prefix: "192.0.2.0/24", Source: "cdn:before"}}); err != nil { t.Fatal(err) } err := p.UpdateModulePrefixSnapshot(ctx, tenant, module, func(cur *store.ModulePrefixSnapshot) (*store.ModulePrefixSnapshot, error) { if cur == nil || len(cur.Prefixes) != 1 { return nil, errors.New("missing prior snapshot") } invalidCommunity := "not-a-uuid" return &store.ModulePrefixSnapshot{InputHash: "after", Prefixes: []store.PrefixRow{{Prefix: "198.51.100.0/24", CommunityID: &invalidCommunity, Source: "cdn:after"}}}, nil }) if err == nil { t.Fatal("invalid copy unexpectedly committed") } snap, ok, err := p.GetModulePrefixSnapshot(ctx, tenant, module) if err != nil || !ok || snap.InputHash != "before" || len(snap.Prefixes) != 1 || snap.Prefixes[0].Prefix != "192.0.2.0/24" { t.Fatalf("partial write visible: %+v %v %v", snap, ok, err) } var hash string if err := p.pool.QueryRow(ctx, `SELECT input_hash FROM module WHERE id=$1`, module).Scan(&hash); err != nil || hash != "before" { t.Fatalf("module input hash changed after rollback: %q %v", hash, err) } } func TestPostgresModuleSnapshotKilledBackendRollsBackIntegration(t *testing.T) { p, tenant := queueIntegration(t) module := snapshotIntegrationModule(t, p, tenant) ctx := context.Background() if err := p.SetModulePrefixSnapshot(ctx, tenant, module, "before", []store.PrefixRow{{Prefix: "192.0.2.0/24"}}); err != nil { t.Fatal(err) } cfg := p.pool.Config() cfg.MaxConns = 1 cfg.ConnConfig.RuntimeParams["application_name"] = "snapshot-kill-" + uuid.NewString() smallPool, err := pgxpool.NewWithConfig(ctx, cfg) if err != nil { t.Fatal(err) } defer smallPool.Close() small := &Postgres{pool: smallPool} killedBackend := false err = small.UpdateModulePrefixSnapshot(ctx, tenant, module, func(*store.ModulePrefixSnapshot) (*store.ModulePrefixSnapshot, error) { var pid int if err := p.pool.QueryRow(ctx, `SELECT pid FROM pg_stat_activity WHERE application_name=$1 AND state='idle in transaction' ORDER BY pid LIMIT 1`, cfg.ConnConfig.RuntimeParams["application_name"]).Scan(&pid); err != nil { return nil, err } var killed bool if err := p.pool.QueryRow(ctx, `SELECT pg_terminate_backend($1)`, pid).Scan(&killed); err != nil || !killed { return nil, fmt.Errorf("terminate snapshot backend: killed=%v err=%w", killed, err) } killedBackend = true return &store.ModulePrefixSnapshot{InputHash: "after", Prefixes: []store.PrefixRow{{Prefix: "198.51.100.0/24"}}}, nil }) if !killedBackend { t.Fatalf("backend kill was not exercised: %v", err) } if err == nil { t.Fatal("killed backend committed snapshot") } verifyCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() snap, ok, err := p.GetModulePrefixSnapshot(verifyCtx, tenant, module) if err != nil || !ok || snap.InputHash != "before" || snap.Prefixes[0].Prefix != "192.0.2.0/24" { t.Fatalf("snapshot changed after backend loss: %+v %v %v", snap, ok, err) } var hash string if err := p.pool.QueryRow(ctx, `SELECT input_hash FROM module WHERE id=$1`, module).Scan(&hash); err != nil || hash != "before" { t.Fatalf("module input hash changed after backend loss: %q %v", hash, err) } }