pangolin/server/internal/usage/aggregator.go

package usage

import (
	"context"
	"sync"
	"time"

	"github.com/redis/go-redis/v9"
)

// Report is one usage delta pushed by the data plane (#5's ReportUsage path).
// It identifies the user only by data-plane UUID; the aggregator resolves it to
// a user_id internally.  No destination, domain, or DNS data is ever carried.
type Report struct {
	// DPUUID is the data-plane credential UUID the node observed.
	DPUUID string
	// BytesUp / BytesDown are the incremental byte counts since the last report.
	BytesUp   uint64
	BytesDown uint64
	// Minutes is the incremental session minutes since the last report.
	Minutes int
	// At is the event time used for UTC day bucketing.  Zero = time.Now().
	At time.Time
	// BatchID is #5's upload batch identifier, used for replay-idempotent
	// de-duplication.  Empty disables dedup for this report.
	BatchID string
}

// UsageReporter is the callback interface #5 invokes when a node reports usage.
// Aggregator satisfies it, so #5 can depend on the interface and be wired to a
// real or mock implementation.
type UsageReporter interface {
	ReportUsage(ctx context.Context, r Report) error
}

// Aggregator implements UsageReporter: it resolves dp_uuid → user_id (with a
// short in-memory + Redis cache), de-duplicates replayed batches, and folds the
// delta into usage_daily by UTC date.
type Aggregator struct {
	store    *Store
	rdb      *redis.Client
	cache    *dpCache
	batchTTL time.Duration
}

// NewAggregator builds an Aggregator.  rdb may be nil (batch dedup and the
// Redis tier of the dp_uuid cache are then disabled; the in-memory cache and
// DB lookups still work).  dpCacheTTL <= 0 defaults to 5 minutes; batchTTL <= 0
// defaults to 10 minutes.
func NewAggregator(store *Store, rdb *redis.Client, dpCacheTTL, batchTTL time.Duration) *Aggregator {
	if dpCacheTTL <= 0 {
		dpCacheTTL = 5 * time.Minute
	}
	if batchTTL <= 0 {
		batchTTL = 10 * time.Minute
	}
	return &Aggregator{
		store:    store,
		rdb:      rdb,
		cache:    newDPCache(dpCacheTTL),
		batchTTL: batchTTL,
	}
}

// redisKeyBatch returns the Redis key marking a processed upload batch.
func redisKeyBatch(batchID string) string { return "usage:batch:" + batchID }

// redisKeyDP returns the Redis key caching a dp_uuid → user_id resolution.
func redisKeyDP(dpUUID string) string { return "usage:dp:" + dpUUID }

// ReportUsage folds one report into usage_daily.  It is safe for concurrent
// use and idempotent across replays of the same BatchID.
func (a *Aggregator) ReportUsage(ctx context.Context, r Report) error {
	if r.DPUUID == "" {
		return nil
	}
	if r.BytesUp == 0 && r.BytesDown == 0 && r.Minutes == 0 {
		return nil
	}

	// Replay guard: a batch already seen is a no-op.
	if dup, err := a.batchSeen(ctx, r.BatchID); err != nil {
		return err
	} else if dup {
		return nil
	}

	userID, err := a.resolveUser(ctx, r.DPUUID)
	if err != nil {
		return err
	}
	if userID == 0 {
		// Unknown/rotated dp_uuid: drop silently (no log of the mapping).
		return nil
	}

	at := r.At
	if at.IsZero() {
		at = time.Now()
	}
	return a.store.AggregateUsage(ctx, userID, at.UTC(), r.BytesUp, r.BytesDown, r.Minutes)
}

// batchSeen marks the batch as processed and reports whether it was already
// seen.  Returns false when batchID is empty or Redis is unavailable (dedup is
// best-effort; the DB accumulation itself is additive, not idempotent, so #5
// must supply a BatchID + Redis to get exactly-once semantics).
func (a *Aggregator) batchSeen(ctx context.Context, batchID string) (bool, error) {
	if batchID == "" || a.rdb == nil {
		return false, nil
	}
	set, err := a.rdb.SetNX(ctx, redisKeyBatch(batchID), "1", a.batchTTL).Result()
	if err != nil {
		return false, err
	}
	// SetNX returns true when the key was newly set (i.e. not a duplicate).
	return !set, nil
}

// resolveUser maps a dp_uuid to a user_id via the in-memory cache, then Redis,
// then the database (populating both caches on a DB hit).
func (a *Aggregator) resolveUser(ctx context.Context, dpUUID string) (int64, error) {
	if id, ok := a.cache.get(dpUUID); ok {
		return id, nil
	}

	if a.rdb != nil {
		if id, err := a.rdb.Get(ctx, redisKeyDP(dpUUID)).Int64(); err == nil {
			a.cache.set(dpUUID, id)
			return id, nil
		} else if err != redis.Nil {
			return 0, err
		}
	}

	id, err := a.store.LookupUserIDByDPUUID(ctx, dpUUID)
	if err != nil {
		return 0, err
	}
	if id == 0 {
		return 0, nil
	}

	a.cache.set(dpUUID, id)
	if a.rdb != nil {
		_ = a.rdb.Set(ctx, redisKeyDP(dpUUID), id, a.cache.ttl).Err()
	}
	return id, nil
}

// --------------------------------------------------------------------------
// in-memory dp_uuid → user_id cache (short TTL)
// --------------------------------------------------------------------------

type dpCacheEntry struct {
	userID  int64
	expires time.Time
}

type dpCache struct {
	mu  sync.RWMutex
	m   map[string]dpCacheEntry
	ttl time.Duration
}

func newDPCache(ttl time.Duration) *dpCache {
	return &dpCache{m: make(map[string]dpCacheEntry), ttl: ttl}
}

func (c *dpCache) get(dpUUID string) (int64, bool) {
	c.mu.RLock()
	e, ok := c.m[dpUUID]
	c.mu.RUnlock()
	if !ok || time.Now().After(e.expires) {
		return 0, false
	}
	return e.userID, true
}

func (c *dpCache) set(dpUUID string, userID int64) {
	c.mu.Lock()
	c.m[dpUUID] = dpCacheEntry{userID: userID, expires: time.Now().Add(c.ttl)}
	c.mu.Unlock()
}