[client] introduce client-side event aggregation (#6627)

* added an implementation of aggregating memory store Signed-off-by: Dmitri <dmitri.external@netbird.io> * initial support for aggregation of events Signed-off-by: Dmitri <dmitri.external@netbird.io> * added tcp-aggregation test Signed-off-by: Dmitri <dmitri.external@netbird.io> * added manager integration test Signed-off-by: Dmitri <dmitri.external@netbird.io> * added tracking of the number of start-, drop, and end-events in an aggregation window Signed-off-by: Dmitri <dmitri.external@netbird.io> * fixes based on sonarcube checks Signed-off-by: Dmitri <dmitri.external@netbird.io> * regenerated proto files Signed-off-by: Dmitri <dmitri.external@netbird.io> * removed inadvertenly added google proto files Signed-off-by: Dmitri <dmitri.external@netbird.io> * pacifying linter Signed-off-by: Dmitri <dmitri.external@netbird.io> * update test to validate event aggregation over tcp, udp, icmp, and icmpv6 Signed-off-by: Dmitri <dmitri.external@netbird.io> * updated event aggregation test Signed-off-by: Dmitri <dmitri.external@netbird.io> * regenerate protobufs with expected versions of protoc and protoc-gen-go Signed-off-by: Dmitri <dmitri.external@netbird.io> * remove protoc/protoc-gen headers from flow_grpc.pb.go Signed-off-by: Dmitri <dmitri.external@netbird.io> * updated openapi spec Signed-off-by: Dmitri <dmitri.external@netbird.io> * updated openapi NetworkTrafficEvent spec, regenerated types Signed-off-by: Dmitri <dmitri.external@netbird.io> * respond to feedback Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io> * fixed an issue with how we track events that shouldn't be aggregated Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io> * fixed mapping of events to protobuf Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io> * icmp code values in aggregated events do not matter Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io> * regenerate openapi types Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io> * added a comment re: unbounded unacked events Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io> * reset aggregated event type to unknown Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io> * fix event aggregation test Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io> * used the source port of the earliest event Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io> * add tracking of window starts and ends Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io> * updated openapi spec Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io> * reverted changes to generate.sh Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io> * cleanup handling of not-aggregated events + test Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io> * responded to feedback + small fixes Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io> * small fix in a test Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io> * another test Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io> * force setting non-empty rule id on aggregated events Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io> * fixed a couple of issues flagged by coderabbit Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io> * fix spelling Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io> * handle exhausted retry backoffs Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io> --------- Signed-off-by: Dmitri <dmitri.external@netbird.io> Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>
2026-10-20 12:29:09 +02:00 · 2026-07-06 16:04:26 +02:00
parent 91acb8147c
commit 47352e6e45
12 changed files with 1081 additions and 149 deletions
@@ -1,10 +1,15 @@
 package store

 import (
+	"maps"
+	"math/rand"
+	v2 "math/rand/v2"
+	"net/netip"
+	"slices"
 	"sync"
+	"time"

 	"github.com/google/uuid"
-
 	"github.com/netbirdio/netbird/client/internal/netflow/types"
 )

@@ -19,6 +24,13 @@ type Memory struct {
 	events map[uuid.UUID]*types.Event
 }

+type AggregatingMemory struct {
+	Memory
+	WindowStart time.Time
+	WindowEnd   time.Time
+	rnd         *v2.PCG
+}
+
 func (m *Memory) StoreEvent(event *types.Event) {
 	m.mux.Lock()
 	defer m.mux.Unlock()
@@ -48,3 +60,95 @@ func (m *Memory) DeleteEvents(ids []uuid.UUID) {
 		delete(m.events, id)
 	}
 }
+
+func NewAggregatingMemoryStore() *AggregatingMemory {
+	return &AggregatingMemory{WindowStart: time.Now(), Memory: Memory{events: make(map[uuid.UUID]*types.Event)}, rnd: v2.NewPCG(rand.Uint64(), rand.Uint64())}
+}
+
+func (am *AggregatingMemory) ResetAggregationWindow() types.FlowEventAggregator {
+	am.mux.Lock()
+	defer am.mux.Unlock()
+
+	now := time.Now()
+	toret := AggregatingMemory{WindowStart: am.WindowStart, WindowEnd: now, Memory: Memory{events: am.events}, rnd: v2.NewPCG(rand.Uint64(), rand.Uint64())}
+
+	am.events = make(map[uuid.UUID]*types.Event)
+	am.WindowStart = now
+
+	return &toret
+}
+
+type aggregationKey struct {
+	srcAddr   netip.Addr
+	destAddr  netip.Addr
+	destPort  uint16
+	direction int
+	protocol  uint8
+	icmpType  uint8
+	unique    uint64 // used to prevent aggregation on non icmp/udp/tcp events
+}
+
+func (am *AggregatingMemory) GetAggregatedEvents() []*types.Event {
+	am.mux.Lock()
+	defer am.mux.Unlock()
+
+	aggregated := make(map[aggregationKey]*types.Event)
+	for _, v := range am.events {
+		lookupKey := aggregationKey{srcAddr: v.SourceIP, destAddr: v.DestIP, destPort: v.DestPort, direction: int(v.Direction), protocol: uint8(v.Protocol), icmpType: v.ICMPType}
+		if _, ok := aggregated[lookupKey]; !ok {
+			event := v.Clone()
+
+			switch event.Type {
+			case types.TypeStart:
+				event.NumOfStarts += 1
+			case types.TypeDrop:
+				event.NumOfDrops += 1
+			case types.TypeEnd:
+				event.NumOfEnds += 1
+			}
+			event.Type = types.TypeUnknown
+
+			// Please note that ICMPCode field isn't propagated by the manager (see flow/proto/flow.pb.go, FlowFields struct)
+			// so the field value in an icmp event in the "aggregated" doesn't matter
+
+			event.WindowStart = am.WindowStart
+			event.WindowEnd = am.WindowEnd
+
+			if event.Protocol != types.ICMP && event.Protocol != types.ICMPv6 && event.Protocol != types.UDP && event.Protocol != types.TCP {
+				lookupKey.unique = am.rnd.Uint64() // to make the lookup key unique so we don't aggregate on it
+			}
+
+			aggregated[lookupKey] = event
+			continue
+		}
+
+		aggregatedEvent := aggregated[lookupKey]
+		if aggregatedEvent.Protocol != types.ICMP && aggregatedEvent.Protocol != types.ICMPv6 && aggregatedEvent.Protocol != types.UDP && aggregatedEvent.Protocol != types.TCP {
+			continue // we don't aggregate this type of events; shouldn't ever get here
+		}
+
+		// track the number of connections, duration?, open and close events?
+		aggregatedEvent.RxBytes += v.RxBytes
+		aggregatedEvent.RxPackets += v.RxPackets
+		aggregatedEvent.TxBytes += v.TxBytes
+		aggregatedEvent.TxPackets += v.TxPackets
+		switch v.Type {
+		case types.TypeStart:
+			aggregatedEvent.NumOfStarts += 1
+		case types.TypeDrop:
+			aggregatedEvent.NumOfDrops += 1
+		case types.TypeEnd:
+			aggregatedEvent.NumOfEnds += 1
+		}
+		if aggregatedEvent.Timestamp.Compare(v.Timestamp) > 0 {
+			aggregatedEvent.Timestamp = v.Timestamp
+			aggregatedEvent.ID = v.ID
+			aggregatedEvent.SourcePort = v.SourcePort
+		}
+		if len(aggregatedEvent.RuleID) == 0 && len(v.RuleID) != 0 {
+			aggregatedEvent.RuleID = slices.Clone(v.RuleID)
+		}
+	}
+
+	return slices.Collect(maps.Values(aggregated)) // could return an iterator instead here
+}