mirror of
https://github.com/netbirdio/netbird.git
synced 2026-07-20 07:21:27 +02:00
added an implementation of aggregating memory store
Signed-off-by: Dmitri <dmitri.external@netbird.io>
This commit is contained in:
@@ -1,7 +1,11 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"maps"
|
||||
"net/netip"
|
||||
"slices"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
@@ -19,6 +23,10 @@ type Memory struct {
|
||||
events map[uuid.UUID]*types.Event
|
||||
}
|
||||
|
||||
type AggregatingMemory struct {
|
||||
Memory
|
||||
}
|
||||
|
||||
func (m *Memory) StoreEvent(event *types.Event) {
|
||||
m.mux.Lock()
|
||||
defer m.mux.Unlock()
|
||||
@@ -48,3 +56,53 @@ func (m *Memory) DeleteEvents(ids []uuid.UUID) {
|
||||
delete(m.events, id)
|
||||
}
|
||||
}
|
||||
|
||||
func (am *AggregatingMemory) StartAggregationWindow() *AggregatingMemory {
|
||||
am.mux.Lock()
|
||||
defer am.mux.Unlock()
|
||||
|
||||
toret := AggregatingMemory{Memory: Memory{events: am.Memory.events}}
|
||||
am.events = make(map[uuid.UUID]*types.Event)
|
||||
|
||||
return &toret
|
||||
}
|
||||
|
||||
type aggregationKey struct {
|
||||
destAddr netip.Addr
|
||||
destPort uint16
|
||||
protocol uint8
|
||||
icmpType uint8
|
||||
ts int64 // used to prevent aggregation on non icmp/udp/tcp events
|
||||
}
|
||||
|
||||
func (am *AggregatingMemory) GetAggregatedEvents() []*types.Event {
|
||||
aggregated := make(map[aggregationKey]*types.Event)
|
||||
for _, v := range am.events {
|
||||
lookupKey := aggregationKey{destAddr: v.DestIP, destPort: v.DestPort, protocol: uint8(v.Protocol), icmpType: v.ICMPCode}
|
||||
if aggregatedEvent, ok := aggregated[lookupKey]; ok {
|
||||
switch aggregatedEvent.Protocol {
|
||||
case types.ICMP, types.ICMPv6, types.UDP, types.TCP:
|
||||
aggregatedEvent.RxBytes += v.RxBytes
|
||||
aggregatedEvent.RxPackets += v.RxPackets
|
||||
aggregatedEvent.TxBytes += v.TxBytes
|
||||
aggregatedEvent.TxPackets += v.TxPackets
|
||||
if aggregatedEvent.Timestamp.Compare(v.Timestamp) < 0 {
|
||||
aggregatedEvent.Timestamp = v.Timestamp
|
||||
}
|
||||
// do we aggregate icmp by code?
|
||||
default:
|
||||
// shouldn't get here
|
||||
}
|
||||
} else {
|
||||
switch v.Protocol {
|
||||
case types.ICMP, types.ICMPv6, types.TCP, types.UDP:
|
||||
aggregated[lookupKey] = v
|
||||
default:
|
||||
lookupKey.ts = time.Now().UnixNano()
|
||||
aggregated[lookupKey] = v
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return slices.Collect(maps.Values(aggregated)) // could return an iterator instead here
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user