feat: respect cgroup limits in total memory

worker/src/build.rs

@@ -1,4 +1,4 @@
-use crate::Args;
+use crate::{get_memory_bytes, Args};
 use chrono::Local;
 use common::{JobOk, WorkerJobUpdateRequest, WorkerPollRequest, WorkerPollResponse};
 use log::{error, info, warn};
@@ -301,7 +301,7 @@ async fn build_worker_inner(args: &Args) -> anyhow::Result<()> {
         hostname: gethostname::gethostname().to_string_lossy().to_string(),
         arch: args.arch.clone(),
         worker_secret: args.worker_secret.clone(),
-        memory_bytes: sysinfo::System::new_all().total_memory() as i64,
+        memory_bytes: get_memory_bytes(),
         disk_free_space_bytes: fs2::free_space(std::env::current_dir()?)? as i64,
         logical_cores: num_cpus::get() as i32,
     };

worker/src/heartbeat.rs

@@ -1,4 +1,4 @@
-use crate::Args;
+use crate::{get_memory_bytes, Args};
 use common::WorkerHeartbeatRequest;
 use log::{info, warn};
 use std::{
@@ -40,7 +40,7 @@ pub async fn heartbeat_worker_inner(args: &Args) -> anyhow::Result<()> {
                 arch: args.arch.clone(),
                 worker_secret: args.worker_secret.clone(),
                 git_commit: env!("VERGEN_GIT_DESCRIBE").to_string(),
-                memory_bytes: sysinfo::System::new_all().total_memory() as i64,
+                memory_bytes: get_memory_bytes(),
                 disk_free_space_bytes: fs2::free_space(std::env::current_dir()?)? as i64,
                 logical_cores: num_cpus::get() as i32,
                 performance: args.worker_performance,

worker/src/lib.rs

@@ -1,5 +1,6 @@
 use clap::Parser;
 use std::path::PathBuf;
+use sysinfo::System;
 
 pub mod build;
 pub mod heartbeat;
@@ -49,3 +50,12 @@ pub struct Args {
     #[arg(short = 'p', long, env = "BUILDIT_WORKER_PERFORMANCE")]
     pub worker_performance: Option<i64>,
 }
+
+pub fn get_memory_bytes() -> i64 {
+    let system = System::new_all();
+    if let Some(limits) = system.cgroup_limits() {
+        limits.total_memory as i64
+    } else {
+        system.total_memory() as i64
+    }
+}

worker/src/main.rs

@@ -1,5 +1,6 @@
 use clap::Parser;
 use log::info;
+use sysinfo::System;
 use worker::{build::build_worker, heartbeat::heartbeat_worker, Args};
 
 #[tokio::main]
@@ -9,6 +10,10 @@ async fn main() -> anyhow::Result<()> {
     let args = Args::parse();
     info!("Starting AOSC BuildIt! worker");
 
+    // Refresh memory usage for get_memory_bytes()
+    let mut s = System::new();
+    s.refresh_memory();
+
     tokio::spawn(heartbeat_worker(args.clone()));
 
     build_worker(args.clone()).await;