feat userver: set trx_tracker start/end in transactions

Awaiting long requests during an active transaction (e.g. PostgreSQL, YDB, Clickhouse) is a frequently occuring bug that can easily use up the DB connection pool.

This PR adds calls of StartTransaction and EndTransaction in Begin/Commit/Rollback methods of transactional databases (MySQL, PostgreSQL, SQLite, YDB).
commit_hash:435e51a1fd2734b964359d980c171b02503409ff

Author: p34k1n

.mapping.json

@@ -1970,8 +1970,6 @@
   "core/src/utils/text.cpp":"taxi/uservices/userver/core/src/utils/text.cpp",
   "core/src/utils/text_test.cpp":"taxi/uservices/userver/core/src/utils/text_test.cpp",
   "core/src/utils/trx_tracker.cpp":"taxi/uservices/userver/core/src/utils/trx_tracker.cpp",
-  "core/src/utils/trx_tracker_internal.cpp":"taxi/uservices/userver/core/src/utils/trx_tracker_internal.cpp",
-  "core/src/utils/trx_tracker_internal.hpp":"taxi/uservices/userver/core/src/utils/trx_tracker_internal.hpp",
   "core/src/utils/trx_tracker_test.cpp":"taxi/uservices/userver/core/src/utils/trx_tracker_test.cpp",
   "core/src/utils/userver_info.cpp":"taxi/uservices/userver/core/src/utils/userver_info.cpp",
   "core/static_configs/dns_client.yaml":"taxi/uservices/userver/core/static_configs/dns_client.yaml",

core/include/userver/utils/trx_tracker.hpp

@@ -3,6 +3,9 @@
 /// @file
 /// @brief Tracking for heavy operations while having active transactions.
 
+#include <optional>
+#include <thread>
+
 #include <userver/utils/impl/source_location.hpp>
 #include <userver/utils/statistics/rate.hpp>
 
