routewatch/internal/routewatch/asnfetcher.go

package routewatch

import (
	"context"
	"errors"
	"log/slog"
	"sync"
	"time"

	"git.eeqj.de/sneak/routewatch/internal/database"
	"git.eeqj.de/sneak/routewatch/internal/server"
	"git.eeqj.de/sneak/routewatch/internal/whois"
)

// ASN fetcher configuration constants.
const (
	// baseInterval is the starting interval between fetch attempts.
	baseInterval = 15 * time.Second

	// minInterval is the minimum interval after successes (rate limit).
	minInterval = 1 * time.Second

	// maxInterval is the maximum interval after failures (backoff cap).
	maxInterval = 5 * time.Minute

	// backoffMultiplier is how much to multiply interval on failure.
	backoffMultiplier = 2

	// whoisStaleThreshold is how old WHOIS data can be before refresh.
	whoisStaleThreshold = 30 * 24 * time.Hour // 30 days

	// immediateQueueSize is the buffer size for immediate fetch requests.
	immediateQueueSize = 100

	// statsWindow is how long to keep stats for.
	statsWindow = time.Hour
)

// ASNFetcher handles background WHOIS lookups for ASNs.
type ASNFetcher struct {
	db             database.Store
	whoisClient    *whois.Client
	logger         *slog.Logger
	immediateQueue chan int
	stopCh         chan struct{}
	done           chan struct{} // closed when the fetcher goroutine returns

	// fetchMu ensures only one fetch runs at a time
	fetchMu sync.Mutex

	// interval tracking with mutex protection
	intervalMu       sync.Mutex
	currentInterval  time.Duration
	consecutiveFails int

	// hourly stats tracking; now is the clock, time.Now except in tests
	now          func() time.Time
	statsMu      sync.Mutex
	successTimes []time.Time
	errorTimes   []time.Time
}

// NewASNFetcher creates a new ASN fetcher.
func NewASNFetcher(db database.Store, logger *slog.Logger) *ASNFetcher {
	return &ASNFetcher{
		db:              db,
		whoisClient:     whois.NewClient(),
		logger:          logger.With("component", "asn_fetcher"),
		immediateQueue:  make(chan int, immediateQueueSize),
		stopCh:          make(chan struct{}),
		done:            make(chan struct{}),
		currentInterval: baseInterval,
		now:             time.Now,
		successTimes:    make([]time.Time, 0),
		errorTimes:      make([]time.Time, 0),
	}
}

// Start begins the background ASN fetcher goroutine. It runs until Stop:
// the end of ctx does not stop it.
func (f *ASNFetcher) Start(ctx context.Context) {
	go f.run(context.WithoutCancel(ctx))

	f.logger.Info("ASN fetcher started",
		"base_interval", baseInterval,
		"min_interval", minInterval,
		"max_interval", maxInterval,
	)
}

// Stop stops the fetcher and waits, until ctx ends, for a lookup in progress to
// finish.
func (f *ASNFetcher) Stop(ctx context.Context) {
	close(f.stopCh)

	select {
	case <-f.done:
		f.logger.Info("ASN fetcher stopped")
	case <-ctx.Done():
		f.logger.Warn("Stop deadline passed before the ASN fetcher stopped")
	}
}

// QueueImmediate queues an ASN for immediate WHOIS lookup.
// Non-blocking - if queue is full, the request is dropped.
func (f *ASNFetcher) QueueImmediate(asn int) {
	select {
	case f.immediateQueue <- asn:
		f.logger.Debug("Queued immediate WHOIS lookup", "asn", asn)
	default:
		f.logger.Debug("Immediate queue full, dropping request", "asn", asn)
	}
}

// GetStats returns statistics about fetcher activity.
func (f *ASNFetcher) GetStats() server.ASNFetcherStats {
	f.statsMu.Lock()
	defer f.statsMu.Unlock()

	f.intervalMu.Lock()
	interval := f.currentInterval
	fails := f.consecutiveFails
	f.intervalMu.Unlock()

	// Prune old entries and count
	cutoff := f.now().Add(-statsWindow)
	f.successTimes = pruneOldTimes(f.successTimes, cutoff)
	f.errorTimes = pruneOldTimes(f.errorTimes, cutoff)

	return server.ASNFetcherStats{
		SuccessesLastHour: len(f.successTimes),
		ErrorsLastHour:    len(f.errorTimes),
		CurrentInterval:   interval,
		ConsecutiveFails:  fails,
	}
}

// pruneOldTimes drops the times at or before cutoff from the front of times,
// which holds times in the order they were recorded, and returns the rest. It
// copies nothing, so it is cheap enough to run on every append.
func pruneOldTimes(times []time.Time, cutoff time.Time) []time.Time {
	for len(times) > 0 && !times[0].After(cutoff) {
		times = times[1:]
	}

	return times
}

// getInterval returns the current fetch interval.
func (f *ASNFetcher) getInterval() time.Duration {
	f.intervalMu.Lock()
	defer f.intervalMu.Unlock()

	return f.currentInterval
}

// recordSuccess decreases the interval on successful fetch.
func (f *ASNFetcher) recordSuccess() {
	f.intervalMu.Lock()
	f.consecutiveFails = 0

	// Decrease interval by half, but not below minimum
	newInterval := max(f.currentInterval/backoffMultiplier, minInterval)

	if newInterval != f.currentInterval {
		f.logger.Debug("Decreased fetch interval",
			"old_interval", f.currentInterval,
			"new_interval", newInterval,
		)
		f.currentInterval = newInterval
	}
	f.intervalMu.Unlock()

	// Record success time for stats, dropping those older than statsWindow
	f.statsMu.Lock()
	now := f.now()
	f.successTimes = pruneOldTimes(f.successTimes, now.Add(-statsWindow))
	f.successTimes = append(f.successTimes, now)
	f.statsMu.Unlock()
}

