respond to feedback

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

client/internal/netflow/manager.go

@@ -249,7 +249,7 @@ func (m *Manager) receiveACKs(ctx context.Context, client *client.GRPCClient) {
 }
 
 func (m *Manager) startRetries(ctx context.Context) {
-	ticker := time.NewTimer(m.retryInterval)
+	timer := time.NewTimer(m.retryInterval)
 	retryBackoff := backoff.WithContext(&backoff.ExponentialBackOff{
 		InitialInterval:     1 * time.Second,
 		RandomizationFactor: 0.5,
@@ -259,13 +259,13 @@ func (m *Manager) startRetries(ctx context.Context) {
 		Stop:                backoff.Stop,
 		Clock:               backoff.SystemClock,
 	}, ctx)
-	defer ticker.Stop()
+	defer timer.Stop()
 
 	for {
 		select {
 		case <-ctx.Done():
 			return
-		case <-ticker.C:
+		case <-timer.C:
 			for _, e := range m.eventsWithoutAcks.GetEvents() {
 				if e.Timestamp.Add(time.Second).After(time.Now()) {
 					// grace period on retries to avoid early retries
@@ -273,12 +273,12 @@ func (m *Manager) startRetries(ctx context.Context) {
 					continue
 				}
 				if err := m.send(e); err != nil {
-					ticker = time.NewTimer(retryBackoff.NextBackOff()) //nolint:staticcheck,wastedassign
+					timer = time.NewTimer(retryBackoff.NextBackOff()) //nolint:staticcheck,wastedassign
 					break
 				}
 			}
 			retryBackoff.Reset()
-			ticker = time.NewTimer(time.Second)
+			timer = time.NewTimer(time.Second)
 		}
 	}
 }

client/internal/netflow/store/event_aggregation_test.go

@@ -17,22 +17,32 @@ func TestFlowAggregation(t *testing.T) {
 	var protocols = []types.Protocol{types.ICMP, types.ICMPv6, types.TCP, types.UDP}
 	var tests = []struct {
 		description string
+		addresses   [][]netip.Addr
+		dstPort     uint16
 		eventTypes  []types.Type
 	}{
 		{
 			description: "start and stop",
+			addresses:   [][]netip.Addr{{ipAddr("1.1.1.1"), ipAddr("2.2.2.2")}, {ipAddr("3.3.3.3"), ipAddr("2.2.2.2")}},
+			dstPort:     uint16(random.Uint32() >> 16),
 			eventTypes:  []types.Type{types.TypeStart, types.TypeEnd},
 		},
 		{
 			description: "start and drop",
+			addresses:   [][]netip.Addr{{ipAddr("1.1.1.1"), ipAddr("2.2.2.2")}, {ipAddr("3.3.3.3"), ipAddr("2.2.2.2")}},
+			dstPort:     uint16(random.Uint32() >> 16),
 			eventTypes:  []types.Type{types.TypeStart, types.TypeDrop},
 		},
 		{
 			description: "start only",
+			addresses:   [][]netip.Addr{{ipAddr("1.1.1.1"), ipAddr("2.2.2.2")}, {ipAddr("3.3.3.3"), ipAddr("2.2.2.2")}},
+			dstPort:     uint16(random.Uint32() >> 16),
 			eventTypes:  []types.Type{types.TypeStart},
 		},
 		{
 			description: "drop only",
+			addresses:   [][]netip.Addr{{ipAddr("1.1.1.1"), ipAddr("2.2.2.2")}, {ipAddr("3.3.3.3"), ipAddr("2.2.2.2")}},
+			dstPort:     uint16(random.Uint32() >> 16),
 			eventTypes:  []types.Type{types.TypeDrop},
 		}}
 
@@ -42,8 +52,8 @@ func TestFlowAggregation(t *testing.T) {
 				store := NewAggregatingMemoryStore()
 				allExpected := make([]*types.Event, 0)
 
-				for i := 0; i < 2; i++ {
+				for _, srcAndDst := range tt.addresses {
-					inEvents, expected := generateEvents(tt.eventTypes, protocol, types.Ingress, 0)
+					inEvents, expected := generateEvents(srcAndDst[0], srcAndDst[1], tt.dstPort, tt.eventTypes, protocol, types.Ingress, 0)
 					for _, e := range inEvents {
 						store.StoreEvent(e)
 					}
@@ -63,22 +73,27 @@ func TestIcmpEventAggregation(t *testing.T) {
 
 	var tests = []struct {
 		description string
+		addresses   [][]netip.Addr
 		eventTypes  []types.Type
 	}{
 		{
 			description: "start and stop",
+			addresses:   [][]netip.Addr{{ipAddr("1.1.1.1"), ipAddr("2.2.2.2")}},
 			eventTypes:  []types.Type{types.TypeStart, types.TypeEnd},
 		},
 		{
 			description: "start and drop",
+			addresses:   [][]netip.Addr{{ipAddr("1.1.1.1"), ipAddr("2.2.2.2")}},
 			eventTypes:  []types.Type{types.TypeStart, types.TypeDrop},
 		},
 		{
 			description: "start only",
+			addresses:   [][]netip.Addr{{ipAddr("1.1.1.1"), ipAddr("2.2.2.2")}},
 			eventTypes:  []types.Type{types.TypeStart},
 		},
 		{
 			description: "drop only",
+			addresses:   [][]netip.Addr{{ipAddr("1.1.1.1"), ipAddr("2.2.2.2")}},
 			eventTypes:  []types.Type{types.TypeDrop},
 		}}
 
@@ -88,7 +103,7 @@ func TestIcmpEventAggregation(t *testing.T) {
 				store := NewAggregatingMemoryStore()
 				allExpected := make([]*types.Event, 0)
 				for _, icmpType := range icmpTypes {
-					events, expected := generateEvents(tt.eventTypes, protocol, types.Ingress, icmpType)
+					events, expected := generateEvents(tt.addresses[0][0], tt.addresses[0][1], 0, tt.eventTypes, protocol, types.Ingress, icmpType)
 					for _, e := range events {
 						store.StoreEvent(e)
 					}
@@ -107,15 +122,12 @@ func ipAddr(a string) netip.Addr {
 	return addr
 }
 
-func generateEvents(eventTypes []types.Type, protocol types.Protocol, direction types.Direction, icmpType uint8) ([]*types.Event, *types.Event) {
+func generateEvents(srcIp, dstIp netip.Addr, dstPort uint16, eventTypes []types.Type, protocol types.Protocol, direction types.Direction, icmpType uint8) ([]*types.Event, *types.Event) {
 	var rxPackets, txPackets, rxBytes, txBytes uint64
 	inEvents := make([]*types.Event, 0)
 	ts := time.Now()
 	flowId := uuid.New()
-	srcIp := ipAddr("1.1.1.1")
 	srcPort := uint16(random.Uint32() >> 16)
-	dstIp := ipAddr("2.2.2.2")
-	dstPort := uint16(random.Uint32() >> 16)
 
 	for idx, eventType := range eventTypes {
 		e := &types.Event{

client/internal/netflow/store/memory.go

@@ -71,17 +71,22 @@ func (am *AggregatingMemory) ResetAggregationWindow() types.FlowEventAggregator
 }
 
 type aggregationKey struct {
-	destAddr netip.Addr
+	srcAddr   netip.Addr
-	destPort uint16
+	destAddr  netip.Addr
-	protocol uint8
+	destPort  uint16
-	icmpType uint8
+	direction int
-	unique   int64 // used to prevent aggregation on non icmp/udp/tcp events
+	protocol  uint8
+	icmpType  uint8
+	unique    int64 // 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{destAddr: v.DestIP, destPort: v.DestPort, protocol: uint8(v.Protocol), icmpType: v.ICMPType}
+		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 {
 			aggregated[lookupKey] = v.Clone()
 			event := aggregated[lookupKey]