From 257969b85c4360116f6f44e625dec22f16dedd69 Mon Sep 17 00:00:00 2001
From: Dmitriy Musatkin <musatkd@amazon.com>
Date: Fri, 20 Oct 2023 15:56:38 -0700
Subject: [PATCH] basic mem usage functions

---
 CMakeLists.txt                            |  3 ++-
 include/aws/common/system_resource_util.h | 26 ++++++++++++++++++++++
 source/posix/system_resource_utils.c      | 25 +++++++++++++++++++++
 source/windows/system_resource_utils.c    | 27 +++++++++++++++++++++++
 tests/CMakeLists.txt                      |  2 ++
 tests/system_resource_util_test.c         | 22 ++++++++++++++++++
 6 files changed, 104 insertions(+), 1 deletion(-)
 create mode 100644 include/aws/common/system_resource_util.h
 create mode 100644 source/posix/system_resource_utils.c
 create mode 100644 source/windows/system_resource_utils.c
 create mode 100644 tests/system_resource_util_test.c

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 0ff1d9e7e..85ebf6bf1 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -82,7 +82,8 @@ if (WIN32)
     endif ()
 
     list(APPEND PLATFORM_DEFINES WINDOWS_KERNEL_LIB=${WINDOWS_KERNEL_LIB})
-    list(APPEND PLATFORM_LIBS BCrypt ${WINDOWS_KERNEL_LIB} Ws2_32 Shlwapi)
+    list(APPEND PLATFORM_DEFINES PSAPI_VERSION=1)
+    list(APPEND PLATFORM_LIBS BCrypt ${WINDOWS_KERNEL_LIB} Ws2_32 Shlwapi Psapi)
 else ()
     file(GLOB AWS_COMMON_OS_HEADERS
         "include/aws/common/posix/*"
diff --git a/include/aws/common/system_resource_util.h b/include/aws/common/system_resource_util.h
new file mode 100644
index 000000000..db63608f8
--- /dev/null
+++ b/include/aws/common/system_resource_util.h
@@ -0,0 +1,26 @@
+#ifndef AWS_COMMON_SYSTEM_RESOURCE_UTIL_H
+#define AWS_COMMON_SYSTEM_RESOURCE_UTIL_H
+
+/**
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+#include <aws/common/common.h>
+
+AWS_PUSH_SANE_WARNING_LEVEL
+
+AWS_EXTERN_C_BEGIN
+
+struct aws_resource_usage {
+    size_t maxrss;
+
+    size_t _reserved[8];
+};
+
+AWS_COMMON_API int aws_resource_usage_for_current_process(struct aws_resource_usage *resource_usage);
+
+
+AWS_EXTERN_C_END
+AWS_POP_SANE_WARNING_LEVEL
+
+#endif /* AWS_COMMON_SYSTEM_RESOURCE_UTIL_H */
diff --git a/source/posix/system_resource_utils.c b/source/posix/system_resource_utils.c
new file mode 100644
index 000000000..042d4a182
--- /dev/null
+++ b/source/posix/system_resource_utils.c
@@ -0,0 +1,25 @@
+/**
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+
+#include <aws/common/system_resource_util.h>
+
+#include <sys/resource.h>
+
+int aws_resource_usage_for_current_process(struct aws_resource_usage *ru) {
+    AWS_PRECONDITION(ru);
+
+    struct rusage usage;
+
+    if (getrusage(RUSAGE_SELF, &usage)) {
+        return aws_raise_error(AWS_ERROR_SYS_CALL_FAILURE);
+    }
+
+#if defined(AWS_OS_APPLE)
+    ru->maxrss = usage.ru_maxrss / 1024;
+#else
+    ru->maxrss = usage.ru_maxrss;
+#endif 
+    return AWS_OP_SUCCESS;
+}
diff --git a/source/windows/system_resource_utils.c b/source/windows/system_resource_utils.c
new file mode 100644
index 000000000..00c0e7502
--- /dev/null
+++ b/source/windows/system_resource_utils.c
@@ -0,0 +1,27 @@
+/**
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+
+#include <aws/common/system_resource_util.h>
+
+#include <windows.h>
+#include <psapi.h>
+
+int aws_resource_usage_for_current_process(struct aws_resource_usage *resource_usage) {
+    AWS_PRECONDITION(resource_usage);
+
+    HANDLE hProcess =  GetCurrentProcess();;
+    PROCESS_MEMORY_COUNTERS pmc;
+
+    BOOL ret = GetProcessMemoryInfo(hProcess, &pmc, sizeof(pmc));
+    CloseHandle( hProcess ); 
+    
+    if (!ret) {
+        return aws_raise_error(AWS_ERROR_SYS_CALL_FAILURE);
+    }
+
+    ru->maxrss = pmc.PeakWorkingSetSize;
+
+    return AWS_OP_SUCCESS;
+}
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index 4681e04c3..907432b58 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -402,6 +402,8 @@ add_test_case(ring_buffer_acquire_up_to_multi_threaded_test)
 add_test_case(string_to_log_level_success_test)
 add_test_case(string_to_log_level_failure_test)
 
+add_test_case(test_resource_usage_maxrss)
+
 if(NOT ANDROID)
     add_test_case(test_logging_filter_at_AWS_LL_NONE_s_logf_all_levels)
     add_test_case(test_logging_filter_at_AWS_LL_FATAL_s_logf_all_levels)
diff --git a/tests/system_resource_util_test.c b/tests/system_resource_util_test.c
new file mode 100644
index 000000000..119538073
--- /dev/null
+++ b/tests/system_resource_util_test.c
@@ -0,0 +1,22 @@
+/**
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+
+#include <aws/common/system_resource_util.h>
+
+#include <aws/testing/aws_test_harness.h>
+
+static int s_test_resource_usage_maxrss(struct aws_allocator *allocator, void *ctx) {
+    (void)ctx;
+
+    struct aws_resource_usage ru;
+    AWS_ZERO_STRUCT(ru);
+    ASSERT_SUCCESS(aws_resource_usage_for_current_process(&ru));
+
+    ASSERT_TRUE(ru.maxrss > 0);
+
+    return 0;
+}
+
+AWS_TEST_CASE(test_resource_usage_maxrss, s_test_resource_usage_maxrss)
\ No newline at end of file