agent: add get volume stats handler in agent

retrieve the stats of direct-assigned volumes from the guest

Fixes: #3454

Signed-off-by: shuochen0311 <shuo.chen@databricks.com>

src/agent/Cargo.lock

@@ -214,6 +214,12 @@ dependencies = [
  "syn",
 ]
 
+[[package]]
+name = "core-foundation-sys"
+version = "0.8.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5827cebf4670468b8772dd191856768aedcb1b0278a04f989f7766351917b9dc"
+
 [[package]]
 name = "crc32fast"
 version = "1.3.0"
@@ -233,6 +239,30 @@ dependencies = [
  "crossbeam-utils",
 ]
 
+[[package]]
+name = "crossbeam-deque"
+version = "0.8.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6455c0ca19f0d2fbf751b908d5c55c1f5cbc65e03c4225427254b46890bdde1e"
+dependencies = [
+ "cfg-if 1.0.0",
+ "crossbeam-epoch",
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "crossbeam-epoch"
+version = "0.9.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "97242a70df9b89a65d0b6df3c4bf5b9ce03c5b7309019777fbde37e7537f8762"
+dependencies = [
+ "cfg-if 1.0.0",
+ "crossbeam-utils",
+ "lazy_static",
+ "memoffset",
+ "scopeguard",
+]
+
 [[package]]
 name = "crossbeam-utils"
 version = "0.8.5"
@@ -565,6 +595,7 @@ dependencies = [
  "slog",
  "slog-scope",
  "slog-stdlog",
+ "sysinfo",
  "tempfile",
  "thiserror",
  "tokio",
@@ -1238,6 +1269,31 @@ dependencies = [
  "rand_core",
 ]
 
+[[package]]
+name = "rayon"
+version = "1.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c06aca804d41dbc8ba42dfd964f0d01334eceb64314b9ecf7c5fad5188a06d90"
+dependencies = [
+ "autocfg",
+ "crossbeam-deque",
+ "either",
+ "rayon-core",
+]
+
+[[package]]
+name = "rayon-core"
+version = "1.9.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d78120e2c850279833f1dd3582f730c4ab53ed95aeaaaa862a2a5c71b1656d8e"
+dependencies = [
+ "crossbeam-channel",
+ "crossbeam-deque",
+ "crossbeam-utils",
+ "lazy_static",
+ "num_cpus",
+]
+
 [[package]]
 name = "redox_syscall"
 version = "0.2.10"
@@ -1518,6 +1574,21 @@ dependencies = [
  "unicode-xid",
 ]
 
+[[package]]
+name = "sysinfo"
+version = "0.23.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9e757000a4bed2b1be9be65a3f418b9696adf30bb419214c73997422de73a591"
+dependencies = [
+ "cfg-if 1.0.0",
+ "core-foundation-sys",
+ "libc",
+ "ntapi",
+ "once_cell",
+ "rayon",
+ "winapi",
+]
+
 [[package]]
 name = "take_mut"
 version = "0.2.2"

src/agent/Cargo.toml

@@ -20,6 +20,7 @@ scopeguard = "1.0.0"
 thiserror = "1.0.26"
 regex = "1.5.4"
 serial_test = "0.5.1"
+sysinfo = "0.23.0"
 
 # Async helpers
 async-trait = "0.1.42"

src/agent/src/rpc.rs

@@ -23,9 +23,10 @@ use oci::{LinuxNamespace, Root, Spec};
 use protobuf::{Message, RepeatedField, SingularPtrField};
 use protocols::agent::{
     AddSwapRequest, AgentDetails, CopyFileRequest, GuestDetailsResponse, Interfaces, Metrics,
-    OOMEvent, ReadStreamResponse, Routes, StatsContainerResponse, WaitProcessResponse,
-    WriteStreamResponse,
+    OOMEvent, ReadStreamResponse, Routes, StatsContainerResponse, VolumeStatsRequest,
+    WaitProcessResponse, WriteStreamResponse,
 };
+use protocols::csi::{VolumeCondition, VolumeStatsResponse, VolumeUsage, VolumeUsage_Unit};
 use protocols::empty::Empty;
 use protocols::health::{
     HealthCheckResponse, HealthCheckResponse_ServingStatus, VersionCheckResponse,
@@ -43,6 +44,8 @@ use nix::sys::stat;
 use nix::unistd::{self, Pid};
 use rustjail::process::ProcessOperations;
 
+use sysinfo::{DiskExt, System, SystemExt};
+
 use crate::device::{
     add_devices, get_virtio_blk_pci_device_name, update_device_cgroup, update_env_pci,
 };
@@ -68,6 +71,7 @@ use tracing::instrument;
 use libc::{self, c_char, c_ushort, pid_t, winsize, TIOCSWINSZ};
 use std::convert::TryFrom;
 use std::fs;
+use std::os::unix::fs::MetadataExt;
 use std::os::unix::prelude::PermissionsExt;
 use std::process::{Command, Stdio};
 use std::time::Duration;
@@ -1254,6 +1258,47 @@ impl protocols::agent_ttrpc::AgentService for AgentService {
         Err(ttrpc_error!(ttrpc::Code::INTERNAL, ""))
     }
 
+    async fn get_volume_stats(
+        &self,
+        ctx: &TtrpcContext,
+        req: VolumeStatsRequest,
+    ) -> ttrpc::Result<VolumeStatsResponse> {
+        trace_rpc_call!(ctx, "get_volume_stats", req);
+        is_allowed!(req);
+
+        info!(sl!(), "get volume stats!");
+        let mut resp = VolumeStatsResponse::new();
+
+        let mut condition = VolumeCondition::new();
+
+        match File::open(&req.volume_guest_path) {
+            Ok(_) => {
+                condition.abnormal = false;
+                condition.message = String::from("OK");
+            }
+            Err(e) => {
+                info!(sl!(), "failed to open the volume");
+                return Err(ttrpc_error!(ttrpc::Code::INTERNAL, e));
+            }
+        };
+
+        let mut usage_vec = Vec::new();
+
+        // to get volume capacity stats
+        get_volume_capacity_stats(&req.volume_guest_path)
+            .map(|u| usage_vec.push(u))
+            .map_err(|e| ttrpc_error!(ttrpc::Code::INTERNAL, e))?;
+
+        // to get volume inode stats
+        get_volume_inode_stats(&req.volume_guest_path)
+            .map(|u| usage_vec.push(u))
+            .map_err(|e| ttrpc_error!(ttrpc::Code::INTERNAL, e))?;
+
+        resp.usage = RepeatedField::from_vec(usage_vec);
+        resp.volume_condition = SingularPtrField::some(condition);
+        Ok(resp)
+    }
+
     async fn add_swap(
         &self,
         ctx: &TtrpcContext,
@@ -1341,6 +1386,48 @@ fn get_memory_info(block_size: bool, hotplug: bool) -> Result<(u64, bool)> {
     Ok((size, plug))
 }
 
+fn get_volume_capacity_stats(path: &str) -> Result<VolumeUsage> {
+    let mut usage = VolumeUsage::new();
+
+    let s = System::new();
+    for disk in s.disks() {
+        if let Some(v) = disk.name().to_str() {
+            if v.to_string().eq(path) {
+                usage.available = disk.available_space();
+                usage.total = disk.total_space();
+                usage.used = usage.total - usage.available;
+                usage.unit = VolumeUsage_Unit::BYTES; // bytes
+                break;
+            }
+        } else {
+            return Err(anyhow!(nix::Error::EINVAL));
+        }
+    }
+
+    Ok(usage)
+}
+
+fn get_volume_inode_stats(path: &str) -> Result<VolumeUsage> {
+    let mut usage = VolumeUsage::new();
+
+    let s = System::new();
+    for disk in s.disks() {
+        if let Some(v) = disk.name().to_str() {
+            if v.to_string().eq(path) {
+                let meta = fs::metadata(disk.mount_point())?;
+                let inode = meta.ino();
+                usage.used = inode;
+                usage.unit = VolumeUsage_Unit::INODES;
+                break;
+            }
+        } else {
+            return Err(anyhow!(nix::Error::EINVAL));
+        }
+    }
+
+    Ok(usage)
+}
+
 pub fn have_seccomp() -> bool {
     if cfg!(feature = "seccomp") {
         return true;

src/tools/agent-ctl/src/client.rs

@@ -171,6 +171,11 @@ static AGENT_CMDS: &[AgentCmd] = &[
         st: ServiceType::Agent,
         fp: agent_cmd_sandbox_get_oom_event,
     },
+    AgentCmd {
+        name: "GetVolumeStats",
+        st: ServiceType::Agent,
+        fp: agent_cmd_sandbox_get_volume_stats,
+    },
     AgentCmd {
         name: "ListInterfaces",
         st: ServiceType::Agent,
@@ -1641,6 +1646,29 @@ fn agent_cmd_sandbox_get_oom_event(
     Ok(())
 }
 
+fn agent_cmd_sandbox_get_volume_stats(
+    ctx: &Context,
+    client: &AgentServiceClient,
+    _health: &HealthClient,
+    _options: &mut Options,
+    args: &str,
+) -> Result<()> {
+    let req: VolumeStatsRequest = utils::make_request(args)?;
+
+    let ctx = clone_context(ctx);
+
+    debug!(sl!(), "sending request"; "request" => format!("{:?}", req));
+
+    let reply = client
+        .get_volume_stats(ctx, &req)
+        .map_err(|e| anyhow!(e).context(ERR_API_FAILED))?;
+
+    info!(sl!(), "response received";
+        "response" => format!("{:?}", reply));
+
+    Ok(())
+}
+
 fn agent_cmd_sandbox_copy_file(
     ctx: &Context,
     client: &AgentServiceClient,