@@ -22,13 +25,62 @@ USERVER_NAMESPACE_BEGIN
 /// @snippet utils/trx_tracker_test.cpp  Sample TransactionTracker usage
 namespace utils::trx_tracker {
 
-/// @brief Increment transaction counter.
-void StartTransaction();
+namespace impl {
+
+/// @brief Global enabler for transaction tracker.
+class GlobalEnabler final {
+public:
+    explicit GlobalEnabler(bool enable = true);
+    ~GlobalEnabler();
+
+    GlobalEnabler(const GlobalEnabler&) = delete;
+    GlobalEnabler& operator=(const GlobalEnabler&) = delete;
+};
+
+/// @brief Check if transaction tracker is enabled.
+bool IsEnabled() noexcept;
 
-/// @brief Decrement transaction counter.
+/// @brief Unique ID for every task.
 ///
-/// If called before StartTransaction, the behavior is undefined.
-void EndTransaction() noexcept;
+/// Sometimes transactions start and end in different coroutines.
+/// To prevent transaction from incrementing and decrementing different
+/// transaction counters, TransactionLock stores TaskId on Lock and
+/// checks that stored TaskId is the same as current TaskId in Unlock.
+class TaskId final {
+public:
+    TaskId();
+
+    bool operator==(const TaskId& other) const;
+    bool operator!=(const TaskId& other) const;
+
+private:
+    std::thread::id created_thread_id_;
+    std::uint64_t thread_local_counter_;
+};
+
+}  // namespace impl
+
+/// @brief Class for incrementing and decrementing transaction counter.
+class TransactionLock final {
+public:
+    TransactionLock() = default;
+    TransactionLock(const TransactionLock&) = delete;
+    TransactionLock(TransactionLock&&) noexcept;
+    TransactionLock operator=(const TransactionLock&) = delete;
+    TransactionLock& operator=(TransactionLock&&) noexcept;
+
+    /// @brief Decrement transaction counter on destruction.
+    ~TransactionLock();
+
+    /// @brief Manually increment transaction counter.
+    void Lock() noexcept;
+
+    /// @brief Manually decrement transaction counter.
+    void Unlock() noexcept;
+
+private:
+    std::optional<impl::TaskId> task_id_;
+};
 
 /// @brief Check for active transactions.
 void CheckNoTransactions(utils::impl::SourceLocation location = utils::impl::SourceLocation::Current());
@@ -45,7 +97,7 @@ void CheckNoTransactions(std::string_view location);
 class CheckDisabler final {
 public:
     /// @brief Disable check for active transactions.
-    CheckDisabler();
+    explicit CheckDisabler();
 
     /// @brief Reenable check for active transactions on destruction.
     ~CheckDisabler();

core/src/components/run.cpp

@@ -27,6 +27,7 @@
 #include <userver/utils/strerror.hpp>
 #include <userver/utils/string_literal.hpp>
 #include <userver/utils/traceful_exception.hpp>
+#include <userver/utils/trx_tracker.hpp>
 
 #include <components/manager.hpp>
 #include <components/manager_config.hpp>
@@ -36,7 +37,6 @@
 #include <utils/ignore_signal_scope.hpp>
 #include <utils/jemalloc.hpp>
 #include <utils/signal_catcher.hpp>
-#include <utils/trx_tracker_internal.hpp>
 
 USERVER_NAMESPACE_BEGIN
 
@@ -203,7 +203,7 @@ void DoRun(
     utils::impl::UserverExperimentsScope experiments_scope;
     std::optional<impl::Manager> manager;
 
-    const utils::trx_tracker::GlobalEnabler enabler{manager_config.enable_trx_tracker};
+    const utils::trx_tracker::impl::GlobalEnabler enabler{manager_config.enable_trx_tracker};
 
     try {
         experiments_scope.EnableOnly(manager_config.enabled_experiments);

core/src/utils/trx_tracker.cpp

@@ -1,7 +1,6 @@
 #include <userver/utils/trx_tracker.hpp>
 
-#include "trx_tracker_internal.hpp"
-
+#include <userver/compiler/thread_local.hpp>
 #include <userver/engine/task/local_variable.hpp>
 #include <userver/formats/json.hpp>
 #include <userver/logging/log.hpp>
@@ -16,9 +15,12 @@ namespace utils::trx_tracker {
 
 namespace {
 
+compiler::ThreadLocal local_task_counter = []() -> std::uint64_t { return 0; };
+
 struct TransactionTracker final {
     std::uint32_t trx_count = 0;
     std::uint32_t disabler_count = 0;
+    impl::TaskId task_id;
 };
 
 struct TransactionTrackerStatisticsInternal final {
@@ -29,12 +31,12 @@ engine::TaskLocalVariable<TransactionTracker> transaction_tracker{};
 TransactionTrackerStatisticsInternal transaction_tracker_statistics{};
 
 void CheckNoTransactions(utils::function_ref<std::string()> get_location) {
-    if (!IsEnabled()) {
+    if (!impl::IsEnabled()) {
         return;
     }
 
     auto* tracker = transaction_tracker.GetOptional();
-    if (tracker && !tracker->disabler_count && tracker->trx_count) {
+    if (tracker && !tracker->disabler_count && tracker->trx_count > 0) {
         const std::string location = get_location();
         logging::LogExtra log_extra;
         log_extra.Extend("location", location);
@@ -49,19 +51,84 @@ void CheckNoTransactions(utils::function_ref<std::string()> get_location) {
     }
 }
 
-}  // namespace
-
-void StartTransaction() {
-    if (IsEnabled()) {
+std::optional<impl::TaskId> StartTransaction() {
+    if (impl::IsEnabled()) {
         ++transaction_tracker->trx_count;
+        return transaction_tracker->task_id;
     }
+    return std::nullopt;
 }
 
-void EndTransaction() noexcept {
-    if (IsEnabled()) {
-        UASSERT(transaction_tracker.GetOptional() != nullptr);
+bool EndTransaction(const impl::TaskId& task_id) noexcept {
+    if (impl::IsEnabled()) {
+        if (transaction_tracker.GetOptional() == nullptr || transaction_tracker->task_id != task_id) {
+            // EndTransaction is called in different task than StartTransaction
+            return false;
+        }
         --transaction_tracker->trx_count;
+        return true;
     }
+    return true;
+}
+
+bool trx_tracker_enabled = false;
+bool trx_tracker_enabler_exists = false;
+
+}  // namespace
+
+namespace impl {
+
+GlobalEnabler::GlobalEnabler(bool enable) {
+    UASSERT(!trx_tracker_enabler_exists);
+    trx_tracker_enabled = enable;
+    trx_tracker_enabler_exists = true;
+}
+
+GlobalEnabler::~GlobalEnabler() {
+    trx_tracker_enabled = false;
+    trx_tracker_enabler_exists = false;
+}
+
+bool IsEnabled() noexcept { return trx_tracker_enabled; }
+
+TaskId::TaskId() : created_thread_id_{std::this_thread::get_id()} {
+    auto counter_scope{local_task_counter.Use()};
+    thread_local_counter_ = (*counter_scope)++;
+}
+
+bool TaskId::operator==(const TaskId& other) const {
+    return created_thread_id_ == other.created_thread_id_ && thread_local_counter_ == other.thread_local_counter_;
+}
+
+bool TaskId::operator!=(const TaskId& other) const { return !(*this == other); }
+
+}  // namespace impl
+
+TransactionLock::~TransactionLock() { Unlock(); }
+
+void TransactionLock::Lock() noexcept {
+    if (!task_id_.has_value()) {
+        task_id_ = StartTransaction();
+    }
+}
+
+void TransactionLock::Unlock() noexcept {
+    if (task_id_.has_value() && EndTransaction(task_id_.value())) {
+        task_id_ = std::nullopt;
+    }
+}
+
+TransactionLock::TransactionLock(TransactionLock&& other) noexcept
+    : task_id_(std::exchange(other.task_id_, std::nullopt)) {}
+
+TransactionLock& TransactionLock::operator=(TransactionLock&& other) noexcept {
+    if (&other == this) {
+        return *this;
+    }
+
+    Unlock();
+    task_id_ = std::exchange(other.task_id_, std::nullopt);
+    return *this;
 }
 
 void CheckNoTransactions(utils::impl::SourceLocation location) {