yoloyolo/aggregate/window.go

package aggregate

import (
	"time"

	"github.com/netsynth/netsynth/classify"
)

// AggregatePcap reads all ClassifiedPackets (with Timestamp set), assigns to time windows
// using packet timestamps per D-01, fills gap windows with empty snapshots per D-02,
// and returns all WindowSnapshots. Calls onSnapshot for each if non-nil (D-07: --verbose).
// Returns synchronously after channel closes (pcap processing is finite).
func AggregatePcap(events <-chan classify.ClassifiedPacket, windowMs int,
	onSnapshot func(classify.WindowSnapshot)) []classify.WindowSnapshot {

	// Drain the events channel into a slice.
	var all []classify.ClassifiedPacket
	for ev := range events {
		all = append(all, ev)
	}
	if len(all) == 0 {
		return nil
	}

	// Find min and max timestamps (handle non-monotonic pcaps per Pitfall 5).
	minTS := all[0].Timestamp
	maxTS := all[0].Timestamp
	for _, ev := range all[1:] {
		if ev.Timestamp.Before(minTS) {
			minTS = ev.Timestamp
		}
		if ev.Timestamp.After(maxTS) {
			maxTS = ev.Timestamp
		}
	}

	// Calculate the maximum window index.
	maxWindowIdx := int(maxTS.Sub(minTS).Milliseconds()) / windowMs

	// Create all snapshots including gaps (D-02: gaps are silent).
	snapshots := make([]classify.WindowSnapshot, maxWindowIdx+1)
	for i := range snapshots {
		snapshots[i] = classify.WindowSnapshot{
			Counts:      make(map[classify.TrafficClass]int64),
			WindowIndex: i,
		}
	}

	// Assign each packet to its window.
	for _, ev := range all {
		idx := int(ev.Timestamp.Sub(minTS).Milliseconds()) / windowMs
		// Clamp to [0, maxWindowIdx] for safety.
		if idx < 0 {
			idx = 0
		}
		if idx > maxWindowIdx {
			idx = maxWindowIdx
		}
		snapshots[idx].Counts[ev.Class]++
		snapshots[idx].TotalPackets++
	}

	// Fire onSnapshot callback for each snapshot in order (D-07: --verbose).
	if onSnapshot != nil {
		for _, s := range snapshots {
			onSnapshot(s)
		}
	}

	return snapshots
}

// DefaultWindowMs is the default time window duration in milliseconds.
const DefaultWindowMs = 500

// Aggregate reads ClassifiedPackets from events, accumulates counts per TrafficClass
// in time windows of windowMs duration, and emits a WindowSnapshot per window.
// When done is closed, flushes the current partial window and closes the output channel.
// If onSnapshot is non-nil, it is called with each snapshot before sending to the output channel
// (used for --verbose stderr output).
func Aggregate(done <-chan struct{}, events <-chan classify.ClassifiedPacket,
	windowMs int, onSnapshot func(classify.WindowSnapshot)) <-chan classify.WindowSnapshot {

	out := make(chan classify.WindowSnapshot, 8)
	go func() {
		defer close(out)
		ticker := time.NewTicker(time.Duration(windowMs) * time.Millisecond)
		defer ticker.Stop()

		counts := make(map[classify.TrafficClass]int64)
		var totalPackets int64
		windowIndex := 0

		flush := func() {
			snap := classify.WindowSnapshot{
				Counts:       counts,
				TotalPackets: totalPackets,
				WindowIndex:  windowIndex,
			}
			if onSnapshot != nil {
				onSnapshot(snap)
			}
			out <- snap
			counts = make(map[classify.TrafficClass]int64)
			totalPackets = 0
			windowIndex++
		}

		for {
			select {
			case <-done:
				flush() // flush final partial window
				return
			case <-ticker.C:
				flush()
			case ev, ok := <-events:
				if !ok {
					flush()
					return
				}
				counts[ev.Class]++
				totalPackets++
			}
		}
	}()
	return out
}