EvoBGP/internal/repository/telemetry.go

package repository

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"evobgp/internal/store"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
)

const telemetryColumns = `tenant_id::text,target_id,COALESCE(speaker_id::text,''),kind,local_authority,config_fingerprint,bgp_json,health_json,next_poll_at,failure_count,discovery_cursor`
const telemetryClaimColumns = `t.tenant_id::text,t.target_id,COALESCE(t.speaker_id::text,''),t.kind,t.local_authority,t.config_fingerprint,t.bgp_json,t.health_json,t.next_poll_at,t.failure_count,t.discovery_cursor,t.fencing_token`

func scanTelemetry(row pgx.Row) (store.TelemetrySnapshot, error) {
	var s store.TelemetrySnapshot
	var bgp, health []byte
	err := row.Scan(&s.TenantID, &s.TargetID, &s.SpeakerID, &s.Kind, &s.LocalAuthority, &s.ConfigFingerprint, &bgp, &health, &s.NextPollAt, &s.FailureCount, &s.DiscoveryCursor)
	if err != nil {
		return s, err
	}
	if err := json.Unmarshal(bgp, &s.BGP); err != nil {
		return s, fmt.Errorf("repository: decode telemetry BGP: %w", err)
	}
	if err := json.Unmarshal(health, &s.Health); err != nil {
		return s, fmt.Errorf("repository: decode telemetry health: %w", err)
	}
	return s, nil
}

// ClaimTelemetry inserts at most one bounded batch of new targets and leases due rows.
func (p *Postgres) ClaimTelemetry(ctx context.Context, owner string, limit int, lease time.Duration, localAuthority string) ([]store.TelemetryClaim, error) {
	if owner == "" || limit <= 0 {
		return nil, store.ErrInvalidInput
	}
	if limit > 64 {
		limit = 64
	}
	if lease <= 0 {
		lease = 15 * time.Second
	}
	_, err := p.pool.Exec(ctx, `INSERT INTO telemetry_snapshot(tenant_id,target_id,speaker_id,kind)
	 SELECT s.tenant_id,s.id::text,s.id,'remote_agent' FROM bgp_speaker s
	 WHERE s.role<>'master' AND COALESCE(s.meta_json->>'agent_domain','')<>'' AND COALESCE(s.meta_json->>'agent_secret','')<>''
	 AND NOT EXISTS (SELECT 1 FROM telemetry_snapshot t WHERE t.tenant_id=s.tenant_id AND t.target_id=s.id::text)
	 ORDER BY s.id LIMIT $1 ON CONFLICT DO NOTHING`, limit)
	if err != nil {
		return nil, fmt.Errorf("repository: materialize telemetry targets: %w", err)
	}
	if localAuthority != "" {
		_, err = p.pool.Exec(ctx, `INSERT INTO telemetry_snapshot(tenant_id,target_id,kind,local_authority)
		 SELECT t.id,'cp-local/'||t.id::text||'/'||$3,'local_bird',$3 FROM tenant t
		 WHERE t.id<>$2::uuid AND NOT EXISTS (SELECT 1 FROM telemetry_snapshot x WHERE x.tenant_id=t.id AND x.target_id='cp-local/'||t.id::text||'/'||$3)
		 ORDER BY t.id LIMIT $1 ON CONFLICT DO NOTHING`, limit, store.SystemJobTenantID, localAuthority)
		if err != nil {
			return nil, fmt.Errorf("repository: materialize local telemetry targets: %w", err)
		}
	}
	rows, err := p.pool.Query(ctx, `WITH picked AS (
	 SELECT tenant_id,target_id FROM telemetry_snapshot WHERE next_poll_at<=clock_timestamp()
	 AND (lease_until IS NULL OR lease_until<=clock_timestamp()) AND (kind<>'local_bird' OR local_authority=$4)
	 AND (kind='local_bird' OR EXISTS (SELECT 1 FROM bgp_speaker s WHERE s.id=telemetry_snapshot.speaker_id AND s.role<>'master'
	 AND COALESCE(s.meta_json->>'agent_domain','')<>'' AND COALESCE(s.meta_json->>'agent_secret','')<>''))
	 ORDER BY next_poll_at,tenant_id,target_id FOR UPDATE SKIP LOCKED LIMIT $1
	) UPDATE telemetry_snapshot t SET owner_id=$2,lease_until=clock_timestamp()+$3::bigint*interval '1 millisecond',fencing_token=t.fencing_token+1
	 FROM picked WHERE t.tenant_id=picked.tenant_id AND t.target_id=picked.target_id
	 RETURNING `+telemetryClaimColumns, limit, owner, lease.Milliseconds(), localAuthority)
	if err != nil {
		return nil, fmt.Errorf("repository: claim telemetry: %w", err)
	}
	defer rows.Close()
	claims := make([]store.TelemetryClaim, 0, limit)
	for rows.Next() {
		var c store.TelemetryClaim
		var bgp, health []byte
		if err = rows.Scan(&c.TenantID, &c.TargetID, &c.SpeakerID, &c.Kind, &c.LocalAuthority, &c.ConfigFingerprint, &bgp, &health, &c.NextPollAt, &c.FailureCount, &c.DiscoveryCursor, &c.FencingToken); err != nil {
			return nil, fmt.Errorf("repository: scan telemetry claim: %w", err)
		}
		if err = json.Unmarshal(bgp, &c.BGP); err != nil {
			return nil, fmt.Errorf("repository: decode telemetry BGP: %w", err)
		}
		if err = json.Unmarshal(health, &c.Health); err != nil {
			return nil, fmt.Errorf("repository: decode telemetry health: %w", err)
		}
		c.OwnerID = owner
		claims = append(claims, c)
	}
	if err = rows.Err(); err != nil {
		return nil, fmt.Errorf("repository: iterate telemetry claims: %w", err)
	}
	rows.Close()
	for i := range claims {
		c := &claims[i]
		if c.Kind == "remote_agent" {
			c.Speaker, err = p.getTelemetrySpeaker(ctx, c.TenantID, c.SpeakerID)
		} else {
			var spID string
			err = p.pool.QueryRow(ctx, `SELECT id::text FROM bgp_speaker WHERE tenant_id=$1::uuid AND role='master' ORDER BY id LIMIT 1`, c.TenantID).Scan(&spID)
			if err == pgx.ErrNoRows {
				err = nil
			} else if err == nil {
				c.SpeakerID = spID
				c.Speaker, err = p.getTelemetrySpeaker(ctx, c.TenantID, spID)
			}
		}
		if err != nil {
			return nil, fmt.Errorf("repository: resolve telemetry speaker: %w", err)
		}
		c.ConfigFingerprint = store.TelemetrySpeakerFingerprint(c.Speaker)
	}
	return claims, nil
}

