feat: add traffic flow deduplication with 24h retention and update-in-place for existing flows

Add traffic_flow_key helper to generate unique flow identifier from node/IPs/protocol/ports/decision, implement 24-hour retention cutoff to delete old flows instead of all flows on heartbeat, build existing_flows lookup map from database with composite key matching, update agent_heartbeat to check for existing flows and update bytes/packets/state/observed_at/raw in-place instead of creating duplicates, extend
This commit is contained in:
2026-07-09 21:19:26 +02:00
parent da60155710
commit 0d07349de0
+43 -6
View File
@@ -536,6 +536,19 @@ def raw_flow_decision(flow: TrafficFlow) -> str | None:
return None return None
def traffic_flow_key(node_id: str, raw_flow: dict[str, object]) -> tuple[object, ...]:
decision = str(raw_flow.get("decision") or raw_flow.get("state") or "observed").lower()
return (
node_id,
str(raw_flow.get("source_ip") or ""),
str(raw_flow.get("destination_ip") or ""),
str(raw_flow.get("protocol") or "unknown"),
flow_int(raw_flow.get("source_port"), 0) or None,
flow_int(raw_flow.get("destination_port"), 0) or None,
decision,
)
def policy_read_payload(policy: Policy, deployment_status: dict[str, object] | None = None) -> dict[str, object]: def policy_read_payload(policy: Policy, deployment_status: dict[str, object] | None = None) -> dict[str, object]:
return { return {
"id": policy.id, "id": policy.id,
@@ -1357,16 +1370,39 @@ def agent_heartbeat(payload: AgentHeartbeat, authorization: str | None = Header(
agent.version = payload.version agent.version = payload.version
agent.last_seen_at = datetime.utcnow() agent.last_seen_at = datetime.utcnow()
agent.last_payload = payload.model_dump(mode="json") agent.last_payload = payload.model_dump(mode="json")
for old_flow in db.scalars(select(TrafficFlow).where(TrafficFlow.node_id == node.id)).all(): retention_cutoff = datetime.utcnow() - timedelta(hours=24)
for old_flow in db.scalars(select(TrafficFlow).where(TrafficFlow.node_id == node.id, TrafficFlow.updated_at < retention_cutoff)).all():
db.delete(old_flow) db.delete(old_flow)
existing_flows = {
(
flow.node_id,
flow.source_ip,
flow.destination_ip,
flow.protocol,
flow.source_port,
flow.destination_port,
str(flow.state or "observed").lower(),
): flow
for flow in db.scalars(select(TrafficFlow).where(TrafficFlow.node_id == node.id)).all()
}
for raw_flow in payload.flows[:1000]: for raw_flow in payload.flows[:1000]:
source_ip = str(raw_flow.get("source_ip") or "") source_ip = str(raw_flow.get("source_ip") or "")
destination_ip = str(raw_flow.get("destination_ip") or "") destination_ip = str(raw_flow.get("destination_ip") or "")
if not source_ip or not destination_ip: if not source_ip or not destination_ip:
continue continue
observed_at = payload.collected_at or datetime.utcnow() observed_at = payload.collected_at or datetime.utcnow()
db.add( observed_at = observed_at.replace(tzinfo=None) if observed_at.tzinfo else observed_at
TrafficFlow( key = traffic_flow_key(node.id, raw_flow)
existing = existing_flows.get(key)
if existing:
existing.bytes = flow_int(raw_flow.get("bytes"))
existing.packets = flow_int(raw_flow.get("packets"))
existing.state = str(raw_flow.get("decision") or raw_flow.get("state") or "") or None
existing.observed_at = observed_at
existing.raw = raw_flow
existing.updated_at = datetime.utcnow()
continue
flow = TrafficFlow(
node_id=node.id, node_id=node.id,
source_ip=source_ip, source_ip=source_ip,
destination_ip=destination_ip, destination_ip=destination_ip,
@@ -1376,10 +1412,11 @@ def agent_heartbeat(payload: AgentHeartbeat, authorization: str | None = Header(
bytes=flow_int(raw_flow.get("bytes")), bytes=flow_int(raw_flow.get("bytes")),
packets=flow_int(raw_flow.get("packets")), packets=flow_int(raw_flow.get("packets")),
state=str(raw_flow.get("decision") or raw_flow.get("state") or "") or None, state=str(raw_flow.get("decision") or raw_flow.get("state") or "") or None,
observed_at=observed_at.replace(tzinfo=None) if observed_at.tzinfo else observed_at, observed_at=observed_at,
raw=raw_flow, raw=raw_flow,
) )
) db.add(flow)
existing_flows[key] = flow
commit_or_400(db) commit_or_400(db)
return {"status": "ok", "node_id": node.id} return {"status": "ok", "node_id": node.id}
@@ -1416,7 +1453,7 @@ async def workload_insights(workload_id: str, _: CurrentUser, db: Session = Depe
flows = db.scalars( flows = db.scalars(
select(TrafficFlow) select(TrafficFlow)
.where((TrafficFlow.source_ip.in_(workload_ips)) | (TrafficFlow.destination_ip.in_(workload_ips))) .where((TrafficFlow.source_ip.in_(workload_ips)) | (TrafficFlow.destination_ip.in_(workload_ips)))
.order_by(TrafficFlow.updated_at.desc()) .order_by((TrafficFlow.state == "blocked").desc(), TrafficFlow.updated_at.desc())
.limit(50) .limit(50)
).all() ).all()
for flow in flows: for flow in flows: