EvoBGP/internal/store/memory_telemetry.go

package store

import (
	"context"
	"sort"
	"strings"
	"time"

	"github.com/google/uuid"
)

type telemetryMemoryRow struct {
	snapshot       TelemetrySnapshot
	owner          string
	leaseUntil     time.Time
	token          int64
	lastRequested  time.Time
	refreshPending bool
}

func telemetryKey(tenantID, targetID string) string { return tenantID + "/" + targetID }

func (m *Memory) ensureTelemetryRows(localAuthority string) {
	if m.telemetry == nil {
		m.telemetry = make(map[string]telemetryMemoryRow)
	}
	for _, sp := range m.speakers {
		if sp == nil || !SpeakerNeedsRemoteDispatch(sp.Role, ParseSpeakerMeta(sp.MetaJSON)) {
			continue
		}
		key := telemetryKey(sp.TenantID, sp.ID)
		if _, ok := m.telemetry[key]; !ok {
			m.telemetry[key] = telemetryMemoryRow{snapshot: TelemetrySnapshot{TenantID: sp.TenantID, TargetID: sp.ID, SpeakerID: sp.ID, Kind: "remote_agent"}}
		}
	}
	if localAuthority != "" {
		for tenantID := range m.tenants {
			if tenantID == SystemJobTenantID {
				continue
			}
			targetID := TelemetryLocalTargetID(tenantID, localAuthority)
			key := telemetryKey(tenantID, targetID)
			if _, ok := m.telemetry[key]; !ok {
				m.telemetry[key] = telemetryMemoryRow{snapshot: TelemetrySnapshot{TenantID: tenantID, TargetID: targetID, Kind: "local_bird", LocalAuthority: localAuthority}}
			}
		}
	}
}

// ClaimTelemetry acquires at most limit due targets with a short fencing lease.
func (m *Memory) ClaimTelemetry(_ context.Context, owner string, limit int, lease time.Duration, localAuthority string) ([]TelemetryClaim, error) {
	if strings.TrimSpace(owner) == "" || limit <= 0 {
		return nil, ErrInvalidInput
	}
	if limit > 64 {
		limit = 64
	}
	if lease <= 0 {
		lease = 15 * time.Second
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	m.ensureTelemetryRows(localAuthority)
	keys := make([]string, 0, len(m.telemetry))
	for key := range m.telemetry {
		keys = append(keys, key)
	}
	sort.Slice(keys, func(i, j int) bool {
		a, b := m.telemetry[keys[i]], m.telemetry[keys[j]]
		if !a.snapshot.NextPollAt.Equal(b.snapshot.NextPollAt) {
			return a.snapshot.NextPollAt.Before(b.snapshot.NextPollAt)
		}
		return keys[i] < keys[j]
	})
	now := time.Now().UTC()
	out := make([]TelemetryClaim, 0, limit)
	for _, key := range keys {
		row := m.telemetry[key]
		if row.snapshot.Kind == "local_bird" && row.snapshot.LocalAuthority != localAuthority || row.snapshot.NextPollAt.After(now) || row.leaseUntil.After(now) {
			continue
		}
		var speaker *Speaker
		if row.snapshot.Kind == "remote_agent" {
			sp := m.speakers[row.snapshot.SpeakerID]
			if sp == nil || !SpeakerNeedsRemoteDispatch(sp.Role, ParseSpeakerMeta(sp.MetaJSON)) {
				delete(m.telemetry, key)
				continue
			}
			cp := *sp
			speaker = &cp
		} else {
			for _, id := range sortedSpeakerIDs(m.speakers) {
				sp := m.speakers[id]
				if sp.TenantID == row.snapshot.TenantID && strings.EqualFold(sp.Role, "master") {
					cp := *sp
					speaker = &cp
					break
				}
			}
		}
		row.owner = owner
		row.token++
		row.leaseUntil = now.Add(lease)
		m.telemetry[key] = row
		claim := TelemetryClaim{TelemetrySnapshot: row.snapshot, Speaker: speaker, OwnerID: owner, FencingToken: row.token}
		if speaker != nil {
			claim.SpeakerID = speaker.ID
		}
		claim.ConfigFingerprint = TelemetrySpeakerFingerprint(speaker)
		out = append(out, claim)
		if len(out) == limit {
			break
		}
	}
	return out, nil
}

// CompleteTelemetry commits only the current live lease and preserves last success on failure.
func (m *Memory) CompleteTelemetry(_ context.Context, claim TelemetryClaim, result TelemetryResult, nextPoll time.Time) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	key := telemetryKey(claim.TenantID, claim.TargetID)
	row, ok := m.telemetry[key]
	now := time.Now().UTC()
	if !ok || row.owner != claim.OwnerID || row.token != claim.FencingToken || !row.leaseUntil.After(now) {
		return ErrTelemetryLeaseLost
	}
	if claim.Speaker != nil {
		current := m.speakers[claim.Speaker.ID]
		if TelemetrySpeakerFingerprint(current) != claim.ConfigFingerprint {
			return ErrTelemetryLeaseLost
		}
	}
	if claim.Kind == "local_bird" {
		masterID := ""
		for _, id := range sortedSpeakerIDs(m.speakers) {
			sp := m.speakers[id]
			if sp.TenantID == claim.TenantID && strings.EqualFold(sp.Role, "master") {
				masterID = id
				break
			}
		}
		if masterID != claim.SpeakerID {
			return ErrTelemetryLeaseLost
		}
	}
	if row.snapshot.ConfigFingerprint != claim.ConfigFingerprint {
		row.snapshot.BGP = TelemetryBGP{}
		row.snapshot.Health = TelemetryHealth{}
		row.snapshot.DiscoveryCursor = 0
	}
	row.snapshot.ConfigFingerprint = claim.ConfigFingerprint
	row.snapshot.SpeakerID = claim.SpeakerID
	if result.BGP.Source != "" {
		if result.BGP.Error != "" {
			result.BGP.Sessions = nil
			result.BGP.LastSuccessAt = row.snapshot.BGP.LastSuccessAt
		} else {
			result.BGP.LastSuccessAt = result.BGP.ReceivedAt
		}
		row.snapshot.BGP = result.BGP
	}
	if result.Health.Source != "" {
		if !result.Health.OK {
			result.Health.LastSuccessAt = row.snapshot.Health.LastSuccessAt
		} else {
			result.Health.LastSuccessAt = result.Health.ReceivedAt
		}
		row.snapshot.Health = result.Health
	}
	if row.snapshot.BGP.Error != "" || row.snapshot.Health.Source != "" && !row.snapshot.Health.OK {
		row.snapshot.FailureCount++
	} else {
		row.snapshot.FailureCount = 0
	}
	row.snapshot.NextPollAt = nextPoll.UTC()
	if row.refreshPending {
		row.snapshot.NextPollAt = now
		row.refreshPending = false
	}
	row.owner = ""
	row.leaseUntil = time.Time{}
	if result.BGP.Source != "" && result.BGP.Error == "" {
		row.snapshot.DiscoveryCursor = m.writeTelemetryDiscoveriesLocked(claim, result.BGP, row.snapshot.DiscoveryCursor)
	}
	m.telemetry[key] = row
	return nil
}