func (p *Postgres) getTelemetrySpeaker(ctx context.Context, tenantID, speakerID string) (*store.Speaker, error) {
	var sp store.Speaker
	err := p.pool.QueryRow(ctx, `SELECT id::text,tenant_id::text,role,COALESCE(endpoint,''),COALESCE(meta_json::text,'{}') FROM bgp_speaker WHERE tenant_id=$1::uuid AND id=$2::uuid`, tenantID, speakerID).Scan(&sp.ID, &sp.TenantID, &sp.Role, &sp.Endpoint, &sp.MetaJSON)
	if err != nil {
		return nil, err
	}
	return &sp, nil
}

// CompleteTelemetry records an attempt only while the claim still owns its lease.
func (p *Postgres) CompleteTelemetry(ctx context.Context, claim store.TelemetryClaim, result store.TelemetryResult, nextPoll time.Time) error {
	tx, err := p.pool.Begin(ctx)
	if err != nil {
		return fmt.Errorf("repository: begin telemetry completion: %w", err)
	}
	defer func() { _ = tx.Rollback(ctx) }()
	var oldBGP, oldHealth []byte
	var token int64
	var owner string
	var leaseUntil time.Time
	var pending bool
	var failures int
	var oldFingerprint string
	var discoveryCursor int
	err = tx.QueryRow(ctx, `SELECT bgp_json,health_json,fencing_token,COALESCE(owner_id,''),COALESCE(lease_until,'epoch'::timestamptz),refresh_pending,failure_count,config_fingerprint,discovery_cursor
	 FROM telemetry_snapshot WHERE tenant_id=$1::uuid AND target_id=$2 FOR UPDATE`, claim.TenantID, claim.TargetID).Scan(&oldBGP, &oldHealth, &token, &owner, &leaseUntil, &pending, &failures, &oldFingerprint, &discoveryCursor)
	if err == pgx.ErrNoRows || err == nil && (token != claim.FencingToken || owner != claim.OwnerID) {
		return store.ErrTelemetryLeaseLost
	}
	if err != nil {
		return fmt.Errorf("repository: read telemetry lease: %w", err)
	}
	var dbNow time.Time
	if err = tx.QueryRow(ctx, `SELECT clock_timestamp()`).Scan(&dbNow); err != nil {
		return fmt.Errorf("repository: telemetry clock: %w", err)
	}
	if !leaseUntil.After(dbNow) {
		return store.ErrTelemetryLeaseLost
	}
	if claim.Speaker != nil {
		var current store.Speaker
		err = tx.QueryRow(ctx, `SELECT id::text,role,COALESCE(endpoint,''),COALESCE(meta_json::text,'{}') FROM bgp_speaker WHERE id=$1::uuid AND tenant_id=$2::uuid FOR SHARE`, claim.Speaker.ID, claim.TenantID).Scan(&current.ID, &current.Role, &current.Endpoint, &current.MetaJSON)
		if err == pgx.ErrNoRows || err == nil && store.TelemetrySpeakerFingerprint(&current) != claim.ConfigFingerprint {
			return store.ErrTelemetryLeaseLost
		}
		if err != nil {
			return fmt.Errorf("repository: verify telemetry config: %w", err)
		}
	}
	if claim.Kind == "local_bird" {
		var masterID string
		err = tx.QueryRow(ctx, `SELECT id::text FROM bgp_speaker WHERE tenant_id=$1::uuid AND role='master' ORDER BY id LIMIT 1`, claim.TenantID).Scan(&masterID)
		if err == pgx.ErrNoRows {
			masterID = ""
			err = nil
		}
		if err != nil {
			return fmt.Errorf("repository: verify local telemetry master: %w", err)
		}
		if masterID != claim.SpeakerID {
			return store.ErrTelemetryLeaseLost
		}
	}
	var bgp store.TelemetryBGP
	var health store.TelemetryHealth
	if err = json.Unmarshal(oldBGP, &bgp); err != nil {
		return fmt.Errorf("repository: decode previous telemetry BGP: %w", err)
	}
	if err = json.Unmarshal(oldHealth, &health); err != nil {
		return fmt.Errorf("repository: decode previous telemetry health: %w", err)
	}
	if oldFingerprint != claim.ConfigFingerprint {
		bgp = store.TelemetryBGP{}
		health = store.TelemetryHealth{}
		discoveryCursor = 0
	}
	if result.BGP.Source != "" {
		if result.BGP.Error != "" {
			result.BGP.Sessions = nil
			result.BGP.LastSuccessAt = bgp.LastSuccessAt
		} else {
			result.BGP.LastSuccessAt = result.BGP.ReceivedAt
		}
		bgp = result.BGP
	}
	if result.Health.Source != "" {
		if result.Health.OK {
			result.Health.LastSuccessAt = result.Health.ReceivedAt
		} else {
			result.Health.LastSuccessAt = health.LastSuccessAt
		}
		health = result.Health
	}
	if bgp.Error != "" || health.Source != "" && !health.OK {
		failures++
	} else {
		failures = 0
	}
	if pending {
		nextPoll = dbNow
	}
	bgpJSON, err := json.Marshal(bgp)
	if err != nil {
		return fmt.Errorf("repository: encode telemetry BGP: %w", err)
	}
	healthJSON, err := json.Marshal(health)
	if err != nil {
		return fmt.Errorf("repository: encode telemetry health: %w", err)
	}
	_, err = tx.Exec(ctx, `UPDATE telemetry_snapshot SET bgp_json=$5::jsonb,health_json=$6::jsonb,next_poll_at=$7,failure_count=$8,
	 config_fingerprint=$9,speaker_id=$10::uuid
	 WHERE tenant_id=$1::uuid AND target_id=$2 AND owner_id=$3 AND fencing_token=$4`, claim.TenantID, claim.TargetID, claim.OwnerID, claim.FencingToken, bgpJSON, healthJSON, nextPoll, failures, claim.ConfigFingerprint, nullTelemetryUUID(claim.SpeakerID))
	if err != nil {
		return fmt.Errorf("repository: write telemetry: %w", err)
	}
	if result.BGP.Source != "" && result.BGP.Error == "" {
		if discoveryCursor, err = writeTelemetryDiscoveries(ctx, tx, claim, result.BGP, discoveryCursor); err != nil {
			return err
		}
	}
	tag, err := tx.Exec(ctx, `UPDATE telemetry_snapshot SET owner_id=NULL,lease_until=NULL,refresh_pending=false,discovery_cursor=$5 WHERE tenant_id=$1::uuid AND target_id=$2 AND owner_id=$3 AND fencing_token=$4 AND lease_until>clock_timestamp()`, claim.TenantID, claim.TargetID, claim.OwnerID, claim.FencingToken, discoveryCursor)
	if err != nil {
		return fmt.Errorf("repository: release telemetry lease: %w", err)
	}
	if tag.RowsAffected() != 1 {
		return store.ErrTelemetryLeaseLost
	}
	if err = tx.Commit(ctx); err != nil {
		return fmt.Errorf("repository: commit telemetry: %w", err)
	}
	return nil
}

