agent: add liveness probe component

Author: rvql

agent/config/deepflow-agent.yaml

@@ -39,6 +39,11 @@ controller-ips:
 ## kubernetes api watcher
 #async-worker-thread-number: 16
 
+## Enable HTTP liveness probe endpoint
+#liveness-probe-enabled: false
+## HTTP liveness probe listen port
+#liveness-probe-port: 39090
+
 ## Type of agent identifier, choose from [ip-and-mac, ip], defaults to "ip-and-mac"
 #agent-unique-identifier: ip-and-mac
 

agent/src/config/config.rs

@@ -135,6 +135,8 @@ pub struct Config {
     pub pid_file: String,
     pub team_id: String,
     pub cgroups_disabled: bool,
+    pub liveness_probe_enabled: bool,
+    pub liveness_probe_port: u16,
 }
 
 impl Config {
@@ -297,6 +299,8 @@ impl Default for Config {
             pid_file: Default::default(),
             team_id: "".into(),
             cgroups_disabled: false,
+            liveness_probe_enabled: false,
+            liveness_probe_port: 39090,
         }
     }
 }
@@ -3975,6 +3979,8 @@ mod tests {
             .expect("failed loading config file");
         assert_eq!(c.controller_ips.len(), 1);
         assert_eq!(&c.controller_ips[0], "127.0.0.1");
+        assert!(!c.liveness_probe_enabled);
+        assert_eq!(c.liveness_probe_port, 39090);
     }
 
     #[test]

agent/src/dispatcher/analyzer_mode_dispatcher.rs

@@ -20,7 +20,7 @@ use std::{
     ops::Add,
     sync::{atomic::Ordering, Arc, RwLock},
     thread::{self, JoinHandle},
-    time::Duration,
+    time::{Duration, Instant},
 };
 
 use arc_swap::access::Access;
@@ -47,6 +47,7 @@ use crate::{
     },
     flow_generator::{flow_map::Config, FlowMap},
     handler::{MiniPacket, PacketHandler},
