Add InitAPI call count to workaround multiple init mis-use

aws · Oct 12, 2023 · 1f49f91 · 1f49f91
1 parent b0ee1c0
commit 1f49f91
Show file tree

Hide file tree

Showing 2 changed files with 127 additions and 0 deletions.
diff --git a/src/aws-cpp-sdk-core/source/Aws.cpp b/src/aws-cpp-sdk-core/source/Aws.cpp
@@ -26,8 +26,20 @@ namespace Aws
 {
     static const char* ALLOCATION_TAG = "Aws_Init_Cleanup";
 
+    static std::mutex s_initShutdownMutex;
+    static size_t s_initCount = 0;
+
     void InitAPI(const SDKOptions &options)
     {
+        std::unique_lock<std::mutex> lock(s_initShutdownMutex);
+        s_initCount += 1;
+        if(s_initCount != 1)
+        {
+            AWS_LOGSTREAM_ERROR(ALLOCATION_TAG, "AWS-SDK-CPP is already initialized " << s_initCount - 1 << " times. "
+                                                "Consequent calls to InitAPI are ignored.");
+            return;
+        }
+
 #ifdef USE_AWS_MEMORY_MANAGEMENT
         if(options.memoryManagementOptions.memoryManager)
         {
@@ -172,6 +184,20 @@ namespace Aws
 
     void ShutdownAPI(const SDKOptions& options)
     {
+        std::unique_lock<std::mutex> lock(s_initShutdownMutex);
+        if(s_initCount != 1)
+        {
+            if(!s_initCount) {
+                AWS_LOGSTREAM_ERROR(ALLOCATION_TAG, "Unable to ShutdownAPI of AWS-SDK-CPP: the SDK was not initialized.");
+            } else {
+                AWS_LOGSTREAM_ERROR(ALLOCATION_TAG, "AWS-SDK-CPP: this call to ShutdownAPI is ignored, current init count = " << s_initCount);
+                s_initCount -= 1;
+            }
+            return;
+        } else {
+            AWS_LOGSTREAM_INFO(ALLOCATION_TAG, "Shutdown AWS SDK for C++.");
+        }
+        s_initCount -= 1;
         Aws::Utils::ComponentRegistry::TerminateAllComponents();
         Aws::Utils::ComponentRegistry::ShutdownComponentRegistry();
         Aws::Monitoring::CleanupMonitoring();

diff --git a/tests/aws-cpp-sdk-core-tests/utils/memory/AwsMemoryManagementTest.cpp b/tests/aws-cpp-sdk-core-tests/utils/memory/AwsMemoryManagementTest.cpp
@@ -40,6 +40,107 @@ TEST(AwsMemoryManagementTest, InitShutdownRepeatable)
     ASSERT_TRUE(memorySystem.IsClean());
 }
 
+TEST(AwsMemoryManagementTest, MultiInit)
+{
+    Aws::SDKOptions options;
+    options.loggingOptions.logLevel = Aws::Utils::Logging::LogLevel::Trace;
+    options.httpOptions.installSigPipeHandler = true;
+
+    const uint32_t BUCKET_COUNT = 1024;
+    const uint32_t TRACKER_PER_BLOCK = 128;
+    ExactTestMemorySystem memorySystem(BUCKET_COUNT, TRACKER_PER_BLOCK); \
+    options.memoryManagementOptions.memoryManager = &memorySystem;
+
+    for (size_t ii = 0; ii < 10; ++ii)
+    {
+        Aws::InitAPI(options);
+        for (unsigned jj = 0; jj < 10; ++jj) {
+            AWS_LOG_WARN("AwsMemoryManagementTest", "MultiInit test warn level");
+            AWS_LOGF_TRACE(AWS_LS_COMMON_GENERAL, "AwsMemoryManagementTest::%s test log", "MultiInit");
+        }
+
+        Aws::InitAPI(options);
+        for (unsigned jj = 0; jj < 10; ++jj) {
+            AWS_LOG_WARN("AwsMemoryManagementTest", "MultiInit test warn level");
+            AWS_LOGF_TRACE(AWS_LS_COMMON_GENERAL, "AwsMemoryManagementTest::%s test log", "MultiInit");
+        }
+
+        Aws::ShutdownAPI(options);
+        for (unsigned jj = 0; jj < 10; ++jj) {
+            AWS_LOG_WARN("AwsMemoryManagementTest", "MultiInit after first shutdown test warn level");
+            AWS_LOGF_TRACE(AWS_LS_COMMON_GENERAL, "AwsMemoryManagementTest::%s after first shutdown test log", "MultiInit");
+        }
+        Aws::ShutdownAPI(options);
+    }
+
+    ASSERT_EQ(memorySystem.GetCurrentOutstandingAllocations(), 0ULL);
+    ASSERT_EQ(memorySystem.GetCurrentBytesAllocated(), 0ULL);
+    ASSERT_TRUE(memorySystem.IsClean());
+}
+
+TEST(AwsMemoryManagementTest, MultiInitParallel)
+{
+    Aws::SDKOptions options;
+    options.loggingOptions.logLevel = Aws::Utils::Logging::LogLevel::Trace;
+    options.httpOptions.installSigPipeHandler = true;
+
+    const uint32_t BUCKET_COUNT = 1024;
+    const uint32_t TRACKER_PER_BLOCK = 128;
+    ExactTestMemorySystem memorySystem(BUCKET_COUNT, TRACKER_PER_BLOCK); \
+    options.memoryManagementOptions.memoryManager = &memorySystem;
+
+    std::atomic<bool> isRunning;
+    isRunning.store(false);
+    auto InitShutdown = [&]()
+    {
+        while(isRunning)
+        {
+            Aws::ShutdownAPI(options);
+            Aws::InitAPI(options);
+            size_t sleepMs = 0 + std::rand() % ( 200 - 0 ); // random from 0 to 200
+            std::this_thread::sleep_for(std::chrono::milliseconds(sleepMs));
+            Aws::InitAPI(options);
+            AWS_LOG_FLUSH();
+
+            Aws::ShutdownAPI(options);
+            Aws::ShutdownAPI(options);
+        }
+        return true;
+    };
+
+    auto threadLogger = [&]()
+    {
+      while(isRunning)
+      {
+        AWS_LOG_WARN("AwsMemoryManagementTest", "MultiInitParallel test logging in thread");
+        AWS_LOGF_TRACE(AWS_LS_COMMON_GENERAL, "AwsMemoryManagementTest::%s CRT logging in thread", "MultiInitParallel");
+        AWS_LOG_FLUSH();
+        std::this_thread::sleep_for(std::chrono::milliseconds(5));
+      }
+      return true;
+    };
+
+    for (size_t ii = 0; ii < 5; ++ii)
+    {
+      isRunning.store(true);
+      std::future<bool> initShutdownFuture = std::async(std::launch::async, InitShutdown);
+
+      std::future<bool> loggerThread1Future = std::async(std::launch::async, threadLogger);
+      std::future<bool> loggerThread2Future = std::async(std::launch::async, threadLogger);
+      size_t sleepMs = 500 + std::rand() % ( 800 - 500 ); // random from 500 to 800
+      std::this_thread::sleep_for(std::chrono::milliseconds(sleepMs));
+      isRunning.store(false);
+
+      ASSERT_TRUE(initShutdownFuture.get());
+      ASSERT_TRUE(loggerThread1Future.get());
+      ASSERT_TRUE(loggerThread2Future.get());
+    }
+
+    ASSERT_EQ(memorySystem.GetCurrentOutstandingAllocations(), 0ULL);
+    ASSERT_EQ(memorySystem.GetCurrentBytesAllocated(), 0ULL);
+    ASSERT_TRUE(memorySystem.IsClean());
+}
+
 #ifdef USE_AWS_MEMORY_MANAGEMENT
 class TestMemoryManager : public Aws::Utils::Memory::MemorySystemInterface
 {