[Agent] Support for the Agent to modify its own container resource limits

agent/Cargo.toml

@@ -14,6 +14,7 @@ arc-swap = "1.5.0"
 base64 = "0.21"
 bincode = "2.0.0-rc.1"
 bitflags = "1.3.2"
+bollard = "0.16.1"
 bson = "2.7.0"
 bytesize = "1.1.0"
 cadence = "0.27.0"

agent/src/config/config.rs

@@ -1159,7 +1159,7 @@ pub struct RuntimeConfig {
     pub vtap_group_id: String,
     #[serde(skip)]
     pub enabled: bool,
-    pub max_cpus: u32,
+    pub max_millicpus: u32,
     pub max_memory: u64,
     pub sync_interval: u64,          // unit(second)
     pub platform_sync_interval: u64, // unit(second)
@@ -1307,7 +1307,7 @@ impl RuntimeConfig {
         Self {
             vtap_group_id: Default::default(),
             enabled: true,
-            max_cpus: 1,
+            max_millicpus: 1000,
             max_memory: 768,
             sync_interval: 60,
             platform_sync_interval: 10,
@@ -1487,7 +1487,16 @@ impl TryFrom<trident::Config> for RuntimeConfig {
         let rc = Self {
             vtap_group_id: Default::default(),
             enabled: conf.enabled(),
-            max_cpus: conf.max_cpus(),
+            max_millicpus: {
+                // Compatible with max_cpus and max_millicpus, take the smaller value in milli-cores.
+                let max_cpus = conf.max_cpus() * 1000;
+                let max_millicpus = conf.max_millicpus.unwrap_or(max_cpus); // conf.max_millicpus may be None, handle the case where conf.max_millicpus is None
+                if max_cpus != 0 && max_millicpus != 0 {
+                    max_cpus.min(max_millicpus)
+                } else {
+                    max_cpus | max_millicpus
+                }
+            },
             max_memory: (conf.max_memory() as u64) << 20,
             sync_interval: conf.sync_interval() as u64,
             platform_sync_interval: conf.platform_sync_interval() as u64,

agent/src/config/handler.rs

@@ -65,14 +65,17 @@ use crate::{
     handler::PacketHandlerBuilder,
     metric::document::TapSide,
     trident::{AgentComponents, RunningMode},
-    utils::environment::{free_memory_check, get_container_mem_limit, running_in_container},
+    utils::environment::{free_memory_check, running_in_container},
 };
 #[cfg(any(target_os = "linux", target_os = "android"))]
 use crate::{
     dispatcher::recv_engine::af_packet::OptTpacketVersion,
     ebpf::CAP_LEN_MAX,
     platform::ProcRegRewrite,
-    utils::environment::{get_ctrl_ip_and_mac, is_tt_workload},
+    utils::environment::{
+        get_container_resource_limits, get_ctrl_ip_and_mac, is_tt_workload,
+        set_container_resource_limit,
+    },
 };
 #[cfg(target_os = "linux")]
 use crate::{
@@ -194,7 +197,7 @@ impl fmt::Debug for CollectorConfig {
 #[derive(Clone, Debug, PartialEq)]
 pub struct EnvironmentConfig {
     pub max_memory: u64,
-    pub max_cpus: u32,
+    pub max_millicpus: u32,
     pub process_threshold: u32,
     pub thread_threshold: u32,
     pub sys_free_memory_limit: u32,
@@ -1202,10 +1205,7 @@ impl TryFrom<(Config, RuntimeConfig)> for ModuleConfig {
     type Error = ConfigError;
 
     fn try_from(conf: (Config, RuntimeConfig)) -> Result<Self, Self::Error> {
-        let (static_config, mut conf) = conf;
-        if running_in_container() {
-            conf.max_memory = get_container_mem_limit().unwrap_or(conf.max_memory);
-        }
+        let (static_config, conf) = conf;
         let controller_ip = static_config.controller_ips[0].parse::<IpAddr>().unwrap();
         let dest_ip = if conf.analyzer_ip.len() > 0 {
             conf.analyzer_ip.clone()
@@ -1231,7 +1231,7 @@ impl TryFrom<(Config, RuntimeConfig)> for ModuleConfig {
             },
             environment: EnvironmentConfig {
                 max_memory: conf.max_memory,
-                max_cpus: conf.max_cpus,
+                max_millicpus: conf.max_millicpus,
                 process_threshold: conf.process_threshold,
                 thread_threshold: conf.thread_threshold,
                 sys_free_memory_limit: conf.sys_free_memory_limit,
@@ -1574,6 +1574,8 @@ impl TryFrom<(Config, RuntimeConfig)> for ModuleConfig {
 pub struct ConfigHandler {
     pub ctrl_ip: IpAddr,
     pub ctrl_mac: MacAddr,
+    pub container_cpu_limit: u32, // unit: milli-core
+    pub container_mem_limit: u64, // unit: bytes
     pub logger_handle: Option<LoggerHandle>,
     // need update
     pub static_config: Config,
@@ -1587,10 +1589,17 @@ impl ConfigHandler {
             ModuleConfig::try_from((config.clone(), RuntimeConfig::default())).unwrap();
         let current_config = Arc::new(ArcSwap::from_pointee(candidate_config.clone()));
 
+        #[cfg(any(target_os = "linux", target_os = "android"))]
+        let (container_cpu_limit, container_mem_limit) = get_container_resource_limits();
+        #[cfg(target_os = "windows")]
+        let (container_cpu_limit, container_mem_limit) = (0, 0);
+
         Self {
             static_config: config,
             ctrl_ip,
             ctrl_mac,
+            container_cpu_limit,
+            container_mem_limit,
             candidate_config,
             current_config,
             logger_handle: None,
@@ -2194,25 +2203,43 @@ impl ConfigHandler {
                 candidate_config.environment.max_memory = new_config.environment.max_memory;
             }
 
-            if candidate_config.environment.max_cpus != new_config.environment.max_cpus {
-                info!("cpu limit set to {}", new_config.environment.max_cpus);
-                candidate_config.environment.max_cpus = new_config.environment.max_cpus;
+            if candidate_config.environment.max_millicpus != new_config.environment.max_millicpus {
+                info!(
+                    "cpu limit set to {}",
+                    new_config.environment.max_millicpus as f64 / 1000.0
+                );
+                candidate_config.environment.max_millicpus = new_config.environment.max_millicpus;
+            }
+            #[cfg(target_os = "linux")]
+            if running_in_container() {
+                if self.container_cpu_limit != candidate_config.environment.max_millicpus
+                    || self.container_mem_limit != candidate_config.environment.max_memory
+                {
+                    info!("current container cpu limit: {}, memory limit: {}bytes, set cpu limit {} and memory limit {}bytes", self.container_cpu_limit as f64 / 1000.0, self.container_mem_limit, candidate_config.environment.max_millicpus as f64 / 1000.0, candidate_config.environment.max_memory);
+                    if let Err(e) = runtime.block_on(set_container_resource_limit(
+                        candidate_config.environment.max_millicpus,
+                        candidate_config.environment.max_memory,
+                    )) {
+                        warn!("set container resources limit failed: {:?}", e);
+                    };
+                }
             }
         } else {
             let mut system = sysinfo::System::new();
             system.refresh_memory();
             let max_memory = system.total_memory();
             system.refresh_cpu();
             let max_cpus = 1.max(system.cpus().len()) as u32;
+            let max_millicpus = max_cpus * 1000;
 
             if candidate_config.environment.max_memory != max_memory {
                 info!("memory set ulimit when tap_mode=analyzer");
                 candidate_config.environment.max_memory = max_memory;
             }
 
-            if candidate_config.environment.max_cpus != max_cpus {
+            if candidate_config.environment.max_millicpus != max_millicpus {
                 info!("cpu set ulimit when tap_mode=analyzer");
-                candidate_config.environment.max_cpus = max_cpus;
+                candidate_config.environment.max_millicpus = max_millicpus;
             }
         }
 

agent/src/monitor.rs

@@ -224,9 +224,9 @@ impl RefCountable for SysStatusBroker {
             CounterValue::Unsigned(self.config.load().max_memory as u64),
         ));
         metrics.push((
-            "max_cpus",
+            "max_millicpus",
             CounterType::Gauged,
-            CounterValue::Unsigned(self.config.load().max_cpus as u64),
+            CounterValue::Unsigned(self.config.load().max_millicpus as u64),
         ));
         metrics.push((
             "system_free_memory_limit",

agent/src/trident.rs

@@ -90,7 +90,7 @@ use crate::{
         command::get_hostname,
         environment::{
             check, controller_ip_check, free_memory_check, free_space_checker, get_ctrl_ip_and_mac,
-            get_env, kernel_check, running_in_container, tap_interface_check,
+            get_env, kernel_check, running_in_container, running_in_k8s, tap_interface_check,
             trident_process_check,
         },
         guard::Guard,
@@ -428,7 +428,6 @@ impl Trident {
                 "use K8S_NODE_IP_FOR_DEEPFLOW env ip as destination_ip({})",
                 ctrl_ip
             );
-            warn!("When running in a container, the cpu and memory limits notified by deepflow-server will be ignored, please make sure to use K8s or docker for resource limits.");
         }
 
         #[cfg(target_os = "linux")]
@@ -498,7 +497,7 @@ impl Trident {
         if matches!(
             config_handler.static_config.agent_mode,
             RunningMode::Managed
-        ) && running_in_container()
+        ) && running_in_k8s()
             && config_handler
                 .static_config
                 .kubernetes_cluster_id

agent/src/utils/cgroups/linux.rs

@@ -133,18 +133,18 @@ impl Cgroups {
 
         let environment_config = self.config.clone();
         let running = self.running.clone();
-        let mut last_cpu = 0;
+        let mut last_millicpus = 0;
         let mut last_memory = 0;
         let cgroup = self.cgroup.clone();
         let thread = thread::Builder::new()
             .name("cgroups-controller".to_owned())
             .spawn(move || {
                 loop {
                     let environment = environment_config.load();
-                    let max_cpus = environment.max_cpus;
+                    let max_millicpus = environment.max_millicpus;
                     let max_memory = environment.max_memory;
-                    if max_cpus != last_cpu || max_memory != last_memory {
-                        if let Err(e) = Self::apply(cgroup.clone(), max_cpus, max_memory) {
+                    if max_millicpus != last_millicpus || max_memory != last_memory {
+                        if let Err(e) = Self::apply(cgroup.clone(), max_millicpus, max_memory) {
                             warn!(
                                 "apply cgroups resource failed, {}, deepflow-agent restart...",
                                 e
@@ -153,7 +153,7 @@ impl Cgroups {
                             break;
                         }
                     }
-                    last_cpu = max_cpus;
+                    last_millicpus = max_millicpus;
                     last_memory = max_memory;
 
                     let (running, timer) = &*running;
@@ -175,9 +175,9 @@ impl Cgroups {
     }
 
     /// 更改资源限制
-    pub fn apply(cgroup: Cgroup, max_cpus: u32, max_memory: u64) -> Result<(), Error> {
+    pub fn apply(cgroup: Cgroup, max_millicpus: u32, max_memory: u64) -> Result<(), Error> {
         let mut resources = Resources::default();
-        let cpu_quota = max_cpus * DEFAULT_CPU_CFS_PERIOD_US;
+        let cpu_quota = max_millicpus * 100; // The unit of cpu_quota is 100_000 us. Convert max_millicpus to the unit of cpu_quota
         let cpu_resources = CpuResources {
             quota: Some(cpu_quota as i64),
             period: Some(DEFAULT_CPU_CFS_PERIOD_US as u64),