Add "kernel_heap_monitor" feature

.github/workflows/main.yml

@@ -68,6 +68,12 @@ jobs:
       if: ${{ failure() }}
       run: docker exec ${{ github.job }} bash -c "cat /root/occlum/build/test/.fail"
 
+    - name: Integration test with optional Occlum features
+      run: |
+        docker exec ${{ github.job }} bash -c 'source /opt/intel/sgxsdk/environment; cd /root/occlum; make clean && LIBOS_FEATURES="kernel_heap_monitor" make install'
+        docker exec ${{ github.job }} bash -c "cd /root/occlum; SGX_MODE=SIM make test-glibc"
+      shell: bash
+
   # Make_test_on_centos:
   #   runs-on: ubuntu-18.04
 

src/libos/Cargo.toml

@@ -50,6 +50,7 @@ dcap = []               # DCAP support. The compilation relies on DCAP package.
 cov = ["sgx_cov"]       # Enable coverage colletcion.
 hyper_mode = []         # For running in hyper mode.
 pku = []                # PKU Support
+kernel_heap_monitor = []# Kernel heap usage tracking. With overhead.
 
 [target.'cfg(not(target_env = "sgx"))'.dependencies]
 sgx_types = { path = "../../deps/rust-sgx-sdk/sgx_types" }

src/libos/Makefile

@@ -45,7 +45,7 @@ LIBOS_LOG ?= error
 
 LIBOS_SONAME := libocclum-libos.so.$(MAJOR_VER_NUM)
 
-LIBOS_FEATURES :=
+LIBOS_FEATURES := $(LIBOS_FEATURES)
 
 ifeq ($(SGX_MODE), HW)
 	LIBOS_CORE_LIB_NAME := occlum-libos-core

src/libos/src/fs/procfs/meminfo.rs

@@ -1,4 +1,5 @@
 use super::*;
+use crate::util::kernel_alloc::KernelAlloc;
 use crate::vm::USER_SPACE_VM_MANAGER;
 
 pub struct MemInfoINode;
@@ -15,13 +16,26 @@ impl ProcINode for MemInfoINode {
     fn generate_data_in_bytes(&self) -> vfs::Result<Vec<u8>> {
         let total_ram = USER_SPACE_VM_MANAGER.get_total_size();
         let free_ram = USER_SPACE_VM_MANAGER.get_precise_free_size();
+        let kernel_heap_total = KernelAlloc::get_kernel_heap_config();
+        let kernel_heap_peak_used = KernelAlloc::get_kernel_heap_peak_used();
+        let kernel_heap_in_use = if let Some(bytes) = KernelAlloc::get_kernel_mem_size() {
+            format!("{} kB", bytes / KB)
+        } else {
+            "Feature not enabled".to_string()
+        };
         Ok(format!(
-            "MemTotal:       {} kB\n\
-             MemFree:        {} kB\n\
-             MemAvailable:   {} kB\n",
+            "MemTotal:              {} kB\n\
+             MemFree:               {} kB\n\
+             MemAvailable:          {} kB\n\
+             KernelHeapTotal:       {} kB\n\
+             KernelHeapPeakUsed:    {} kB\n\
+             KernelHeapInUse:       {}\n",
             total_ram / KB,
             free_ram / KB,
             free_ram / KB,
+            kernel_heap_total / KB,
+            kernel_heap_peak_used / KB,
+            kernel_heap_in_use,
         )
         .into_bytes())
     }

src/libos/src/util/kernel_alloc.rs

@@ -0,0 +1,81 @@
+use std::alloc::{GlobalAlloc, Layout, System};
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+// This file provides "KernelAlloc", a wrapper for sgx_std "System" allocator, is used as
+// the global allocator for Occlum kernel. Currently, this can provides the ability to
+// monitor the kernel heap usage.
+
+pub struct KernelAlloc {
+    size: AtomicUsize,
+}
+
+impl KernelAlloc {
+    pub fn get_kernel_mem_size() -> Option<usize> {
+        cfg_if! {
+            if #[cfg(feature = "kernel_heap_monitor")] {
+                Some(ALLOC.size.load(Ordering::Relaxed))
+            } else {
+                None
+            }
+        }
+    }
+
+    pub fn get_kernel_heap_config() -> usize {
+        std::enclave::get_heap_size()
+    }
+
+    pub fn get_kernel_heap_peak_used() -> usize {
+        sgx_trts::enclave::rsgx_get_peak_heap_used()
+    }
+}
+
+unsafe impl GlobalAlloc for KernelAlloc {
+    #[inline]
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        let ptr = System.alloc(layout);
+        if !ptr.is_null() {
+            self.size.fetch_add(layout.size(), Ordering::Relaxed);
+        }
+
+        ptr
+    }
+
+    #[inline]
+    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+        let ptr = System.alloc_zeroed(layout);
+        if !ptr.is_null() {
+            self.size.fetch_add(layout.size(), Ordering::Relaxed);
+        }
+
+        ptr
+    }
+
+    #[inline]
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        System.dealloc(ptr, layout);
+        self.size.fetch_sub(layout.size(), Ordering::Relaxed);
+    }
+
+    #[inline]
+    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+        let ptr = System.realloc(ptr, layout, new_size);
+        if !ptr.is_null() {
+            let old_size = layout.size();
+            if new_size > old_size {
+                // grow
+                self.size.fetch_add(new_size - old_size, Ordering::Relaxed);
+            } else if new_size < old_size {
+                // shrink
+                self.size.fetch_sub(old_size - new_size, Ordering::Relaxed);
+            }
+        }
+
+        ptr
+    }
+}
+
+#[cfg(feature = "kernel_heap_monitor")]
+#[global_allocator]
+static ALLOC: KernelAlloc = KernelAlloc {
+    size: AtomicUsize::new(0),
+};

src/libos/src/util/mod.rs

@@ -3,6 +3,7 @@ use super::*;
 pub mod dirty;
 pub mod host_file_util;
 pub mod hosts_parser_util;
+pub mod kernel_alloc;
 pub mod log;
 pub mod mem_util;
 pub mod mpx_util;