Prevent SDK allocator being used before InitAPI in case of custom mem…

…ory management
aws · Apr 18, 2024 · f9b3244 · f9b3244
1 parent 26287d2
commit f9b3244
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 9 deletions.
diff --git a/src/aws-cpp-sdk-core/include/aws/core/utils/memory/stl/AWSAllocator.h b/src/aws-cpp-sdk-core/include/aws/core/utils/memory/stl/AWSAllocator.h
@@ -103,6 +103,10 @@ namespace Aws
     template<typename T, typename ...ArgTypes>
     std::shared_ptr<T> MakeShared(const char* allocationTag, ArgTypes&&... args)
     {
+#ifdef USE_AWS_MEMORY_MANAGEMENT
+        Aws::Utils::Memory::MemorySystemInterface* memorySystem = Aws::Utils::Memory::GetMemorySystem();
+        assert(memorySystem && "Memory system is not initialized");
+#endif
         AWS_UNREFERENCED_PARAM(allocationTag);
 
         return std::allocate_shared<T, Aws::Allocator<T>>(Aws::Allocator<T>(), std::forward<ArgTypes>(args)...);

diff --git a/src/aws-cpp-sdk-core/source/utils/memory/AWSMemory.cpp b/src/aws-cpp-sdk-core/source/utils/memory/AWSMemory.cpp
@@ -65,6 +65,10 @@ MemorySystemInterface* GetMemorySystem()
 void* Malloc(const char* allocationTag, size_t allocationSize)
 {
     Aws::Utils::Memory::MemorySystemInterface* memorySystem = Aws::Utils::Memory::GetMemorySystem();
+#ifdef USE_AWS_MEMORY_MANAGEMENT
+    // Was InitAPI forgotten or ShutdownAPI already called?
+    assert(memorySystem && "Memory system is not initialized.");
+#endif
 
     void* rawMemory = nullptr;
     if(memorySystem != nullptr)

diff --git a/tests/aws-cpp-sdk-s3-unit-tests/S3UnitTests.cpp b/tests/aws-cpp-sdk-s3-unit-tests/S3UnitTests.cpp
@@ -19,18 +19,14 @@ class S3UnitTest : public testing::Test {
 protected:
   static void SetUpTestSuite() {
 #ifdef USE_AWS_MEMORY_MANAGEMENT
-    _testMemorySystem = Aws::MakeShared<ExactTestMemorySystem>(ALLOCATION_TAG, 1024, 128);
+    _testMemorySystem.reset(new ExactTestMemorySystem(1024, 128));
     _options.memoryManagementOptions.memoryManager = _testMemorySystem.get();
 #endif
+    InitAPI(_options);
     _mockClientFactory = Aws::MakeShared<MockHttpClientFactory>(ALLOCATION_TAG);
     _mockHttpClient = Aws::MakeShared<MockHttpClient>(ALLOCATION_TAG);
     _mockClientFactory->SetClient(_mockHttpClient);
-    HttpOptions httpOptions;
-    httpOptions.httpClientFactory_create_fn = []() -> std::shared_ptr<HttpClientFactory> {
-      return _mockClientFactory;
-    };
-    _options.httpOptions = httpOptions;
-    InitAPI(_options);
+    SetHttpClientFactory(_mockClientFactory);
     AWSCredentials credentials{"mock", "credentials"};
     const auto epProvider = Aws::MakeShared<S3EndpointProvider>(ALLOCATION_TAG);
     S3ClientConfiguration s3Config;
@@ -62,7 +58,7 @@ class S3UnitTest : public testing::Test {
   static std::shared_ptr<MockHttpClientFactory> _mockClientFactory;
   static std::shared_ptr<S3Client> _s3Client;
 #ifdef USE_AWS_MEMORY_MANAGEMENT
-  static std::shared_ptr<ExactTestMemorySystem> _testMemorySystem;
+  static std::unique_ptr<ExactTestMemorySystem> _testMemorySystem;
 #endif
 };
 
@@ -71,7 +67,8 @@ std::shared_ptr<MockHttpClient> S3UnitTest::_mockHttpClient = nullptr;
 std::shared_ptr<MockHttpClientFactory> S3UnitTest::_mockClientFactory = nullptr;
 std::shared_ptr<S3Client> S3UnitTest::_s3Client = nullptr;
 #ifdef USE_AWS_MEMORY_MANAGEMENT
-std::shared_ptr<ExactTestMemorySystem> S3UnitTest::_testMemorySystem = nullptr;
+// this must be std:: because this is an utility to track allocations in the SDK and should not rely on SDK
+std::unique_ptr<ExactTestMemorySystem> S3UnitTest::_testMemorySystem = nullptr;
 #endif