EvoBGP/internal/repository/module_snapshot_atomic_integration_test.go

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)
	}
}