EvoBGP/internal/repository/deployment.go

package repository

import (
	"context"
	"errors"

	"evobgp/internal/store"
	"github.com/jackc/pgx/v5"
)

const deploymentColumns = `job_id::text,tenant_id::text,speaker_id::text,revision_id::text,generation,request_sequence,state,confirmation,error_message,dispatch_attempts,retry_round,health_received_at,health_revision_id,health_source_at,health_error,health_observation_token`

func scanDeployment(row pgx.Row) (store.DeploymentSpeaker, error) {
	var d store.DeploymentSpeaker
	err := row.Scan(&d.JobID, &d.TenantID, &d.SpeakerID, &d.RevisionID, &d.Generation, &d.RequestSequence, &d.State, &d.Confirmation, &d.Error, &d.DispatchAttempts, &d.RetryRound, &d.HealthReceivedAt, &d.HealthRevisionID, &d.HealthSourceAt, &d.HealthError, &d.ObservationToken)
	return d, err
}

func deploymentFence(ctx context.Context, tx pgx.Tx, f store.JobFence) (int64, error) {
	var seq int64
	err := tx.QueryRow(ctx, `SELECT request_sequence FROM job_audit WHERE id=$1::uuid AND tenant_id=$2::uuid AND owner_id=$3 AND fencing_token=$4 AND status='running' AND NOT cancel_requested AND lease_until>clock_timestamp() FOR UPDATE`, f.JobID, f.TenantID, f.OwnerID, f.Token).Scan(&seq)
	if errors.Is(err, pgx.ErrNoRows) {
		return 0, store.ErrJobLeaseLost
	}
	return seq, err
}