+    liveness::{DebugInfo, LivenessHandle},
     rpc::get_timestamp,
     utils::{
         bytes::read_u32_be,
@@ -153,6 +154,7 @@ pub(super) struct AnalyzerModeDispatcher {
     pub(super) pool_raw_size: usize,
     pub(super) flow_generator_thread_handler: Option<JoinHandle<()>>,
     pub(super) pipeline_thread_handler: Option<JoinHandle<()>>,
+    pub(super) flow_generator_liveness: LivenessHandle,
     pub(super) queue_debugger: Arc<QueueDebugger>,
     pub(super) stats_collector: Arc<stats::Collector>,
     pub(super) inner_queue_size: usize,
@@ -270,6 +272,7 @@ impl AnalyzerModeDispatcher {
         let collector_config = base.collector_config.clone();
         let packet_sequence_output_queue = base.packet_sequence_output_queue.clone(); // Enterprise Edition Feature: packet-sequence
         let stats = base.stats.clone();
+        let liveness = self.flow_generator_liveness.clone();
         #[cfg(any(target_os = "linux", target_os = "android"))]
         let cpu_set = base.options.lock().unwrap().cpu_set;
 
@@ -280,6 +283,7 @@ impl AnalyzerModeDispatcher {
                     let mut timestamp_map: HashMap<CaptureNetworkType, Duration> = HashMap::new();
                     let mut batch = Vec::with_capacity(HANDLER_BATCH_SIZE);
                     let mut output_batch = Vec::with_capacity(HANDLER_BATCH_SIZE);
+                    let mut last_liveness = Instant::now();
                     let mut flow_map = FlowMap::new(
                         id as u32,
                         Some(flow_output_queue),
@@ -298,8 +302,13 @@ impl AnalyzerModeDispatcher {
                             warn!("CPU Affinity({:?}) bind error: {:?}.", &cpu_set, e);
                         }
                     }
+                    liveness.start(DebugInfo::new("running"));
 
                     while !terminated.load(Ordering::Relaxed) {
+                        if last_liveness.elapsed() >= Duration::from_secs(1) {
+                            liveness.heartbeat(DebugInfo::new("running"));
+                            last_liveness = Instant::now();
+                        }
                         let config = Config {
                             flow: &flow_map_config.load(),
                             log_parser: &log_parser_config.load(),
@@ -433,6 +442,7 @@ impl AnalyzerModeDispatcher {
                             output_batch.clear();
                         }
                     }
+                    liveness.stop(DebugInfo::new("stopped"));
                 })
                 .unwrap(),
         );
@@ -543,6 +553,8 @@ impl AnalyzerModeDispatcher {
         let id = base.id;
         let mut batch = Vec::with_capacity(HANDLER_BATCH_SIZE);
         let mut allocator = Allocator::new(self.raw_packet_block_size);
+        let mut last_liveness = Instant::now();
+        base.liveness_handle.start(DebugInfo::new("running"));
         #[cfg(any(target_os = "linux", target_os = "android"))]
         let cpu_set = base.options.lock().unwrap().cpu_set;
         #[cfg(any(target_os = "linux", target_os = "android"))]
@@ -553,6 +565,10 @@ impl AnalyzerModeDispatcher {
         }
 
         while !base.terminated.load(Ordering::Relaxed) {
+            if last_liveness.elapsed() >= Duration::from_secs(1) {
+                base.liveness_handle.heartbeat(DebugInfo::new("running"));
+                last_liveness = Instant::now();
+            }
             if base.reset_whitelist.swap(false, Ordering::Relaxed) {
                 base.tap_interface_whitelist.reset();
             }
@@ -615,6 +631,7 @@ impl AnalyzerModeDispatcher {
             let _ = handler.join();
         }
 
+        base.liveness_handle.stop(DebugInfo::new("stopped"));
         self.base.terminate_handler();
         info!("Stopped dispatcher {}", self.base.is.log_id);
     }

agent/src/dispatcher/base_dispatcher.rs

@@ -54,6 +54,7 @@ use crate::{
     exception::ExceptionHandler,
     flow_generator::AppProto,
     handler::PacketHandlerBuilder,
+    liveness::LivenessHandle,
     policy::PolicyGetter,
     rpc::get_timestamp,
     utils::{bytes::read_u16_be, stats::Collector},
@@ -139,6 +140,7 @@ pub(super) struct InternalState {
     // dispatcher id for easy debugging
     pub log_id: String,
     pub promisc_if_indices: Vec<i32>,
+    pub liveness_handle: LivenessHandle,
 }
 
 impl BaseDispatcher {

agent/src/dispatcher/local_mode_dispatcher.rs

@@ -21,7 +21,7 @@ use std::str;
 use std::sync::atomic::Ordering;
 #[cfg(target_os = "linux")]
 use std::sync::Arc;
-use std::time::Duration;
+use std::time::{Duration, Instant};
 
 use arc_swap::access::Access;
 use log::{debug, info, log_enabled, warn};
@@ -37,6 +37,7 @@ use super::{
     error::Result,
     TunnelTypeBitmap,
 };
+use crate::liveness::DebugInfo;
 
 #[cfg(target_os = "linux")]
 use crate::platform::{GenericPoller, LibvirtXmlExtractor, Poller};
@@ -201,6 +202,8 @@ impl LocalModeDispatcher {
         info!("Start dispatcher {}", base.log_id);
         let time_diff = base.ntp_diff.load(Ordering::Relaxed);
         let mut prev_timestamp = get_timestamp(time_diff);
+        let mut last_liveness = Instant::now();
+        base.liveness_handle.start(DebugInfo::new("running"));
         #[cfg(any(target_os = "linux", target_os = "android"))]
         let cpu_set = base.options.lock().unwrap().cpu_set;
         #[cfg(any(target_os = "linux", target_os = "android"))]
@@ -228,6 +231,10 @@ impl LocalModeDispatcher {
         let mut collector_config = base.collector_config.load().clone();
 
         while !base.terminated.load(Ordering::Relaxed) {
+            if last_liveness.elapsed() >= Duration::from_secs(1) {
+                base.liveness_handle.heartbeat(DebugInfo::new("running"));
+                last_liveness = Instant::now();
+            }
             if base.need_reload_config.swap(false, Ordering::Relaxed) {
                 info!("dispatcher reload config");
                 flow_config = base.flow_map_config.load().clone();
@@ -299,6 +306,7 @@ impl LocalModeDispatcher {
             base.check_and_update_bpf(&mut self.base.engine);
         }
 
+        base.liveness_handle.stop(DebugInfo::new("stopped"));
         self.base.terminate_handler();
         info!("Stopped dispatcher {}", self.base.is.log_id);
     }

agent/src/dispatcher/local_multins_mode_dispatcher.rs

@@ -7,7 +7,7 @@ use std::{
         Arc, Mutex,
     },
     thread::{self, JoinHandle},
-    time::Duration,
+    time::{Duration, Instant},
 };
 
 use arc_swap::access::Access;
@@ -33,6 +33,7 @@ use crate::{
     config::handler::DispatcherAccess,
     exception::ExceptionHandler,
     flow_generator::{flow_map::Config, FlowMap},
+    liveness::{DebugInfo, LivenessHandle},
     rpc::get_timestamp,
     utils::stats::QueueStats,
 };
@@ -44,13 +45,15 @@ pub struct LocalMultinsModeDispatcher {
     base: BaseDispatcher,
 
     receiver_manager: Option<JoinHandle<()>>,
+    receiver_manager_liveness: LivenessHandle,
 }
 
 impl LocalMultinsModeDispatcher {
-    pub fn new(base: BaseDispatcher) -> Self {
+    pub fn new(base: BaseDispatcher, receiver_manager_liveness: LivenessHandle) -> Self {
         Self {
             base,
             receiver_manager: None,
+            receiver_manager_liveness,
         }
     }
 
@@ -84,6 +87,7 @@ impl LocalMultinsModeDispatcher {
             counter: base.counter.clone(),
             ntp_diff: base.ntp_diff.clone(),
             bpf_controls: bpf_controls.clone(),
+            liveness: self.receiver_manager_liveness.clone(),
             output: packet_input,
         };
         self.receiver_manager.replace(
@@ -109,9 +113,15 @@ impl LocalMultinsModeDispatcher {
         let tunnel_type_trim_bitmap = base.tunnel_type_trim_bitmap.clone();
         let mut batch = Vec::with_capacity(PACKET_BATCH_SIZE);
         let mut tap_interface_whitelists: HashMap<u64, TapInterfaceWhitelist> = HashMap::new();
+        let mut last_liveness = Instant::now();
 
         super::set_cpu_affinity(&base.options);
+        base.liveness_handle.start(DebugInfo::new("running"));
         while !base.terminated.load(Ordering::Relaxed) {
+            if last_liveness.elapsed() >= Duration::from_secs(1) {
+                base.liveness_handle.heartbeat(DebugInfo::new("running"));
+                last_liveness = Instant::now();
+            }
             let config = Config {
                 flow: &base.flow_map_config.load(),
                 log_parser: &base.log_parser_config.load(),
@@ -225,6 +235,7 @@ impl LocalMultinsModeDispatcher {
                 });
             }
         }
+        base.liveness_handle.stop(DebugInfo::new("stopped"));
         info!("Stopping local multi-namespace dispatcher");
         info!("Wait for receiver manager to stop");
         self.receiver_manager.take().unwrap().join().unwrap();
@@ -513,6 +524,7 @@ struct ReceiverManager {
     exception_handler: ExceptionHandler,
     counter: Arc<PacketCounter>,
     ntp_diff: Arc<AtomicI64>,
+    liveness: LivenessHandle,
 
     output: DebugSender<Packet>,
 }
@@ -551,11 +563,13 @@ impl ReceiverManager {
             super::set_cpu_affinity(&self.options);
 
             info!("Receiver manager started");
+            self.liveness.start(DebugInfo::new("running"));
 
             let mut loop_count = 0;
             let mut zombie_threads = vec![];
             let mut receiver_threads: HashMap<NsFile, PktReceiverHandle> = HashMap::new();
             while !self.terminated.load(Ordering::Relaxed) {
+                self.liveness.heartbeat(DebugInfo::new("running"));
                 loop_count = (loop_count + 1) % Self::INTERVAL_SECS;
                 if loop_count != 1 {
                     // actual interval is (INTERVAL_SECS - 1) to make this simple and less error prone
@@ -762,6 +776,7 @@ impl ReceiverManager {
                 let _ = handle.join_handle.take().unwrap().join();
             }
 
+            self.liveness.stop(DebugInfo::new("stopped"));
             info!("Receiver manager stopped");
         }
     }

agent/src/dispatcher/local_plus_mode_dispatcher.rs

@@ -23,7 +23,7 @@ use std::process::Command;
 use std::str;
 use std::sync::{atomic::Ordering, Arc};
 use std::thread::{self, JoinHandle};
-use std::time::Duration;
+use std::time::{Duration, Instant};
 
 use arc_swap::access::Access;
 use log::{debug, info, log_enabled, warn};
@@ -49,6 +49,7 @@ use crate::{
     config::DispatcherConfig,
     flow_generator::{flow_map::Config, FlowMap},
     handler::MiniPacket,
+    liveness::{DebugInfo, LivenessHandle},
     rpc::get_timestamp,
     utils::{
         bytes::read_u16_be,
@@ -73,6 +74,7 @@ pub(super) struct LocalPlusModeDispatcher {
     pub(super) stats_collector: Arc<stats::Collector>,
     pub(super) flow_generator_thread_handler: Option<JoinHandle<()>>,
     pub(super) pipeline_thread_handler: Option<JoinHandle<()>>,
+    pub(super) flow_generator_liveness: LivenessHandle,
     pub(super) inner_queue_size: usize,
     pub(super) raw_packet_block_size: usize,
     pub(super) pool_raw_size: usize,
@@ -113,6 +115,7 @@ impl LocalPlusModeDispatcher {
         let npb_dedup_enabled = base.npb_dedup_enabled.clone();
         let pool_raw_size = self.pool_raw_size;
         let tunnel_type_trim_bitmap = base.tunnel_type_trim_bitmap.clone();
+        let liveness = self.flow_generator_liveness.clone();
         #[cfg(any(target_os = "linux", target_os = "android"))]
         let cpu_set = base.options.lock().unwrap().cpu_set;
 
@@ -122,6 +125,7 @@ impl LocalPlusModeDispatcher {
                 .spawn(move || {
                     let mut batch = Vec::with_capacity(HANDLER_BATCH_SIZE);
                     let mut output_batch = Vec::with_capacity(HANDLER_BATCH_SIZE);
+                    let mut last_liveness = Instant::now();
                     let mut flow_map = FlowMap::new(
                         id as u32,
                         Some(flow_output_queue),
@@ -140,8 +144,13 @@ impl LocalPlusModeDispatcher {
                             warn!("CPU Affinity({:?}) bind error: {:?}.", &cpu_set, e);
                         }
                     }
+                    liveness.start(DebugInfo::new("running"));
 
                     while !terminated.load(Ordering::Relaxed) {
+                        if last_liveness.elapsed() >= Duration::from_secs(1) {
+                            liveness.heartbeat(DebugInfo::new("running"));
+                            last_liveness = Instant::now();
+                        }
                         let config = Config {
                             flow: &flow_map_config.load(),
                             log_parser: &log_parser_config.load(),
@@ -291,6 +300,7 @@ impl LocalPlusModeDispatcher {
                             output_batch.clear();
                         }
                     }
+                    liveness.stop(DebugInfo::new("stopped"));
                 })
                 .unwrap(),
         );
@@ -385,8 +395,14 @@ impl LocalPlusModeDispatcher {
         let id = base.id;
         let mut batch = Vec::with_capacity(HANDLER_BATCH_SIZE);
         let mut allocator = Allocator::new(self.raw_packet_block_size);
+        let mut last_liveness = Instant::now();
+        base.liveness_handle.start(DebugInfo::new("running"));
 
         while !base.terminated.load(Ordering::Relaxed) {
+            if last_liveness.elapsed() >= Duration::from_secs(1) {
+                base.liveness_handle.heartbeat(DebugInfo::new("running"));
+                last_liveness = Instant::now();
+            }
             if base.reset_whitelist.swap(false, Ordering::Relaxed) {
                 base.tap_interface_whitelist.reset();
             }
@@ -450,6 +466,7 @@ impl LocalPlusModeDispatcher {
             let _ = handler.join();
         }
 
+        base.liveness_handle.stop(DebugInfo::new("stopped"));
         self.base.terminate_handler();
         info!("Stopped dispatcher {}", self.base.is.log_id);
     }