From 6def2a7994e615bb3dcf5ac5066a4fe964054af8 Mon Sep 17 00:00:00 2001
From: Jiajie Chen <c@jia.je>
Date: Wed, 12 Jun 2024 22:00:40 +0800
Subject: [PATCH] feat: respect cgroup limits in total memory

---
 worker/src/build.rs     |  4 ++--
 worker/src/heartbeat.rs |  4 ++--
 worker/src/lib.rs       | 10 ++++++++++
 worker/src/main.rs      |  5 +++++
 4 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/worker/src/build.rs b/worker/src/build.rs
index 48cfc8b..a6196ad 100644
--- a/worker/src/build.rs
+++ b/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,
     };
diff --git a/worker/src/heartbeat.rs b/worker/src/heartbeat.rs
index f7628dc..4a531ac 100644
--- a/worker/src/heartbeat.rs
+++ b/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,
diff --git a/worker/src/lib.rs b/worker/src/lib.rs
index 1d58fc0..6f03771 100644
--- a/worker/src/lib.rs
+++ b/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
+    }
+}
diff --git a/worker/src/main.rs b/worker/src/main.rs
index 25b63d6..03dc7b9 100644
--- a/worker/src/main.rs
+++ b/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;