Fix memory leak issue with AWS API usage (#6160)

ydb-platform · Aug 7, 2024 · ff0fd33 · ff0fd33
1 parent cb6f3c8
commit ff0fd33
Show file tree

Hide file tree

Showing 25 changed files with 664 additions and 316 deletions.
diff --git a/ydb/core/driver_lib/run/kikimr_services_initializers.cpp b/ydb/core/driver_lib/run/kikimr_services_initializers.cpp
@@ -241,6 +241,25 @@
 
 #include <util/system/hostname.h>
 
+#include <aws/core/Aws.h>
+
+namespace {
+
+struct TAwsApiGuard {
+    TAwsApiGuard() {
+        Aws::InitAPI(Options);
+    }
+
+    ~TAwsApiGuard() {
+        Aws::ShutdownAPI(Options);
+    }
+
+private:
+    Aws::SDKOptions Options;
+};
+
+}
+
 namespace NKikimr {
 
 namespace NKikimrServicesInitializers {
@@ -2810,5 +2829,16 @@ void TGraphServiceInitializer::InitializeServices(NActors::TActorSystemSetup* se
         TActorSetupCmd(NGraph::CreateGraphService(appData->TenantName), TMailboxType::HTSwap, appData->UserPoolId));
 }
 
+TAwsApiInitializer::TAwsApiInitializer(IGlobalObjectStorage& globalObjects)
+    : GlobalObjects(globalObjects)
+{
+}
+
+void TAwsApiInitializer::InitializeServices(NActors::TActorSystemSetup* setup, const NKikimr::TAppData* appData) {
+    Y_UNUSED(setup);
+    Y_UNUSED(appData);
+    GlobalObjects.AddGlobalObject(std::make_shared<TAwsApiGuard>());
+}
+
 } // namespace NKikimrServicesInitializers
 } // namespace NKikimr
diff --git a/ydb/core/driver_lib/run/kikimr_services_initializers.h b/ydb/core/driver_lib/run/kikimr_services_initializers.h
@@ -619,5 +619,14 @@ class TGraphServiceInitializer : public IKikimrServicesInitializer {
     void InitializeServices(NActors::TActorSystemSetup* setup, const NKikimr::TAppData* appData) override;
 };
 
+class TAwsApiInitializer : public IServiceInitializer {
+    IGlobalObjectStorage& GlobalObjects;
+
+public:
+    TAwsApiInitializer(IGlobalObjectStorage& globalObjects);
+
+    void InitializeServices(NActors::TActorSystemSetup* setup, const NKikimr::TAppData* appData) override;
+};
+
 } // namespace NKikimrServicesInitializers
 } // namespace NKikimr
diff --git a/ydb/core/driver_lib/run/run.cpp b/ydb/core/driver_lib/run/run.cpp
@@ -1653,6 +1653,8 @@ TIntrusivePtr<TServiceInitializersList> TKikimrRunner::CreateServiceInitializers
         sil->AddServiceInitializer(new TGraphServiceInitializer(runConfig));
     }
 
+    sil->AddServiceInitializer(new TAwsApiInitializer(*this));
+
     return sil;
 }
 