func sortedSpeakerIDs(speakers map[string]*Speaker) []string {
	ids := make([]string, 0, len(speakers))
	for id := range speakers {
		ids = append(ids, id)
	}
	sort.Strings(ids)
	return ids
}

func (m *Memory) writeTelemetryDiscoveriesLocked(claim TelemetryClaim, bgp TelemetryBGP, cursor int) int {
	sessions, next := TelemetryDiscoveryBatch(bgp.Sessions, cursor, 256)
	known := make(map[string]struct{})
	for _, p := range m.peers {
		if p == nil || p.TenantID != claim.TenantID {
			continue
		}
		if neighbor, ok := NormalizePeerNeighborString(p.Neighbor); ok {
			known[neighbor] = struct{}{}
		}
	}
	for _, sess := range sessions {
		neighbor, ok := NormalizePeerNeighborString(sess.Neighbor)
		if !ok {
			continue
		}
		if _, ok := known[neighbor]; ok {
			continue
		}
		neighborID := strings.TrimSpace(sess.NeighborID)
		d := m.findPeerDiscoveryLocked(claim.TenantID, neighborID, neighbor, sess.NeighborAS)
		if d == nil {
			id := uuid.NewString()
			m.peerDiscoveries[id] = &BGPPeerDiscovery{ID: id, TenantID: claim.TenantID, SpeakerID: claim.SpeakerID, NeighborID: neighborID, Neighbor: neighbor, RemoteASN: sess.NeighborAS, ProtocolName: sess.Name, SessionState: sess.State, Status: PeerDiscoveryPending, FirstSeenAt: bgp.ReceivedAt, LastSeenAt: bgp.ReceivedAt}
			continue
		}
		if d.Status == PeerDiscoveryRejected || d.LastSeenAt.After(bgp.ReceivedAt) {
			continue
		}
		d.ProtocolName = sess.Name
		d.SessionState = sess.State
		d.LastSeenAt = bgp.ReceivedAt
		if claim.SpeakerID != "" {
			d.SpeakerID = claim.SpeakerID
		}
	}
	return next
}

// ListTelemetry reads the last committed snapshots without polling agents.
func (m *Memory) ListTelemetry(_ context.Context, tenantID string) ([]TelemetrySnapshot, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	out := make([]TelemetrySnapshot, 0)
	for _, row := range m.telemetry {
		if row.snapshot.TenantID != tenantID {
			continue
		}
		if row.snapshot.Kind == "remote_agent" {
			sp := m.speakers[row.snapshot.SpeakerID]
			if sp == nil || !SpeakerNeedsRemoteDispatch(sp.Role, ParseSpeakerMeta(sp.MetaJSON)) {
				continue
			}
		}
		snap := row.snapshot
		snap.BGP.Sessions = append(snap.BGP.Sessions[:0:0], snap.BGP.Sessions...)
		out = append(out, snap)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].TargetID < out[j].TargetID })
	return out, nil
}

// RequestTelemetryRefresh coalesces repeated requests per tenant within cooldown.
func (m *Memory) RequestTelemetryRefresh(_ context.Context, tenantID string, cooldown time.Duration) (bool, error) {
	if cooldown <= 0 {
		cooldown = 10 * time.Second
	}
	m.mu.Lock()
	defer m.mu.Unlock()
	now := time.Now().UTC()
	changed := false
	for key, row := range m.telemetry {
		if row.snapshot.TenantID != tenantID || now.Sub(row.lastRequested) < cooldown {
			continue
		}
		row.lastRequested = now
		if row.leaseUntil.After(now) {
			row.refreshPending = true
		} else {
			row.snapshot.NextPollAt = now
		}
		m.telemetry[key] = row
		changed = true
	}
	return changed, nil
}