feat: add config ingester_traffic_overflow_action

and improve traffic overflow

Author: lzf575

agent/crates/public/src/queue/overwrite_queue.rs

@@ -55,6 +55,8 @@ struct OverwriteQueue<T: Sized> {
 
     counter: Counter,
 
+    total_overwritten_count: AtomicU64,
+
     _marker: PhantomData<T>,
 }
 
@@ -78,6 +80,7 @@ impl<T> OverwriteQueue<T> {
             notify: Condvar::new(),
             terminated: AtomicBool::new(false),
             counter: Counter::default(),
+            total_overwritten_count: AtomicU64::new(0),
             _marker: PhantomData,
         }
     }
@@ -130,6 +133,8 @@ impl<T> OverwriteQueue<T> {
                 self.counter
                     .overwritten
                     .fetch_add(to_overwrite as u64, Ordering::Relaxed);
+                self.total_overwritten_count
+                    .fetch_add(to_overwrite as u64, Ordering::Relaxed);
             }
         }
         let free_after_end = self.size - (raw_end & (self.size - 1));
@@ -416,6 +421,13 @@ impl<T> Receiver<T> {
             }
         }
     }
+
+    pub fn total_overwritten_count(&self) -> u64 {
+        self.counter()
+            .queue
+            .total_overwritten_count
+            .load(Ordering::Relaxed)
+    }
 }
 
 impl<T> Drop for Receiver<T> {

agent/src/config/config.rs

@@ -2121,6 +2121,28 @@ impl Default for Ntp {
     }
 }
 
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
+#[repr(u8)]
+pub enum TrafficOverflowAction {
+    #[default]
+    Waiting = 0,
+    Dropping = 1,
+}
+
+fn to_traffic_overflow_action<'de, D>(deserializer: D) -> Result<TrafficOverflowAction, D::Error>
+where
+    D: Deserializer<'de>,
+{
+    match u8::deserialize(deserializer)? {
+        0 => Ok(TrafficOverflowAction::Waiting),
+        1 => Ok(TrafficOverflowAction::Dropping),
+        other => Err(de::Error::invalid_value(
+            Unexpected::Unsigned(other as u64),
+            &"[0-1]",
+        )),
+    }
+}
+
 #[derive(Clone, Debug, Deserialize, PartialEq, Eq)]
 #[serde(default)]
 pub struct Communication {
@@ -2133,6 +2155,8 @@ pub struct Communication {
     #[serde(deserialize_with = "deser_usize_with_mega_unit")]
     pub grpc_buffer_size: usize,
     pub max_throughput_to_ingester: u64,
+    #[serde(deserialize_with = "to_traffic_overflow_action")]
+    pub ingester_traffic_overflow_action: TrafficOverflowAction,
     pub request_via_nat_ip: bool,
     pub proxy_controller_ip: String,
     pub proxy_controller_port: u16,
@@ -2149,6 +2173,7 @@ impl Default for Communication {
             ingester_port: 30033,
             grpc_buffer_size: 5 << 20,
             max_throughput_to_ingester: 100,
+            ingester_traffic_overflow_action: TrafficOverflowAction::Waiting,
             request_via_nat_ip: false,
         }
     }

agent/src/config/handler.rs

@@ -51,7 +51,7 @@ use super::{
         HttpEndpointMatchRule, OracleConfig, PcapStream, PortConfig, ProcessorsFlowLogTunning,
         RequestLogTunning, SessionTimeout, TagFilterOperator, UserConfig,
     },
-    ConfigError, KubernetesPollerType,
+    ConfigError, KubernetesPollerType, TrafficOverflowAction,
 };
 use crate::flow_generator::protocol_logs::decode_new_rpc_trace_context_with_type;
 use crate::rpc::Session;