func writeTelemetryDiscoveries(ctx context.Context, tx pgx.Tx, claim store.TelemetryClaim, bgp store.TelemetryBGP, cursor int) (int, error) {
	sessions, next := store.TelemetryDiscoveryBatch(bgp.Sessions, cursor, 256)
	for _, sess := range sessions {
		neighbor, ok := store.NormalizePeerNeighborString(sess.Neighbor)
		if !ok {
			continue
		}
		var known bool
		if err := tx.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM bgp_peer WHERE tenant_id=$1::uuid AND neighbor=$2::inet)`, claim.TenantID, neighbor).Scan(&known); err != nil {
			return 0, fmt.Errorf("repository: check discovered neighbor: %w", err)
		}
		if known {
			continue
		}
		neighborID := strings.TrimSpace(sess.NeighborID)
		spID := nullTelemetryUUID(claim.SpeakerID)
		seenAt := bgp.ReceivedAt
		_, err := tx.Exec(ctx, `INSERT INTO bgp_peer_discovery(id,tenant_id,speaker_id,neighbor_id,neighbor,remote_asn,protocol_name,session_state,first_seen_at,last_seen_at)
		 VALUES($1,$2::uuid,$3::uuid,$4,$5::inet,$6,$7,$8,$9,$9) ON CONFLICT DO NOTHING`, uuid.NewString(), claim.TenantID, spID, neighborID, neighbor, sess.NeighborAS, sess.Name, sess.State, seenAt)
		if err != nil {
			return 0, fmt.Errorf("repository: insert discovered neighbor: %w", err)
		}
		_, err = tx.Exec(ctx, `UPDATE bgp_peer_discovery SET protocol_name=$5,session_state=$6,last_seen_at=$7,
		 speaker_id=COALESCE($8::uuid,speaker_id),updated_at=clock_timestamp()
		 WHERE tenant_id=$1::uuid AND status<>'rejected' AND last_seen_at<=$7
		 AND ((neighbor_id<>'' AND neighbor_id=$2) OR (neighbor_id='' AND neighbor=$3::inet AND remote_asn=$4))`, claim.TenantID, neighborID, neighbor, sess.NeighborAS, sess.Name, sess.State, seenAt, spID)
		if err != nil {
			return 0, fmt.Errorf("repository: update discovered neighbor: %w", err)
		}
	}
	return next, nil
}

func nullTelemetryUUID(id string) any {
	if strings.TrimSpace(id) == "" {
		return nil
	}
	return id
}

// ListTelemetry returns committed snapshots for one tenant without touching the agents.
func (p *Postgres) ListTelemetry(ctx context.Context, tenantID string) ([]store.TelemetrySnapshot, error) {
	rows, err := p.pool.Query(ctx, `SELECT `+telemetryColumns+` FROM telemetry_snapshot t WHERE tenant_id=$1::uuid AND (kind='local_bird' OR EXISTS
	 (SELECT 1 FROM bgp_speaker s WHERE s.id=t.speaker_id AND s.role<>'master' AND COALESCE(s.meta_json->>'agent_domain','')<>'' AND COALESCE(s.meta_json->>'agent_secret','')<>'')) ORDER BY target_id`, tenantID)
	if err != nil {
		return nil, fmt.Errorf("repository: list telemetry: %w", err)
	}
	defer rows.Close()
	out := make([]store.TelemetrySnapshot, 0)
	for rows.Next() {
		s, err := scanTelemetry(rows)
		if err != nil {
			return nil, err
		}
		out = append(out, s)
	}
	return out, rows.Err()
}

// RequestTelemetryRefresh advances due times once per cooldown window per target.
func (p *Postgres) RequestTelemetryRefresh(ctx context.Context, tenantID string, cooldown time.Duration) (bool, error) {
	if cooldown <= 0 {
		cooldown = 10 * time.Second
	}
	tag, err := p.pool.Exec(ctx, `UPDATE telemetry_snapshot SET last_refresh_requested_at=clock_timestamp(),
	 refresh_pending=CASE WHEN lease_until>clock_timestamp() THEN true ELSE refresh_pending END,
	 next_poll_at=CASE WHEN lease_until>clock_timestamp() THEN next_poll_at ELSE clock_timestamp() END
	 WHERE tenant_id=$1::uuid AND (kind='local_bird' OR speaker_id IS NOT NULL)
	 AND (last_refresh_requested_at IS NULL OR last_refresh_requested_at<=clock_timestamp()-$2::bigint*interval '1 millisecond')`, tenantID, cooldown.Milliseconds())
	if err != nil {
		return false, fmt.Errorf("repository: request telemetry refresh: %w", err)
	}
	return tag.RowsAffected() > 0, nil
}