DDI/agent/core/leases.go

// SPDX-License-Identifier: Apache-2.0

package core

import (
	"context"
	"fmt"
	"net/netip"
	"slices"
	"time"
)

// maxLeasesPerBatch bounds one ReportLeases request (about 100 bytes per
// lease, far below the default gRPC message limit of 4 MiB).
const maxLeasesPerBatch = 5000

// leasesTimeout bounds one LeaseSource.Leases call.
const leasesTimeout = 30 * time.Second

// leaseTracker remembers what the source was last told.
type leaseTracker struct {
	// known holds the leases as last reported, by address.
	known map[netip.Addr]Lease
	// needFull forces the next report to be a full snapshot (at start,
	// after a new connection and after a failed report).
	needFull bool
	lastFull time.Time
}

// leaseLoop polls the role for leases and reports changes. A full snapshot
// is sent at start, after each new connection, after a failed report and
// every LeaseFullSync.
func (a *Agent) leaseLoop(ctx context.Context, ls LeaseSource) {
	t := time.NewTicker(a.cfg.LeaseInterval)
	defer t.Stop()
	tr := &leaseTracker{needFull: true}
	for {
		a.pollLeases(ctx, ls, tr)
		select {
		case <-ctx.Done():
			return
		case <-t.C:
		case <-a.leaseCh:
			tr.needFull = true
		}
	}
}

// pollLeases runs one polling round.
func (a *Agent) pollLeases(ctx context.Context, ls LeaseSource, tr *leaseTracker) {
	if !a.roleMu.TryLock() {
		return // an apply is running; try again next round
	}
	lctx, cancel := context.WithTimeout(ctx, leasesTimeout)
	cur, err := ls.Leases(lctx)
	cancel()
	a.roleMu.Unlock()
	if err != nil {
		if ctx.Err() == nil {
			a.log.Warn("cannot read leases", "error", err)
		}
		return
	}
	full := tr.needFull || time.Since(tr.lastFull) >= a.cfg.LeaseFullSync
	changes, next := diffLeases(tr.known, cur)
	toSend := changes
	if full {
		toSend = sortedLeases(next)
	} else if len(changes) == 0 {
		return
	}
	if err := a.sendLeases(ctx, toSend, full); err != nil {
		if ctx.Err() == nil {
			a.log.Warn("cannot report leases; next report is a full snapshot", "error", err)
		}
		tr.needFull = true
		return
	}
	tr.known = next
	tr.needFull = false
	if full {
		tr.lastFull = time.Now()
	}
}

// sendLeases sends leases in batches of maxLeasesPerBatch. For a full
// snapshot only the first batch carries FullSync; the server replaces its
// view with it and applies the following batches as changes, so a snapshot
// split into several requests ends in the same state. If any batch fails,
// the caller sends a new full snapshot later.
func (a *Agent) sendLeases(ctx context.Context, leases []Lease, full bool) error {
	first := true
	for len(leases) > 0 || first {
		n := min(len(leases), maxLeasesPerBatch)
		seq, err := a.nextLeaseSequence()
		if err != nil {
			return err
		}
		b := LeaseBatch{Sequence: seq, FullSync: full && first, Leases: leases[:n]}
		if err := a.source.ReportLeases(ctx, b); err != nil {
			return err
		}
		a.log.Debug("leases reported", "sequence", seq, "full_sync", b.FullSync, "count", n)
		leases = leases[n:]
		first = false
	}
	return nil
}

// nextLeaseSequence increments and persists the lease sequence before it is
// used, so that a restart never reuses a number.
func (a *Agent) nextLeaseSequence() (uint64, error) {
	a.mu.Lock()
	defer a.mu.Unlock()
	a.st.LeaseSequence++
	a.st.UpdatedAt = time.Now().UTC()
	if err := a.store.SaveState(a.st); err != nil {
		return 0, fmt.Errorf("persist lease sequence: %w", err)
	}
	return a.st.LeaseSequence, nil
}

// diffLeases compares the current leases with the known ones. It returns the
// changes (new or modified leases, and known leases that disappeared, marked
// expired) and the new known set.
func diffLeases(known map[netip.Addr]Lease, cur []Lease) ([]Lease, map[netip.Addr]Lease) {
	next := make(map[netip.Addr]Lease, len(cur))
	for _, l := range cur {
		if l.IP.IsValid() {
			next[l.IP] = l
		}
	}
	var changes []Lease
	for ip, l := range next {
		if old, ok := known[ip]; !ok || !leaseEqual(old, l) {
			changes = append(changes, l)
		}
	}
	for ip, l := range known {
		if _, ok := next[ip]; !ok {
			l.State = LeaseExpired
			changes = append(changes, l)
		}
	}
	slices.SortFunc(changes, func(x, y Lease) int { return x.IP.Compare(y.IP) })
	return changes, next
}

func sortedLeases(m map[netip.Addr]Lease) []Lease {
	out := make([]Lease, 0, len(m))
	for _, l := range m {
		out = append(out, l)
	}
	slices.SortFunc(out, func(x, y Lease) int { return x.IP.Compare(y.IP) })
	return out
}

func leaseEqual(x, y Lease) bool {
	return x.IP == y.IP && x.MAC == y.MAC && x.Hostname == y.Hostname && x.State == y.State &&
		x.Expires.Equal(y.Expires) && x.Updated.Equal(y.Updated) && x.SubnetID == y.SubnetID &&
		x.ClientID == y.ClientID
}