diff --git a/ydb/core/driver_lib/run/ya.make b/ydb/core/driver_lib/run/ya.make
@@ -20,6 +20,7 @@ SRCS(
 )
 
 PEERDIR(
+    contrib/libs/aws-sdk-cpp/aws-cpp-sdk-core
     contrib/libs/protobuf
     ydb/library/actors/core
     ydb/library/actors/dnsresolver

diff --git a/ydb/core/tx/columnshard/ut_schema/ut_columnshard_schema.cpp b/ydb/core/tx/columnshard/ut_schema/ut_columnshard_schema.cpp
@@ -17,6 +17,9 @@
 
 #include <util/system/hostname.h>
 #include <library/cpp/deprecated/atomic/atomic.h>
+#include <library/cpp/testing/hook/hook.h>
+
+#include <aws/core/Aws.h>
 
 namespace NKikimr {
 
@@ -32,6 +35,16 @@ enum class EInitialEviction {
 
 namespace {
 
+Aws::SDKOptions Options;
+
+Y_TEST_HOOK_BEFORE_RUN(InitAwsAPI) {
+    Aws::InitAPI(Options);
+}
+
+Y_TEST_HOOK_AFTER_RUN(ShutdownAwsAPI) {
+    Aws::ShutdownAPI(Options);
+}
+
 static const std::vector<NArrow::NTest::TTestColumn> testYdbSchema = TTestSchema::YdbSchema();
 static const std::vector<NArrow::NTest::TTestColumn> testYdbPk = TTestSchema::YdbPkSchema();
 

diff --git a/ydb/core/tx/columnshard/ut_schema/ya.make b/ydb/core/tx/columnshard/ut_schema/ya.make
@@ -18,6 +18,7 @@ PEERDIR(
     library/cpp/getopt
     library/cpp/regex/pcre
     library/cpp/svnversion
+    contrib/libs/aws-sdk-cpp/aws-cpp-sdk-core
     ydb/core/testlib/default
     ydb/core/tx/columnshard/hooks/abstract
     ydb/core/tx/columnshard/hooks/testing

diff --git a/ydb/core/tx/datashard/import_s3.cpp b/ydb/core/tx/datashard/import_s3.cpp
@@ -14,6 +14,7 @@
 #include <ydb/core/tablet/resource_broker.h>
 #include <ydb/core/wrappers/s3_wrapper.h>
 #include <ydb/core/wrappers/s3_storage.h>
+#include <ydb/core/wrappers/s3_storage_config.h>
 #include <ydb/core/io_formats/ydb_dump/csv_ydb_dump.h>
 #include <ydb/public/lib/scheme_types/scheme_type_id.h>
 

diff --git a/ydb/core/tx/schemeshard/ut_backup/ut_backup.cpp b/ydb/core/tx/schemeshard/ut_backup/ut_backup.cpp
@@ -7,9 +7,27 @@
 #include <util/string/cast.h>
 #include <util/string/printf.h>
 
+#include <library/cpp/testing/hook/hook.h>
+
+#include <aws/core/Aws.h>
+
 using namespace NSchemeShardUT_Private;
 using namespace NKikimr::NWrappers::NTestHelpers;
 
+namespace {
+
+Aws::SDKOptions Options;
+
+Y_TEST_HOOK_BEFORE_RUN(InitAwsAPI) {
+    Aws::InitAPI(Options);
+}
+
+Y_TEST_HOOK_AFTER_RUN(ShutdownAwsAPI) {
+    Aws::ShutdownAPI(Options);
+}
+
+}
+
 Y_UNIT_TEST_SUITE(TBackupTests) {
     using TFillFn = std::function<void(TTestBasicRuntime&)>;
 

diff --git a/ydb/core/tx/schemeshard/ut_backup/ya.make b/ydb/core/tx/schemeshard/ut_backup/ya.make
@@ -20,6 +20,7 @@ IF (NOT OS_WINDOWS)
         library/cpp/getopt
         library/cpp/regex/pcre
         library/cpp/svnversion
+        contrib/libs/aws-sdk-cpp/aws-cpp-sdk-core
         ydb/core/testlib/default
         ydb/core/tx
         ydb/core/tx/schemeshard/ut_helpers

diff --git a/ydb/core/tx/schemeshard/ut_export/ut_export.cpp b/ydb/core/tx/schemeshard/ut_export/ut_export.cpp
@@ -12,13 +12,27 @@
 #include <util/string/printf.h>
 #include <util/system/env.h>
 
+#include <library/cpp/testing/hook/hook.h>
+
+#include <aws/core/Aws.h>
+
 using namespace NSchemeShardUT_Private;
 using namespace NKikimr::NWrappers::NTestHelpers;
 
 using TTablesWithAttrs = TVector<std::pair<TString, TMap<TString, TString>>>;
 
 namespace {
 
+    Aws::SDKOptions Options;
+
+    Y_TEST_HOOK_BEFORE_RUN(InitAwsAPI) {
+        Aws::InitAPI(Options);
+    }
+
+    Y_TEST_HOOK_AFTER_RUN(ShutdownAwsAPI) {
+        Aws::ShutdownAPI(Options);
+    }
+
     void Run(TTestBasicRuntime& runtime, TTestEnv& env, const std::variant<TVector<TString>, TTablesWithAttrs>& tablesVar, const TString& request,
             Ydb::StatusIds::StatusCode expectedStatus = Ydb::StatusIds::SUCCESS,
             const TString& dbName = "/MyRoot", bool serverless = false, const TString& userSID = "") {
@@ -1809,7 +1823,7 @@ partitioning_settings {
             return ev->Get<TEvSchemeShard::TEvModifySchemeTransaction>()->Record
                 .GetTransaction(0).GetOperationType() == NKikimrSchemeOp::ESchemeOpBackup;
         };
-        
+
         THolder<IEventHandle> delayed;
         auto prevObserver = runtime.SetObserverFunc([&](TAutoPtr<IEventHandle>& ev) {
             if (delayFunc(ev)) {
@@ -1976,7 +1990,7 @@ partitioning_settings {
             min_partitions_count: 10
         )"));
     }
-    
+
     Y_UNIT_TEST(UserSID) {
         TTestBasicRuntime runtime;
         TTestEnv env(runtime);

diff --git a/ydb/core/tx/schemeshard/ut_export/ya.make b/ydb/core/tx/schemeshard/ut_export/ya.make
@@ -20,6 +20,7 @@ IF (NOT OS_WINDOWS)
         library/cpp/getopt
         library/cpp/regex/pcre
         library/cpp/svnversion
+        contrib/libs/aws-sdk-cpp/aws-cpp-sdk-core
         ydb/core/testlib/default
         ydb/core/tx
         ydb/core/tx/schemeshard/ut_helpers

diff --git a/ydb/core/tx/schemeshard/ut_export_reboots_s3/ut_export_reboots_s3.cpp b/ydb/core/tx/schemeshard/ut_export_reboots_s3/ut_export_reboots_s3.cpp
@@ -4,10 +4,28 @@
 
 #include <util/string/printf.h>
 
+#include <library/cpp/testing/hook/hook.h>
+
+#include <aws/core/Aws.h>
+
 using namespace NSchemeShardUT_Private;
 using namespace NSchemeShardUT_Private::NExportReboots;
 using namespace NKikimr::NWrappers::NTestHelpers;
 
+namespace {
+
+Aws::SDKOptions Options;
+
+Y_TEST_HOOK_BEFORE_RUN(InitAwsAPI) {
+    Aws::InitAPI(Options);
+}
+
+Y_TEST_HOOK_AFTER_RUN(ShutdownAwsAPI) {
+    Aws::ShutdownAPI(Options);
+}
+
+}
+
 Y_UNIT_TEST_SUITE(TExportToS3WithRebootsTests) {
     using TUnderlying = std::function<void(const TVector<TString>&, const TString&, TTestWithReboots&)>;
 

diff --git a/ydb/core/tx/schemeshard/ut_export_reboots_s3/ya.make b/ydb/core/tx/schemeshard/ut_export_reboots_s3/ya.make
@@ -19,6 +19,7 @@ PEERDIR(
     library/cpp/getopt
     library/cpp/regex/pcre
     library/cpp/svnversion
+    contrib/libs/aws-sdk-cpp/aws-cpp-sdk-core
     ydb/core/testlib/default
     ydb/core/tx
     ydb/core/tx/schemeshard/ut_helpers

diff --git a/ydb/core/tx/schemeshard/ut_restore/ut_restore.cpp b/ydb/core/tx/schemeshard/ut_restore/ut_restore.cpp
@@ -20,8 +20,10 @@
 
 #include <ydb/public/api/protos/ydb_import.pb.h>
 
+#include <aws/core/Aws.h>
 #include <contrib/libs/zstd/include/zstd.h>
 #include <library/cpp/string_utils/quote/quote.h>
+#include <library/cpp/testing/hook/hook.h>
 
 #include <util/datetime/base.h>
 #include <util/generic/size_literals.h>
@@ -37,6 +39,16 @@ using namespace NKikimr::NWrappers::NTestHelpers;
 
 namespace {
 
+    Aws::SDKOptions Options;
+
+    Y_TEST_HOOK_BEFORE_RUN(InitAwsAPI) {
+        Aws::InitAPI(Options);
+    }
+
+    Y_TEST_HOOK_AFTER_RUN(ShutdownAwsAPI) {
+        Aws::ShutdownAPI(Options);
+    }
+
     const TString EmptyYsonStr = R"([[[[];%false]]])";
 
     TString GenerateScheme(const NKikimrSchemeOp::TPathDescription& pathDesc) {
@@ -316,7 +328,6 @@ namespace {
         runtime.SetObserverFunc(prevObserver);
     }
 
-
 } // anonymous
 
 Y_UNIT_TEST_SUITE(TRestoreTests) {

diff --git a/ydb/core/tx/schemeshard/ut_restore/ya.make b/ydb/core/tx/schemeshard/ut_restore/ya.make
@@ -14,6 +14,7 @@ ELSE()
 ENDIF()
 
 PEERDIR(
+    contrib/libs/aws-sdk-cpp/aws-cpp-sdk-core
     contrib/libs/double-conversion
     library/cpp/string_utils/quote
     ydb/core/kqp/ut/common

diff --git a/ydb/core/wrappers/s3_storage.h b/ydb/core/wrappers/s3_storage.h
@@ -3,7 +3,6 @@
 #ifndef KIKIMR_DISABLE_S3_OPS
 
 #include "abstract.h"
-#include "s3_storage_config.h"
 
 #include <ydb/core/protos/flat_scheme_op.pb.h>
 #include <ydb/core/wrappers/events/common.h>
@@ -32,7 +31,7 @@
 
 namespace NKikimr::NWrappers::NExternalStorage {
 
-class TS3ExternalStorage: public IExternalStorageOperator, TS3User {
+class TS3ExternalStorage: public IExternalStorageOperator {
 private:
     THolder<Aws::S3::S3Client> Client;
     const Aws::Client::ClientConfiguration Config;