Fix incorrect aws sdk use (#106)

* add a helper class to init/shutdown aws api properly * add helper class * remove logging in AwsSDKHelper * Use an atomic reference count Use an atomic reference counter for SDK usage count. Delay SDK initialization to later. Move helper class to TSCommunication's cpp file. * add a test case to ensure multi connection uses AWS SDK API appropriately Co-authored-by: James Duong <[email protected]>
aws · Jul 8, 2021 · 54e6def · 54e6def
1 parent 92d3d5d
commit 54e6def
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 10 deletions.
diff --git a/src/odfesqlodbc/communication.cpp b/src/odfesqlodbc/communication.cpp
@@ -24,13 +24,9 @@
 Communication::Communication()
     : m_status(ConnStatusType::CONNECTION_BAD),
       m_client_encoding(m_supported_client_encodings[0]) {
-    LogMsg(LOG_DEBUG, "Initializing AWS API.");
-    Aws::InitAPI(m_sdk_options);
 }
 
 Communication::~Communication() {
-    LogMsg(LOG_DEBUG, "Shutting down AWS API.");
-    Aws::ShutdownAPI(m_sdk_options);
 }
 
 bool Communication::Setup(const runtime_options& options) {

diff --git a/src/odfesqlodbc/communication.h b/src/odfesqlodbc/communication.h
@@ -17,8 +17,6 @@
 #ifndef COMMUNICATION
 #define COMMUNICATION
 
-#include <aws/core/Aws.h>
-
 #include <map>
 #include <memory>
 #include <string>
@@ -135,10 +133,6 @@ class Communication {
      * Current client encoding
      */
     std::string m_client_encoding;
-    /**
-     * AWS sdk options
-     */
-    Aws::SDKOptions m_sdk_options;
     /**
      * Map storing prefetch queues
      */

diff --git a/src/odfesqlodbc/odbc_communication.cpp b/src/odfesqlodbc/odbc_communication.cpp
@@ -25,7 +25,10 @@
 #include "okta_credentials_provider.h"
 #include "version.h"
 #include "mylog.h"
+#include <atomic>
 #include <memory>
+#include <mutex>
+#include <aws/core/Aws.h>
 #include <aws/core/auth/AWSCredentials.h>
 #include <aws/core/auth/AWSCredentialsProvider.h>
 #include <aws/core/client/DefaultRetryStrategy.h>
@@ -34,6 +37,38 @@
 // clang-format on
 
 namespace {
+    /**
+     * A helper class to initialize/shutdown AWS API once per DLL load/unload.
+     */
+    class AwsSdkHelper {
+      public:
+        AwsSdkHelper() :
+          m_reference_count(0) {
+        }
+
+        AwsSdkHelper& operator++() {
+          if (1 == ++m_reference_count) {
+            std::scoped_lock lock(m_mutex);
+            Aws::InitAPI(m_sdk_options);
+          }
+          return *this;
+        }
+
+        AwsSdkHelper& operator--() {
+          if (0 == --m_reference_count) {
+            std::scoped_lock lock(m_mutex);
+            Aws::ShutdownAPI(m_sdk_options);
+          }
+          return *this;
+        }
+
+        Aws::SDKOptions m_sdk_options;
+        std::atomic<int> m_reference_count;
+        std::mutex m_mutex;
+    };
+
+    AwsSdkHelper AWS_SDK_HELPER;
+
     const Aws::String UA_ID_PREFIX = Aws::String("ts-odbc.");
     const std::string DEFAULT_CREATOR_TYPE = "DEFAULT";
 
@@ -93,6 +128,14 @@ namespace {
     };
 }
 
+TSCommunication::TSCommunication() {
+  ++AWS_SDK_HELPER;
+}
+
+TSCommunication::~TSCommunication() {
+  --AWS_SDK_HELPER;
+}
+
 bool TSCommunication::Validate(const runtime_options& options) {
     if (options.auth.region.empty() && options.auth.end_point_override.empty()) {
         throw std::invalid_argument("Both region and end point cannot be empty.");

diff --git a/src/odfesqlodbc/odbc_communication.h b/src/odfesqlodbc/odbc_communication.h
@@ -32,6 +32,10 @@
  */
 class TSCommunication : public Communication {
    public:
+    TSCommunication();
+
+    ~TSCommunication();
+
     /**
      * Validate options
      * @param options const runtime_options&