func (p *Postgres) BeginDeploymentSpeaker(ctx context.Context, f store.JobFence, speaker, revision string) (store.DeploymentSpeaker, error) {
	tx, err := p.pool.Begin(ctx)
	if err != nil {
		return store.DeploymentSpeaker{}, err
	}
	defer rollbackWithTimeout(tx)
	seq, err := deploymentFence(ctx, tx, f)
	if err != nil {
		return store.DeploymentSpeaker{}, err
	}
	var generation, currentSeq int64
	var currentJob, currentRevision string
	err = tx.QueryRow(ctx, `SELECT deployment_generation,deployment_request_sequence,COALESCE(deployment_job_id::text,''),COALESCE(published_revision_id::text,'') FROM bgp_speaker WHERE id=$1::uuid AND tenant_id=$2::uuid FOR UPDATE`, speaker, f.TenantID).Scan(&generation, &currentSeq, &currentJob, &currentRevision)
	if err != nil {
		return store.DeploymentSpeaker{}, err
	}
	if f.ExpectedGeneration != 0 && currentJob != f.JobID && (generation != f.ExpectedGeneration || currentRevision != revision) {
		return store.DeploymentSpeaker{}, store.ErrJobLeaseLost
	}
	reserve := func() error {
		tag, e := tx.Exec(ctx, `UPDATE bgp_speaker SET delivery_owner=$2,delivery_token=$3,delivery_until=clock_timestamp()+($4::bigint*interval '1 millisecond') WHERE id=$1::uuid AND (delivery_until IS NULL OR delivery_until<=clock_timestamp())`, speaker, f.OwnerID, f.Token, store.DeploymentDeliveryWindow.Milliseconds())
		if e != nil {
			return e
		}
		if tag.RowsAffected() == 0 {
			return store.ErrDeploymentBusy
		}
		return nil
	}
	existing, err := scanDeployment(tx.QueryRow(ctx, `SELECT `+deploymentColumns+` FROM deployment_speaker WHERE job_id=$1::uuid AND speaker_id=$2::uuid`, f.JobID, speaker))
	if err == nil {
		if currentJob != f.JobID || generation != existing.Generation || existing.RevisionID != revision {
			return store.DeploymentSpeaker{}, store.ErrJobLeaseLost
		}
		if existing.State != "applied" {
			if err = reserve(); err != nil {
				return store.DeploymentSpeaker{}, err
			}
		}
		if _, err = deploymentFence(ctx, tx, f); err != nil {
			return store.DeploymentSpeaker{}, err
		}
		return existing, tx.Commit(ctx)
	}
	if !errors.Is(err, pgx.ErrNoRows) {
		return store.DeploymentSpeaker{}, err
	}
	if f.ExpectedGeneration != 0 {
		seq = currentSeq
	}
	if currentSeq >= seq && f.ExpectedGeneration == 0 {
		return store.DeploymentSpeaker{}, store.ErrJobLeaseLost
	}
	var valid bool
	err = tx.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM config_revision WHERE id=$1::uuid AND tenant_id=$2::uuid)`, revision, f.TenantID).Scan(&valid)
	if err != nil {
		return store.DeploymentSpeaker{}, err
	}
	if !valid {
		return store.DeploymentSpeaker{}, store.ErrNotFound
	}
	if err = reserve(); err != nil {
		return store.DeploymentSpeaker{}, err
	}
	d, err := scanDeployment(tx.QueryRow(ctx, `INSERT INTO deployment_speaker(job_id,tenant_id,speaker_id,revision_id,generation,request_sequence,state,retry_round) VALUES($1::uuid,$2::uuid,$3::uuid,$4::uuid,$5,$6,'published',$7) RETURNING `+deploymentColumns, f.JobID, f.TenantID, speaker, revision, generation+1, seq, f.RetryRound))
	if err != nil {
		return d, err
	}
	_, err = tx.Exec(ctx, `UPDATE bgp_speaker SET deployment_generation=$2,deployment_request_sequence=$3,deployment_job_id=$4::uuid,published_revision_id=$5::uuid,published_at=now(),updated_at=now() WHERE id=$1::uuid`, speaker, d.Generation, seq, f.JobID, revision)
	if err != nil {
		return d, err
	}
	if _, err = deploymentFence(ctx, tx, f); err != nil {
		return store.DeploymentSpeaker{}, err
	}
	return d, tx.Commit(ctx)
}

func (p *Postgres) RecordDeploymentSpeaker(ctx context.Context, f store.JobFence, d store.DeploymentSpeaker, state, confirmation, message string) error {
	if d.JobID != f.JobID || d.TenantID != f.TenantID {
		return store.ErrJobLeaseLost
	}
	if state == "applied" && confirmation != "agent_reported" && confirmation != "local_applied" {
		return store.ErrInvalidInput
	}
	tx, err := p.pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer rollbackWithTimeout(tx)
	if _, err = deploymentFence(ctx, tx, f); err != nil {
		return err
	}
	var speaker string
	err = tx.QueryRow(ctx, `SELECT id::text FROM bgp_speaker WHERE id=$1::uuid AND tenant_id=$2::uuid AND deployment_job_id=$3::uuid AND deployment_generation=$4 AND published_revision_id=$5::uuid AND delivery_owner=$6 AND delivery_token=$7 AND delivery_until>clock_timestamp() FOR UPDATE`, d.SpeakerID, f.TenantID, f.JobID, d.Generation, d.RevisionID, f.OwnerID, f.Token).Scan(&speaker)
	if errors.Is(err, pgx.ErrNoRows) {
		return store.ErrJobLeaseLost
	}
	if err != nil {
		return err
	}
	_, err = tx.Exec(ctx, `UPDATE deployment_speaker SET state=$3,confirmation=$4,error_message=$5,dispatch_attempts=dispatch_attempts+CASE WHEN $3='dispatched' THEN 1 ELSE 0 END,updated_at=now() WHERE job_id=$1::uuid AND speaker_id=$2::uuid AND generation=$6`, f.JobID, speaker, state, confirmation, message, d.Generation)
	if err != nil {
		return err
	}
	if state != "dispatched" && state != "unknown" {
		if _, err = tx.Exec(ctx, `UPDATE bgp_speaker SET delivery_until=NULL,delivery_owner=NULL WHERE id=$1::uuid`, speaker); err != nil {
			return err
		}
	}
	if state == "applied" {
		_, err = tx.Exec(ctx, `UPDATE bgp_speaker SET last_applied_revision_id=$2::uuid,updated_at=now() WHERE id=$1::uuid`, speaker, d.RevisionID)
		if err != nil {
			return err
		}
	}
	if _, err = deploymentFence(ctx, tx, f); err != nil {
		return err
	}
	return tx.Commit(ctx)
}

func (p *Postgres) ListDeploymentSpeakers(ctx context.Context, tenant, job string) ([]store.DeploymentSpeaker, error) {
	rows, err := p.pool.Query(ctx, `SELECT `+deploymentColumns+` FROM deployment_speaker WHERE tenant_id=$1::uuid AND job_id=$2::uuid ORDER BY speaker_id`, tenant, job)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []store.DeploymentSpeaker
	for rows.Next() {
		d, e := scanDeployment(rows)
		if e != nil {
			return nil, e
		}
		out = append(out, d)
	}
	return out, rows.Err()
}

func (p *Postgres) ListPendingDeployments(ctx context.Context, limit int) ([]store.DeploymentSpeaker, error) {
	if limit <= 0 || limit > 100 {
		limit = 100
	}
	rows, err := p.pool.Query(ctx, `SELECT `+deploymentColumns+` FROM deployment_speaker WHERE (job_id,speaker_id) IN (SELECT d.job_id,d.speaker_id FROM deployment_speaker d JOIN bgp_speaker s ON s.id=d.speaker_id AND s.deployment_job_id=d.job_id AND s.deployment_generation=d.generation JOIN job_audit j ON j.id=d.job_id WHERE d.state<>'applied' AND d.retry_round<3 AND j.status IN ('succeeded','failed') AND NOT j.cancel_requested AND (s.delivery_until IS NULL OR s.delivery_until<=clock_timestamp()) AND d.updated_at < clock_timestamp()-((10*(d.retry_round+1))::bigint*interval '1 second')) ORDER BY updated_at LIMIT $1`, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []store.DeploymentSpeaker
	for rows.Next() {
		d, e := scanDeployment(rows)
		if e != nil {
			return nil, e
		}
		out = append(out, d)
	}
	return out, rows.Err()
}

// ListDeploymentObservations rotates through current generations by oldest CP
// receipt time; the bounded batch prevents a large fleet monopolizing a worker.
func (p *Postgres) ListDeploymentObservations(ctx context.Context, limit int) ([]store.DeploymentSpeaker, error) {
	if limit <= 0 || limit > 100 {
		limit = 10
	}
	rows, err := p.pool.Query(ctx, `WITH picked AS (SELECT d.job_id AS picked_job,d.speaker_id AS picked_speaker FROM deployment_speaker d JOIN bgp_speaker s ON s.id=d.speaker_id AND s.deployment_job_id=d.job_id AND s.deployment_generation=d.generation JOIN job_audit j ON j.id=d.job_id WHERE j.status IN ('succeeded','failed') AND NOT j.cancel_requested AND (s.delivery_until IS NULL OR s.delivery_until<=clock_timestamp()) ORDER BY d.health_received_at NULLS FIRST,d.speaker_id FOR UPDATE OF d SKIP LOCKED LIMIT $1) UPDATE deployment_speaker d SET health_observation_token=d.health_observation_token+1 FROM picked WHERE d.job_id=picked.picked_job AND d.speaker_id=picked.picked_speaker RETURNING `+deploymentColumns, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []store.DeploymentSpeaker
	for rows.Next() {
		d, e := scanDeployment(rows)
		if e != nil {
			return nil, e
		}
		out = append(out, d)
	}
	return out, rows.Err()
}

// ObserveDeployment accepts only a health report for the still-current generation
// outside any delivery reservation. It never treats a legacy report as verified.
func (p *Postgres) ObserveDeployment(ctx context.Context, d store.DeploymentSpeaker, o store.DeploymentObservation) error {
	tx, err := p.pool.Begin(ctx)
	if err != nil {
		return err
	}
	defer rollbackWithTimeout(tx)
	var speaker string
	err = tx.QueryRow(ctx, `SELECT id::text FROM bgp_speaker WHERE id=$1::uuid AND tenant_id=$2::uuid AND deployment_job_id=$3::uuid AND deployment_generation=$4 AND published_revision_id=$5::uuid AND (delivery_until IS NULL OR delivery_until<=clock_timestamp()) FOR UPDATE`, d.SpeakerID, d.TenantID, d.JobID, d.Generation, d.RevisionID).Scan(&speaker)
	if errors.Is(err, pgx.ErrNoRows) {
		return store.ErrJobLeaseLost
	}
	if err != nil {
		return err
	}
	tag, err := tx.Exec(ctx, `UPDATE deployment_speaker SET health_observation_token=health_observation_token+1,health_received_at=clock_timestamp(),health_revision_id=$3,health_source_at=$4,health_error=$5,state=CASE WHEN $5='' AND $3<>'' THEN CASE WHEN $3=revision_id::text THEN 'applied' ELSE 'unknown' END ELSE state END,confirmation=CASE WHEN $5='' AND $3<>'' THEN CASE WHEN $3=revision_id::text THEN 'agent_reported' ELSE 'unconfirmed' END ELSE confirmation END WHERE job_id=$1::uuid AND speaker_id=$2::uuid AND health_observation_token=$6`, d.JobID, speaker, o.RevisionID, o.SourceAt, o.Error, d.ObservationToken)
	if err != nil {
		return err
	}
	if tag.RowsAffected() != 1 {
		return store.ErrJobLeaseLost
	}
	if o.Error == "" && o.RevisionID == d.RevisionID {
		if _, err = tx.Exec(ctx, `UPDATE bgp_speaker SET last_applied_revision_id=$2::uuid,updated_at=now() WHERE id=$1::uuid`, speaker, d.RevisionID); err != nil {
			return err
		}
	}
	return tx.Commit(ctx)
}