@@ -232,6 +232,7 @@ pub struct SenderConfig {
     pub npb_socket_type: agent::SocketType,
     pub multiple_sockets_to_ingester: bool,
     pub max_throughput_to_ingester: u64, // unit: Mbps
+    pub ingester_traffic_overflow_action: TrafficOverflowAction,
     pub collector_socket_type: agent::SocketType,
     pub standalone_data_file_size: u32,
     pub standalone_data_file_dir: String,
@@ -1748,6 +1749,10 @@ impl TryFrom<(Config, UserConfig)> for ModuleConfig {
                     .throughput_monitoring_interval,
                 multiple_sockets_to_ingester: conf.outputs.socket.multiple_sockets_to_ingester,
                 max_throughput_to_ingester: conf.global.communication.max_throughput_to_ingester,
+                ingester_traffic_overflow_action: conf
+                    .global
+                    .communication
+                    .ingester_traffic_overflow_action,
                 collector_socket_type: conf.outputs.socket.data_socket_type,
                 standalone_data_file_size: conf.global.standalone_mode.max_data_file_size,
                 standalone_data_file_dir: conf.global.standalone_mode.data_file_dir.clone(),
@@ -3875,6 +3880,17 @@ impl ConfigHandler {
             );
             communication.max_throughput_to_ingester = new_communication.max_throughput_to_ingester;
         }
+        if communication.ingester_traffic_overflow_action
+            != new_communication.ingester_traffic_overflow_action
+        {
+            info!(
+                "Update global.communication.ingester_traffic_overflow_action from {:?} to {:?}.",
+                communication.ingester_traffic_overflow_action,
+                new_communication.ingester_traffic_overflow_action
+            );
+            communication.ingester_traffic_overflow_action =
+                new_communication.ingester_traffic_overflow_action;
+        }
         if communication.ingester_ip != new_communication.ingester_ip {
             info!(
                 "Update global.communication.ingester_ip from {:?} to {:?}.",

agent/src/config/mod.rs

@@ -19,7 +19,7 @@ pub mod handler;
 
 pub use config::{
     AgentIdType, Config, ConfigError, DpdkSource, KubernetesPollerType, OracleConfig, PcapStream,
-    PrometheusExtraLabels, UserConfig, K8S_CA_CRT_PATH,
+    PrometheusExtraLabels, TrafficOverflowAction, UserConfig, K8S_CA_CRT_PATH,
 };
 #[cfg(any(target_os = "linux", target_os = "android"))]
 pub use config::{ApiResources, ProcessMatcher};

agent/src/sender/uniform_sender.rs

@@ -30,12 +30,15 @@ use std::time::{Duration, Instant, SystemTime};
 use arc_swap::access::Access;
 use lazy_static::lazy_static;
 use log::{debug, error, info, warn};
-use public::sender::{SendMessageType, Sendable};
+use public::{
+    leaky_bucket::LeakyBucket,
+    sender::{SendMessageType, Sendable},
+};
 use rand::{thread_rng, RngCore};
 
 use super::{get_sender_id, QUEUE_BATCH_SIZE};
 
-use crate::config::handler::SenderAccess;
+use crate::config::{handler::SenderAccess, TrafficOverflowAction};
 use crate::exception::ExceptionHandler;
 use crate::trident::SenderEncoder;
 use crate::utils::stats::{
@@ -53,6 +56,7 @@ pub struct SenderCounter {
     pub tx: AtomicU64,
     pub tx_bytes: AtomicU64,
     pub dropped: AtomicU64,
+    pub waited: AtomicU64,
 }
 
 impl RefCountable for SenderCounter {
@@ -89,6 +93,11 @@ impl RefCountable for SenderCounter {
                 CounterType::Counted,
                 CounterValue::Unsigned(self.dropped.swap(0, Ordering::Relaxed)),
             ),
+            (
+                "waited",
+                CounterType::Counted,
+                CounterValue::Unsigned(self.waited.swap(0, Ordering::Relaxed)),
+            ),
         ]
     }
 }
@@ -153,10 +162,10 @@ impl<T: Sendable> Encoder<T> {
                 version: HEADER_VESION,
                 team_id: 0,
                 organization_id: 0,
-                agent_id: agent_id,
+                agent_id,
                 reserved_1: 0,
                 reserved_2: 0,
-                encoder: encoder,
+                encoder,
             },
             _marker: PhantomData,
         }
@@ -340,6 +349,7 @@ lazy_static! {
     static ref TOTAL_SENT_BYTES: Arc<AtomicU64> = Arc::new(AtomicU64::new(0));
     static ref SENT_START_DURATION: Arc<AtomicU64> = Arc::new(AtomicU64::new(0));
     static ref LAST_LOGGING_DURATION: Arc<AtomicU64> = Arc::new(AtomicU64::new(0));
+    static ref LEAKY_BUCKET: LeakyBucket = LeakyBucket::new(Some(0));
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -380,6 +390,7 @@ pub struct UniformSender<T> {
 
     input: Arc<Receiver<T>>,
     counter: Arc<SenderCounter>,
+    overwritten_count: u64,
 
     encoder: Encoder<T>,
     private_conn: Mutex<Connection>,
@@ -389,6 +400,8 @@ pub struct UniformSender<T> {
     multiple_sockets_to_ingester: bool,
     dest_ip: String,
     dest_port: u16,
+    max_throughput_mbps: u64,
+    ingester_traffic_overflow_action: TrafficOverflowAction,
 
     config: SenderAccess,
 
@@ -426,6 +439,7 @@ impl<T: Sendable> UniformSender<T> {
             name,
             input,
             counter: Arc::new(SenderCounter::default()),
+            overwritten_count: 0,
             encoder: Encoder::new(
                 0,
                 SendMessageType::TaggedFlow,
@@ -440,6 +454,8 @@ impl<T: Sendable> UniformSender<T> {
             multiple_sockets_to_ingester: false,
             dest_ip: "127.0.0.1".to_string(),
             dest_port: cfg.dest_port,
+            max_throughput_mbps: 0,
+            ingester_traffic_overflow_action: TrafficOverflowAction::Waiting,
 
             running,
             stats,
@@ -522,6 +538,11 @@ impl<T: Sendable> UniformSender<T> {
     }
 
     fn send_buffer(&mut self) {
+        if self.is_traffic_overflow()
+            && self.ingester_traffic_overflow_action == TrafficOverflowAction::Dropping
+        {
+            return;
+        }
         let mut conn = match self.connection_type {
             ConnectionType::Global => self.global_shared_conn.lock().unwrap(),
             ConnectionType::PrivateShared => {
@@ -626,37 +647,51 @@ impl<T: Sendable> UniformSender<T> {
         }
     }
 
-    fn is_exceed_max_throughput(&mut self, max_throughput_mbps: u64) -> bool {
-        if max_throughput_mbps == 0 {
-            return false;
+    fn log_when_traffic_overflow(&mut self) {
+        let now = SystemTime::now()
+            .duration_since(SystemTime::UNIX_EPOCH)
+            .unwrap();
+        // to prevent frequent log printing, print at least once every 10 seconds
+        if now - Duration::from_nanos(LAST_LOGGING_DURATION.load(Ordering::Relaxed))
+            > Duration::from_secs(10)
+        {
+            warn!(
+                "{} sender dropping message, throughput exceed setting value 'max_throughput_to_ingester' {}Mbps, action {:?}, total overwrittern count {}",
+                self.name, self.max_throughput_mbps, self.ingester_traffic_overflow_action, self.overwritten_count
+            );
+            LAST_LOGGING_DURATION.store(now.as_nanos() as u64, Ordering::Relaxed);
         }
-        let max_throughput_bytes = max_throughput_mbps << 20 >> 3;
-        if TOTAL_SENT_BYTES.load(Ordering::Relaxed) > max_throughput_bytes {
-            let now = SystemTime::now()
-                .duration_since(SystemTime::UNIX_EPOCH)
-                .unwrap();
+    }
 
-            let used = now - Duration::from_nanos(SENT_START_DURATION.load(Ordering::Relaxed));
-            if used > Duration::from_secs(1) {
-                SENT_START_DURATION.store(now.as_nanos() as u64, Ordering::Relaxed);
-                TOTAL_SENT_BYTES.store(0, Ordering::Relaxed);
-            } else {
-                // to prevent frequent log printing, print at least once every 5 seconds
-                if now - Duration::from_nanos(LAST_LOGGING_DURATION.load(Ordering::Relaxed))
-                    > Duration::from_secs(5)
-                {
-                    warn!(
-                        "{} sender dropping message, throughput execeed setting value 'max_throughput_to_ingester' {}Mbps",
-                        self.name, max_throughput_mbps
-                    );
-                    LAST_LOGGING_DURATION.store(now.as_nanos() as u64, Ordering::Relaxed);
+    fn is_traffic_overflow(&mut self) -> bool {
+        if self.max_throughput_mbps == 0 {
+            return false;
+        }
+        let mut overflow = false;
+        if self.ingester_traffic_overflow_action == TrafficOverflowAction::Waiting {
+            while !LEAKY_BUCKET.acquire(self.encoder.buffer_len() as u64) {
+                // LEAKY_BUCKET token is updated every 100ms by default,
+                // wait 10ms each time until the token is acquired
+                thread::sleep(Duration::from_millis(10));
+                if self.input.total_overwritten_count() > self.overwritten_count {
+                    overflow = true;
+                    self.overwritten_count = self.input.total_overwritten_count();
                 }
-                self.exception_handler
-                    .set(Exception::DataBpsThresholdExceeded);
-                return true;
+                self.counter.waited.fetch_add(1, Ordering::Relaxed);
             }
+        } else {
+            if !LEAKY_BUCKET.acquire(self.encoder.buffer_len() as u64) {
+                overflow = true;
+                self.counter.dropped.fetch_add(1, Ordering::Relaxed);
+            }
+        }
+
+        if overflow {
+            self.exception_handler
+                .set(Exception::DataBpsThresholdExceeded);
+            self.log_when_traffic_overflow();
         }
-        return false;
+        overflow
     }
 
     fn check_or_register_counterable(&mut self, message_type: SendMessageType) {
@@ -677,7 +712,12 @@ impl<T: Sendable> UniformSender<T> {
         while self.running.load(Ordering::Relaxed) {
             let config = self.config.load();
             let socket_type = config.collector_socket_type;
-            let max_throughput_mpbs = config.max_throughput_to_ingester;
+            let max_throughput_mbps = config.max_throughput_to_ingester;
+            if self.max_throughput_mbps != max_throughput_mbps {
+                LEAKY_BUCKET.set_rate(Some(max_throughput_mbps << 17)); // Mbit -> byte
+                self.max_throughput_mbps = max_throughput_mbps;
+            }
+            self.ingester_traffic_overflow_action = config.ingester_traffic_overflow_action;
             match self.input.recv_all(
                 &mut batch,
                 Some(Duration::from_secs(Self::QUEUE_READ_TIMEOUT)),
@@ -688,13 +728,6 @@ impl<T: Sendable> UniformSender<T> {
                         start_cached = Instant::now();
                         self.cached = false;
                     }
-                    if self.is_exceed_max_throughput(max_throughput_mpbs) {
-                        self.counter
-                            .dropped
-                            .fetch_add(batch.len() as u64, Ordering::Relaxed);
-                        batch.clear();
-                        continue;
-                    }
                     for send_item in batch.drain(..) {
                         if !self.running.load(Ordering::Relaxed) {
                             break;

server/agent_config/README-CH.md

@@ -1058,9 +1058,43 @@ global:
 **详细描述**:
 
 向 Server 端 Ingester 模块发送可观测性数据的最大允许流量,
-超过此限速时数据将会主动丢弃、且采集器会标记为异常状态并触发告警。
+超限行为参考 `ingester_traffic_overflow_action` 配置描述。
 配置为 0 表示不限速。
 
+### Ingester 流量超限的动作 {#global.communication.ingester_traffic_overflow_action}
+
+**标签**:
+
+`hot_update`
+
+**FQCN**:
+
+`global.communication.ingester_traffic_overflow_action`
+
+**默认值**:
+```yaml
+global:
+  communication:
+    ingester_traffic_overflow_action: 0
+```
+
+**枚举可选值**:
+| Value | Note                         |
+| ----- | ---------------------------- |
+| 0 | WAIT |
+| 1 | DROP |
+
+**模式**:
+| Key  | Value                        |
+| ---- | ---------------------------- |
+| Type | int |
+
+**详细描述**:
+
+Ingester 流量超限的动作
+- WAIT：暂停发送，数据缓存到队列，等待下次发送。
+- DROP：直接丢弃数据，并触发 Agent `数据流量达到限速`异常。
+
 ### 请求 NAT IP 地址 {#global.communication.request_via_nat_ip}
 
 **标签**: