Introduce a tool to display memory usage.

1. minor fix 2. add max heap usage info 3. add a script to display memory usage. Signed-off-by: xiaoyuxlu <[email protected]>
intel · Dec 25, 2023 · 2fb434b · 2fb434b
1 parent 05c41c6
commit 2fb434b
Show file tree

Hide file tree

Showing 6 changed files with 156 additions and 6 deletions.
diff --git a/readme.md b/readme.md
@@ -245,5 +245,15 @@ To run a specific test, use `cargo test <test_func_name>`
 
 To run test with println!() message, use `cargo test -- --nocapture`
 
+## Collect memory usage
+
+To collect memory usage, use
+
+```
+python sh_script/collect_memory_usage.py
+```
+
+This script will display the peak memory used by spdm-emu
+
 ## Known limitation
 This package is only the sample code to show the concept. It does not have a full validation such as robustness functional test and fuzzing test. It does not meet the production quality yet. Any codes including the API definition, the libary and the drivers are subject to change.
diff --git a/sh_script/collect_memory_usage.py b/sh_script/collect_memory_usage.py
@@ -0,0 +1,115 @@
+#
+# Copyright (c) 2023 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0 or MIT
+#
+
+import subprocess
+import threading
+import os
+import re
+import time
+import queue
+
+
+def parse_number(s: str, pattern: re.Pattern) -> int:
+    match = re.search(pattern, s, re.DOTALL)
+    assert match
+    value = match.group(1)
+    return int(value)
+
+
+def run_exec(shell_cmd: str, result_queue) -> str:
+    r = subprocess.run(shell_cmd, shell=True, stdout=subprocess.PIPE, )
+    out = r.stdout.decode().strip()
+    result_queue.put(out)
+    return out
+
+
+def memory_usage(name, requester_cmd, responder_cmd) -> (str, int, int, int):
+    result_responder = queue.Queue()
+    result_requester = queue.Queue()
+
+    responder = threading.Thread(
+        target=run_exec, args=(responder_cmd, result_responder))
+    responder.start()
+    time.sleep(10)
+    requester = threading.Thread(
+        target=run_exec, args=(requester_cmd, result_requester))
+    requester.start()
+
+    requester.join()
+    responder.join()
+
+    out_responder = result_responder.get()
+    out_requester = result_requester.get()
+
+    max_stack_usage_responder = parse_number(
+        out_responder, r"max stack usage: (\d+)")
+    max_heap_usage_responder = parse_number(
+        out_responder, r"max heap usage: (\d+)")
+    max_stack_usage_requester = parse_number(
+        out_requester, r"max stack usage: (\d+)")
+    max_heap_usage_requester = parse_number(
+        out_requester, r"max heap usage: (\d+)")
+
+    return (name, max_stack_usage_requester, max_heap_usage_requester, max_stack_usage_responder, max_heap_usage_responder)
+
+
+def main():
+    test_vector = [
+        (
+            " async + release         ",
+            "cargo run --release -p spdm-requester-emu --no-default-features --features=spdm-ring,hashed-transcript-data,async-executor,test_stack_size,test_heap_size",
+            "cargo run --release -p spdm-responder-emu --no-default-features --features=spdm-ring,hashed-transcript-data,async-executor,test_stack_size,test_heap_size"
+        ),
+        (
+            " tokio + release         ",
+            "cargo run --release -p spdm-requester-emu --no-default-features --features=spdm-ring,hashed-transcript-data,async-tokio,test_stack_size,test_heap_size",
+            "cargo run --release -p spdm-responder-emu --no-default-features --features=spdm-ring,hashed-transcript-data,async-tokio,test_stack_size,test_heap_size"
+        ),
+        (
+            "async + releas + w/o hash",
+            "cargo run --release -p spdm-requester-emu --no-default-features --features=spdm-ring,async-executor,test_stack_size,test_heap_size",
+            "cargo run --release -p spdm-responder-emu --no-default-features --features=spdm-ring,async-executor,test_stack_size,test_heap_size"
+        ),
+        (
+            " sync + release          ",
+            "cargo run --release -p spdm-requester-emu --no-default-features --features=spdm-ring,hashed-transcript-data,async-executor,test_stack_size,test_heap_size,is_sync",
+            "cargo run --release -p spdm-responder-emu --no-default-features --features=spdm-ring,hashed-transcript-data,async-executor,test_stack_size,test_heap_size,is_sync"
+        ),
+        (
+            "sync + release + w/o hash",
+            "cargo run --release -p spdm-requester-emu --no-default-features --features=spdm-ring,test_stack_size,test_heap_size,is_sync",
+            "cargo run --release -p spdm-responder-emu --no-default-features --features=spdm-ring,test_stack_size,test_heap_size,is_sync"
+        ),
+
+        (
+            " async + debug           ",
+            "cargo run -p spdm-requester-emu --no-default-features --features=spdm-ring,hashed-transcript-data,async-executor,test_stack_size,test_heap_size",
+            "cargo run -p spdm-responder-emu --no-default-features --features=spdm-ring,hashed-transcript-data,async-executor,test_stack_size,test_heap_size"
+        ),
+
+        (
+            " sync + debug            ",
+            "cargo run -p spdm-requester-emu --no-default-features --features=spdm-ring,hashed-transcript-data,async-executor,test_stack_size,test_heap_size,is_sync",
+            "cargo run -p spdm-responder-emu --no-default-features --features=spdm-ring,hashed-transcript-data,async-executor,test_stack_size,test_heap_size,is_sync"
+        ),
+
+    ]
+    results = []
+    for t in test_vector:
+        result = memory_usage(*t)
+        results.append(result)
+
+    print("""
+|                         |      Requester        |      Responder        |
+|                         | stack     |  heap     | stack     |  heap     |
+| ----------------------- |-----------|-----------|-----------|-----------|""")
+    for r in results:
+        print(
+            "|{}| {:10}| {:10}| {:10}| {:10}|".format(*r))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/test/spdm-requester-emu/Cargo.toml b/test/spdm-requester-emu/Cargo.toml
@@ -22,9 +22,10 @@ tokio = { version = "1.30.0", features = ["full"] }
 executor = { path = "../../executor" }
 maybe-async = "0.2.7"
 td-benchmark = { git = "https://github.com/confidential-containers/td-shim.git", default-features = false, optional = true }
+dhat = { version = "0.3.2", optional = true }
 
 [features]
-default = ["spdm-emu/default", "async-executor", "test_stack_size"]
+default = ["spdm-emu/default", "async-executor"]
 mut-auth = ["spdm-emu/mut-auth"]
 spdm-ring = ["spdm-emu/spdm-ring"]
 spdm-mbedtls = ["spdm-emu/spdm-mbedtls"]
@@ -33,3 +34,4 @@ async-executor = ["spdm-emu/async-executor"]
 async-tokio = ["spdm-emu/async-tokio"]
 is_sync = ["spdm-emu/is_sync", "spdmlib/is_sync", "maybe-async/is_sync", "idekm/is_sync", "tdisp/is_sync", "mctp_transport/is_sync", "pcidoe_transport/is_sync"]
 test_stack_size = ["td-benchmark"]
+test_heap_size = ["dhat"]
diff --git a/test/spdm-requester-emu/src/main.rs b/test/spdm-requester-emu/src/main.rs
@@ -428,7 +428,7 @@ async fn test_spdm(
     #[cfg(feature = "test_stack_size")]
     {
         let value = td_benchmark::StackProfiling::stack_usage().unwrap();
-        println!("max stack usage {}", value);
+        println!("max stack usage(no idekm): {}", value);
     }
 }
 
@@ -1283,7 +1283,7 @@ fn new_logger_from_env() -> SimpleLogger {
 fn emu_main() {
     td_benchmark::StackProfiling::init(
         0x5aa5_5aa5_5aa5_5aa5,
-        EMU_STACK_SIZE - 0x40000, // main function stack
+        EMU_STACK_SIZE - 0x60000, // main function stack
     );
     emu_main_inner()
 }
@@ -1371,19 +1371,29 @@ fn emu_main_inner() {
     #[cfg(feature = "test_stack_size")]
     {
         let value = td_benchmark::StackProfiling::stack_usage().unwrap();
-        println!("max stack usage {}", value);
+        println!("max stack usage: {}", value);
     }
 }
 
+#[cfg(feature = "test_heap_size")]
+#[global_allocator]
+static ALLOC: dhat::Alloc = dhat::Alloc;
+
 fn main() {
     use std::thread;
 
+    #[cfg(feature = "test_heap_size")]
+    let _profiler = dhat::Profiler::builder().testing().build();
+
     thread::Builder::new()
         .stack_size(EMU_STACK_SIZE)
         .spawn(emu_main)
         .unwrap()
         .join()
         .unwrap();
+
+    #[cfg(feature = "test_heap_size")]
+    log::info!("max heap usage: {}", dhat::HeapStats::get().max_bytes);
 }
 
 pub const MMIO_RANGE_COUNT: usize = 4;

diff --git a/test/spdm-responder-emu/Cargo.toml b/test/spdm-responder-emu/Cargo.toml
@@ -23,9 +23,10 @@ executor = { path = "../../executor" }
 zeroize = { version = "1.5.0", features = ["zeroize_derive"]}
 maybe-async = "0.2.7"
 td-benchmark = { git = "https://github.com/confidential-containers/td-shim.git", default-features = false, optional = true }
+dhat = { version = "0.3.2", optional = true }
 
 [features]
-default = ["spdm-emu/default", "async-executor", "test_stack_size"]
+default = ["spdm-emu/default", "async-executor"]
 mut-auth = ["spdm-emu/mut-auth", "async-executor"]
 mandatory-mut-auth = ["mut-auth", "spdm-emu/mandatory-mut-auth"]
 spdm-ring = ["spdm-emu/spdm-ring"]
@@ -35,3 +36,4 @@ async-executor = ["spdm-emu/async-executor"]
 async-tokio = ["spdm-emu/async-tokio"]
 is_sync = ["spdm-emu/is_sync", "spdmlib/is_sync", "maybe-async/is_sync", "idekm/is_sync", "tdisp/is_sync", "mctp_transport/is_sync", "pcidoe_transport/is_sync"]
 test_stack_size = ["td-benchmark"]
+test_heap_size = ["dhat"]
diff --git a/test/spdm-responder-emu/src/main.rs b/test/spdm-responder-emu/src/main.rs
@@ -123,7 +123,7 @@ fn new_logger_from_env() -> SimpleLogger {
 fn emu_main() {
     td_benchmark::StackProfiling::init(
         0x5aa5_5aa5_5aa5_5aa5,
-        EMU_STACK_SIZE - 0x40000, // main function stack
+        EMU_STACK_SIZE - 0x60000, // main function stack
     );
     emu_main_inner()
 }
@@ -555,8 +555,16 @@ fn pci_idekm_tdisp_rsp_dispatcher(
     }
 }
 
+#[cfg(feature = "test_heap_size")]
+#[global_allocator]
+static ALLOC: dhat::Alloc = dhat::Alloc;
+
 fn main() {
     use std::thread;
+
+    #[cfg(feature = "test_heap_size")]
+    let _profiler = dhat::Profiler::builder().testing().build();
+
     init_device_idekm_instance();
     init_device_tdisp_instance();
 
@@ -566,4 +574,7 @@ fn main() {
         .unwrap()
         .join()
         .unwrap();
+
+    #[cfg(feature = "test_heap_size")]
+    log::info!("max heap usage: {}", dhat::HeapStats::get().max_bytes);
 }