refactor: optimize l7 perf cache

Author: yuanchaoa

agent/benches/lru.rs

@@ -21,7 +21,8 @@ use std::{net::Ipv4Addr, time::Instant};
 
 use criterion::*;
 use deepflow_agent::{
-    _L7PerfCache as L7PerfCache, _LogCache as LogCache, _LogMessageType as LogMessageType,
+    _L7PerfCache as L7PerfCache, _LogCache as LogCache, _LogCacheKey as LogCacheKey,
+    _LogMessageType as LogMessageType,
 };
 use lru::LruCache;
 use rand::prelude::*;
@@ -430,7 +431,9 @@ fn rrt_lru(c: &mut Criterion) {
             let start = Instant::now();
             for item in seeds {
                 cache.rrt_cache.put(
-                    ((item.flow_id as u128) << 64) | item.stream_id.unwrap_or_default() as u128,
+                    LogCacheKey(
+                        ((item.flow_id as u128) << 64) | item.stream_id.unwrap_or_default() as u128,
+                    ),
                     LogCache {
                         msg_type: LogMessageType::Request,
                         time: item.duration.as_micros() as u64,
@@ -456,7 +459,9 @@ fn rrt_lru(c: &mut Criterion) {
             let mut cache = L7PerfCache::new(1000);
             for item in &seeds {
                 cache.rrt_cache.put(
-                    ((item.flow_id as u128) << 64) | item.stream_id.unwrap_or_default() as u128,
+                    LogCacheKey(
+                        ((item.flow_id as u128) << 64) | item.stream_id.unwrap_or_default() as u128,
+                    ),
                     LogCache {
                         msg_type: LogMessageType::Request,
                         time: item.duration.as_micros() as u64,
@@ -466,9 +471,9 @@ fn rrt_lru(c: &mut Criterion) {
             }
             let start = Instant::now();
             for item in &seeds {
-                cache.rrt_cache.get(
-                    &(((item.flow_id as u128) << 64) | item.stream_id.unwrap_or_default() as u128),
-                );
+                cache.rrt_cache.get(&LogCacheKey(
+                    ((item.flow_id as u128) << 64) | item.stream_id.unwrap_or_default() as u128,
+                ));
             }
             start.elapsed()
         })
@@ -488,7 +493,9 @@ fn rrt_lru(c: &mut Criterion) {
             let mut cache = L7PerfCache::new(1000);
             for item in &seeds {
                 cache.rrt_cache.put(
-                    ((item.flow_id as u128) << 64) | item.stream_id.unwrap_or_default() as u128,
+                    LogCacheKey(
+                        ((item.flow_id as u128) << 64) | item.stream_id.unwrap_or_default() as u128,
+                    ),
                     LogCache {
                         msg_type: LogMessageType::Request,
                         time: item.duration.as_micros() as u64,
@@ -498,9 +505,9 @@ fn rrt_lru(c: &mut Criterion) {
             }
             let start = Instant::now();
             for item in &seeds {
-                cache.rrt_cache.get(
-                    &(((item.flow_id as u128) << 64) | item.stream_id.unwrap_or_default() as u128),
-                );
+                cache.rrt_cache.get(&LogCacheKey(
+                    ((item.flow_id as u128) << 64) | item.stream_id.unwrap_or_default() as u128,
+                ));
             }
             start.elapsed()
         })

agent/src/common/l7_protocol_info.rs

@@ -22,7 +22,7 @@ use log::{debug, error, warn};
 use serde::Serialize;
 
 use crate::{
-    common::l7_protocol_log::LogCache,
+    common::l7_protocol_log::{LogCache, LogCacheKey},
     flow_generator::{
         protocol_logs::{
             fastcgi::FastCGIInfo, pb_adapter::L7ProtocolSendLog, AmqpInfo, BrpcInfo, DnsInfo,
@@ -35,7 +35,7 @@ use crate::{
     plugin::CustomInfo,
 };
 
-use super::{ebpf::EbpfType, l7_protocol_log::ParseParam};
+use super::l7_protocol_log::ParseParam;
 
 macro_rules! all_protocol_info {
     ($($name:ident($info_struct:ident)),+$(,)?) => {
@@ -133,35 +133,6 @@ pub trait L7ProtocolInfoInterface: Into<L7ProtocolSendLog> {
         (false, false)
     }
 
-    fn cal_cache_key(&self, param: &ParseParam) -> u128 {
-        /*
-            if session id is some: flow id 64bit | 0 32bit | session id 32bit
-            if session id is none: flow id 64bit | packet_seq 64bit
-        */
-        match self.session_id() {
-            Some(sid) => ((param.flow_id as u128) << 64) | sid as u128,
-            None => {
-                ((param.flow_id as u128) << 64)
-                    | (if param.ebpf_type != EbpfType::None {
-                        // NOTE:
-                        //   In the request-log session aggregation process, for eBPF data, we require that requests and
-                        // responses have consecutive cap_seq to ensure the correctness of session aggregation. However,
-                        // when SR (Segmentation-Reassembly) is enabled, we combine multiple eBPF socket event events
-                        // before parsing the protocol. Therefore, in order to ensure that session aggregation can still
-                        // be performed correctly, we need to retain the cap_seq of the last request and the cap_seq of
-                        // the first response, so that the cap_seq of the request and response can still be consecutive.
-                        if param.direction == PacketDirection::ClientToServer {
-                            param.packet_end_seq + 1
-                        } else {
-                            param.packet_start_seq
-                        }
-                    } else {
-                        0
-                    }) as u128
-            }
-        }
-    }
-
     /*
         calculate rrt
         if have previous log cache:
@@ -176,17 +147,15 @@ pub trait L7ProtocolInfoInterface: Into<L7ProtocolSendLog> {
     */
     fn cal_rrt(&self, param: &ParseParam) -> Option<u64> {
         let mut perf_cache = param.l7_perf_cache.borrow_mut();
-        let cache_key = self.cal_cache_key(param);
-        let previous_log_info = perf_cache.rrt_cache.pop(&cache_key);
+        let cache_key = LogCacheKey::new(param, self.session_id());
+        let previous_log_info = perf_cache.pop(cache_key);
 
         let time = param.time;
         let msg_type: LogMessageType = param.direction.into();
         let timeout = param.rrt_timeout as u64;
 
         if time != 0 {
-            let (in_cached_req, timeout_count) = perf_cache
-                .timeout_cache
-                .get_or_insert_mut(param.flow_id, || (0, 0));
+            let (in_cached_req, timeout_count) = perf_cache.get_or_insert_mut(param.flow_id);
 
             let Some(previous_log_info) = previous_log_info else {
                 if msg_type == LogMessageType::Request {
@@ -225,19 +194,10 @@ pub trait L7ProtocolInfoInterface: Into<L7ProtocolSendLog> {
                 // timeout, save the latest
                 if rrt > timeout {
                     *timeout_count += 1;
-                    perf_cache.rrt_cache.put(
-                        cache_key,
-                        LogCache {
-                            msg_type: param.direction.into(),
-                            time: param.time,
-                            multi_merge_info: None,
-                        },
-                    );
                     None
                 } else {
                     Some(rrt)
                 }
-
             // if previous is resp and current is req and previous time gt current time, likely ebpf disorder,
             // calculate the round trip time.
             } else if previous_log_info.msg_type == LogMessageType::Response
@@ -248,9 +208,7 @@ pub trait L7ProtocolInfoInterface: Into<L7ProtocolSendLog> {
                 if rrt > timeout {
                     // disorder info rrt unlikely have large rrt gt timeout
                     warn!("l7 log info disorder with long time rrt {}", rrt);
-                    // timeout, save latest
                     *timeout_count += 1;
-                    perf_cache.rrt_cache.put(cache_key, previous_log_info);
                     None
                 } else {
                     Some(rrt)
@@ -269,15 +227,15 @@ pub trait L7ProtocolInfoInterface: Into<L7ProtocolSendLog> {
                     if previous_log_info.msg_type == LogMessageType::Request {
                         *in_cached_req += 1;
                     }
-                    perf_cache.rrt_cache.put(cache_key, previous_log_info);
+                    perf_cache.put(cache_key, previous_log_info);
                 } else {
                     if previous_log_info.msg_type == LogMessageType::Request {
                         *timeout_count += 1;
                     }
                     if msg_type == LogMessageType::Request {
                         *in_cached_req += 1;
                     }
-                    perf_cache.rrt_cache.put(
+                    perf_cache.put(
                         cache_key,
                         LogCache {
                             msg_type: param.direction.into(),
@@ -307,8 +265,8 @@ pub trait L7ProtocolInfoInterface: Into<L7ProtocolSendLog> {
         assert!(self.session_id().is_some());
 
         let mut perf_cache = param.l7_perf_cache.borrow_mut();
-        let cache_key = self.cal_cache_key(param);
-        let previous_log_info = perf_cache.rrt_cache.pop(&cache_key);
+        let cache_key = LogCacheKey::new(param, self.session_id());
+        let previous_log_info = perf_cache.pop(cache_key);
 
         let time = param.time;
         let msg_type: LogMessageType = param.direction.into();
@@ -319,9 +277,7 @@ pub trait L7ProtocolInfoInterface: Into<L7ProtocolSendLog> {
             return None;
         }
 
-        let (in_cached_req, timeout_count) = perf_cache
-            .timeout_cache
-            .get_or_insert_mut(param.flow_id, || (0, 0));
+        let (in_cached_req, timeout_count) = perf_cache.get_or_insert_mut(param.flow_id);
 
         let (req_end, resp_end) = {
             let (req_end, resp_end) = self.is_req_resp_end();
@@ -407,15 +363,15 @@ pub trait L7ProtocolInfoInterface: Into<L7ProtocolSendLog> {
             };
 
             if put_back {
-                perf_cache.rrt_cache.put(cache_key, previous_log_info);
+                perf_cache.put(cache_key, previous_log_info);
             }
             r
         } else {
             if previous_log_info.msg_type != msg_type && !merged {
                 *timeout_count += 1;
             }
             if put_back {
-                perf_cache.rrt_cache.put(cache_key, previous_log_info);
+                perf_cache.put(cache_key, previous_log_info);
             }
             None
         }

agent/src/common/l7_protocol_log.rs

@@ -299,38 +299,116 @@ pub struct LogCache {
     pub multi_merge_info: Option<(bool, bool, bool)>,
 }
 
+#[derive(Clone, Copy, Eq, Hash, PartialEq)]
+pub struct LogCacheKey(pub u128);
+
+impl LogCacheKey {
+    pub fn new(param: &ParseParam, session_id: Option<u32>) -> Self {
+        /*
+            if session id is some: flow id 64bit | 0 32bit | session id 32bit
+            if session id is none: flow id 64bit | packet_seq 64bit
+        */
+        let key = match session_id {
+            Some(sid) => ((param.flow_id as u128) << 64) | sid as u128,
+            None => {
+                ((param.flow_id as u128) << 64)
+                    | (if param.ebpf_type != EbpfType::None {
+                        // NOTE:
+                        //   In the request-log session aggregation process, for eBPF data, we require that requests and
+                        // responses have consecutive cap_seq to ensure the correctness of session aggregation. However,
+                        // when SR (Segmentation-Reassembly) is enabled, we combine multiple eBPF socket event events
+                        // before parsing the protocol. Therefore, in order to ensure that session aggregation can still
+                        // be performed correctly, we need to retain the cap_seq of the last request and the cap_seq of
+                        // the first response, so that the cap_seq of the request and response can still be consecutive.
+                        if param.direction == PacketDirection::ClientToServer {
+                            param.packet_end_seq + 1
+                        } else {
+                            param.packet_start_seq
+                        }
+                    } else {
+                        0
+                    }) as u128
+            }
+        };
+
+        Self(key)
+    }
+
+    fn flow_id(&self) -> u64 {
+        (self.0 >> 64) as u64
+    }
+}
+
 pub struct L7PerfCache {
     // lru cache previous rrt
-    pub rrt_cache: LruCache<u128, LogCache>,
+    pub rrt_cache: LruCache<LogCacheKey, LogCache>,
     // LruCache<flow_id, (in_cache_req, count)>
     pub timeout_cache: LruCache<u64, (usize, usize)>,
+    // LruCache<flow_id, LruCache<LogCacheKey, bool>>
+    pub flow_id_map: LruCache<u64, LruCache<LogCacheKey, bool>>,
     // time in microseconds
     pub last_log_time: u64,
 }
 
 impl L7PerfCache {
     // 60 seconds
     const LOG_INTERVAL: u64 = 60_000_000;
+    // When the number of concurrent transactions exceeds this value, the RRT calculation error will occur.
+    const MAX_RRT_CACHE_PER_FLOW: usize = 128;
 
     pub fn new(cap: usize) -> Self {
         L7PerfCache {
             rrt_cache: LruCache::new(cap.try_into().unwrap()),
             timeout_cache: LruCache::new(cap.try_into().unwrap()),
+            flow_id_map: LruCache::new(cap.try_into().unwrap()),
             last_log_time: 0,
         }
     }
 
-    pub fn put(&mut self, key: u128, value: LogCache) -> Option<LogCache> {
+    pub fn put(&mut self, key: LogCacheKey, value: LogCache) -> Option<LogCache> {
         let now = value.time;
         if self.rrt_cache.len() >= usize::from(self.rrt_cache.cap())
             && self.last_log_time + Self::LOG_INTERVAL < now
         {
             self.last_log_time = now;
             debug!("The capacity({}) of the rrt table will be exceeded. please adjust the configuration", self.rrt_cache.cap());
         }
+        if let Some((old, _)) = self
+            .flow_id_map
+            .get_or_insert_mut(key.flow_id(), || {
+                let mut cache = LruCache::new(Self::MAX_RRT_CACHE_PER_FLOW.try_into().unwrap());
+                cache.put(key.clone(), true);
+                cache
+            })
+            .push(key.clone(), true)
+        {
+            // Another cache entry is removed due to the lru's capacity.
+            if key != old {
+                self.rrt_cache.pop(&old);
+                if self.last_log_time + Self::LOG_INTERVAL < now {
+                    self.last_log_time = now;
+                    debug!(
+                        "The capacity({}) of the flow id table will be exceeded, flow id: {}",
+                        Self::MAX_RRT_CACHE_PER_FLOW,
+                        old.flow_id()
+                    );
+                }
+            }
+        }
         self.rrt_cache.put(key, value)
     }
 
+    pub fn pop(&mut self, key: LogCacheKey) -> Option<LogCache> {
+        if let Some(cache) = self.flow_id_map.get_mut(&key.flow_id()) {
+            cache.pop(&key);
+
+            if cache.is_empty() {
+                self.flow_id_map.pop(&key.flow_id());
+            }
+        }
+        self.rrt_cache.pop(&key)
+    }
+
     pub fn pop_timeout_count(&mut self, flow_id: &u64, flow_end: bool) -> usize {
         let (in_cache, t) = self.timeout_cache.pop(flow_id).unwrap_or((0, 0));
         if flow_end {
@@ -340,6 +418,22 @@ impl L7PerfCache {
             t
         }
     }
+
+    pub fn get_or_insert_mut(&mut self, flow_id: u64) -> &mut (usize, usize) {
+        self.timeout_cache.get_or_insert_mut(flow_id, || (0, 0))
+    }
+
+    pub fn remove(&mut self, flow_id: &u64) {
+        self.timeout_cache.pop(flow_id);
+
+        let Some(keys) = self.flow_id_map.pop(flow_id) else {
+            return;
+        };
+
+        for (key, _) in keys {
+            self.rrt_cache.pop(&key);
+        }
+    }
 }
 
 pub struct ParseParam<'a> {

agent/src/flow_generator/flow_map.rs

@@ -2061,6 +2061,7 @@ impl FlowMap {
                 collect_stats = true;
             }
         }
+        self.perf_cache.borrow_mut().remove(&flow.flow_id);
 
         self.size -= 1;
         self.stats_counter

agent/src/flow_generator/protocol_logs/sql/mysql.rs

@@ -444,7 +444,8 @@ impl L7ProtocolParserInterface for MysqlLog {
                 }
                 _ => {}
             }
-            if info.msg_type != LogMessageType::Session {
+            if info.msg_type == LogMessageType::Request || info.msg_type == LogMessageType::Response
+            {
                 info.cal_rrt(param).map(|rrt| {
                     info.rrt = rrt;
                     self.perf_stats.as_mut().map(|p| p.update_rrt(rrt));