// recordFailure increases the interval on failed fetch using exponential backoff.
func (f *ASNFetcher) recordFailure() {
	f.intervalMu.Lock()
	f.consecutiveFails++

	// Exponential backoff: multiply by 2, capped at max
	newInterval := min(f.currentInterval*backoffMultiplier, maxInterval)

	if newInterval != f.currentInterval {
		f.logger.Debug("Increased fetch interval due to failure",
			"old_interval", f.currentInterval,
			"new_interval", newInterval,
			"consecutive_failures", f.consecutiveFails,
		)
		f.currentInterval = newInterval
	}
	f.intervalMu.Unlock()

	// Record error time for stats, dropping those older than statsWindow
	f.statsMu.Lock()
	now := f.now()
	f.errorTimes = pruneOldTimes(f.errorTimes, now.Add(-statsWindow))
	f.errorTimes = append(f.errorTimes, now)
	f.statsMu.Unlock()
}

// run is the main background loop.
func (f *ASNFetcher) run(ctx context.Context) {
	defer close(f.done)

	timer := time.NewTimer(f.getInterval())
	defer timer.Stop()

	for {
		select {
		case <-f.stopCh:
			return

		case asn := <-f.immediateQueue:
			// Process immediate request (respects lock)
			f.tryFetch(ctx, asn)
			// Reset timer after immediate fetch
			timer.Reset(f.getInterval())

		case <-timer.C:
			// Background fetch of stale/missing ASN
			f.fetchNextStale(ctx)
			// Reset timer with potentially updated interval
			timer.Reset(f.getInterval())
		}
	}
}

// tryFetch attempts to fetch and update an ASN, respecting the fetch lock.
// Returns true if fetch was successful.
func (f *ASNFetcher) tryFetch(ctx context.Context, asn int) bool {
	// Try to acquire lock, skip if another fetch is running
	if !f.fetchMu.TryLock() {
		f.logger.Debug("Skipping fetch, another fetch in progress", "asn", asn)

		return false
	}
	defer f.fetchMu.Unlock()

	return f.fetchAndUpdate(ctx, asn)
}

// fetchNextStale finds and fetches the next ASN needing WHOIS data.
func (f *ASNFetcher) fetchNextStale(ctx context.Context) {
	// Try to acquire lock, skip if another fetch is running
	if !f.fetchMu.TryLock() {
		f.logger.Debug("Skipping stale fetch, another fetch in progress")

		return
	}
	defer f.fetchMu.Unlock()

	lookupCtx, cancel := context.WithTimeout(ctx, time.Minute)
	defer cancel()

	asn, err := f.db.GetNextStaleASN(lookupCtx, whoisStaleThreshold)
	if err != nil {
		if !errors.Is(err, database.ErrNoStaleASN) {
			f.logger.Error("Failed to get stale ASN", "error", err)
			f.recordFailure()
		}
		// No stale ASN is not a failure, just nothing to do

		return
	}

	f.fetchAndUpdate(ctx, asn)
}

// fetchAndUpdate performs a WHOIS lookup and updates the database.
// Returns true if successful. A failed lookup leaves the AS's row as it was,
// so the AS stays stale. An AS IANA refers to no RIR is stored with no
// WHOIS data, so it is not looked up again until it is stale; that counts
// as neither a success nor a failure.
func (f *ASNFetcher) fetchAndUpdate(ctx context.Context, asn int) bool {
	ctx, cancel := context.WithTimeout(ctx, time.Minute)
	defer cancel()

	f.logger.Info("Fetching WHOIS data", "asn", asn)

	info, err := f.whoisClient.LookupASN(ctx, asn)
	if errors.Is(err, whois.ErrNoReferral) {
		f.logger.Info("IANA refers the AS to no RIR", "asn", asn)

		err = f.db.UpdateASNWHOIS(ctx, &database.ASNWHOISUpdate{ASN: asn})
		if err != nil {
			f.logger.Error("Failed to update ASN WHOIS data",
				"asn", asn, "error", err)
			f.recordFailure()
		}

		return false
	}

	if err != nil {
		f.logger.Error("WHOIS lookup failed", "asn", asn, "error", err)
		f.recordFailure()

		return false
	}

	// Update database with WHOIS data
	err = f.db.UpdateASNWHOIS(ctx, &database.ASNWHOISUpdate{
		ASN:         asn,
		ASName:      info.ASName,
		OrgName:     info.OrgName,
		OrgID:       info.OrgID,
		Address:     info.Address,
		CountryCode: info.CountryCode,
		AbuseEmail:  info.AbuseEmail,
		AbusePhone:  info.AbusePhone,
		TechEmail:   info.TechEmail,
		TechPhone:   info.TechPhone,
		RIR:         info.RIR,
		RIRRegDate:  info.RegDate,
		RIRLastMod:  info.LastMod,
		WHOISRaw:    info.RawResponse,
	})
	if err != nil {
		f.logger.Error("Failed to update ASN WHOIS data", "asn", asn, "error", err)
		f.recordFailure()

		return false
	}

	f.recordSuccess()
	f.logger.Info("Updated ASN WHOIS data",
		"asn", asn,
		"org_name", info.OrgName,
		"country", info.CountryCode,
		"rir", info.RIR,
		"next_interval", f.getInterval(),
	)

	return true
}

// GetStaleThreshold returns the WHOIS stale threshold duration.
func GetStaleThreshold() time.Duration {
	return whoisStaleThreshold
}