feat userver/pg: coverage: only report sql statement once

userver-framework · Dec 27, 2023 · 150e662 · 150e662
1 parent a702862
commit 150e662
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 7 deletions.
diff --git a/core/include/userver/testsuite/testpoint.hpp b/core/include/userver/testsuite/testpoint.hpp
@@ -19,6 +19,12 @@
 
 USERVER_NAMESPACE_BEGIN
 
+namespace testsuite {
+
+bool AreTestpointsAvailable() noexcept;
+
+}
+
 namespace testsuite::impl {
 
 class TestpointScope final {

diff --git a/core/src/testsuite/testpoint.cpp b/core/src/testsuite/testpoint.cpp
@@ -136,6 +136,11 @@ void TestpointControl::SetClient(TestpointClientBase& client) {
   client_instance = &client;
 }
 
+bool AreTestpointsAvailable() noexcept {
+  if (!client_instance) return false;
+  return true;
+}
+
 }  // namespace testsuite
 
 USERVER_NAMESPACE_END
diff --git a/postgresql/src/storages/postgres/detail/connection_impl.cpp b/postgresql/src/storages/postgres/detail/connection_impl.cpp
@@ -794,14 +794,10 @@ ResultSet ConnectionImpl::ExecuteCommand(const Query& query,
       std::chrono::duration_cast<std::chrono::milliseconds>(
           deadline.TimeLeft());
   auto span = MakeQuerySpan(query, {network_timeout, GetStatementTimeout()});
-  if (query.GetName()) {
-    try {
-      TESTPOINT("sql_statement", formats::json::MakeObject(
-                                     "name", query.GetName()->GetUnderlying()));
-    } catch (const std::exception& e) {
-      LOG_WARNING() << e;
-    }
+  if (testsuite::AreTestpointsAvailable() && query.GetName()) {
+    ReportStatement(query.GetName()->GetUnderlying());
   }
+
   auto scope = span.CreateScopeTime();
   CountExecute count_execute(stats_);
 
@@ -970,6 +966,25 @@ ResultSet ConnectionImpl::WaitResult(const std::string& statement,
 
 void ConnectionImpl::Cancel() { conn_wrapper_.Cancel().Wait(); }
 
+void ConnectionImpl::ReportStatement(const std::string& name) {
+  // Only report statement usage once.
+  {
+    std::unique_lock<engine::Mutex> lock{statements_mutex_};
+    if (statements_reported_.count(name)) return;
+  }
+
+  try {
+    TESTPOINT_CALLBACK(
+        "sql_statement", formats::json::MakeObject("name", name),
+        ([&name, this](auto) {
+          std::unique_lock<engine::Mutex> lock{statements_mutex_};
+          statements_reported_.insert(name);
+        }));
+  } catch (const std::exception& e) {
+    LOG_WARNING() << e;
+  }
+}
+
 }  // namespace storages::postgres::detail
 
 USERVER_NAMESPACE_END
diff --git a/postgresql/src/storages/postgres/detail/connection_impl.hpp b/postgresql/src/storages/postgres/detail/connection_impl.hpp
@@ -5,10 +5,12 @@
 #include <string>
 #include <string_view>
 #include <unordered_map>
+#include <unordered_set>
 
 #include <userver/cache/lru_map.hpp>
 #include <userver/concurrent/background_task_storage_fwd.hpp>
 #include <userver/engine/deadline.hpp>
+#include <userver/engine/mutex.hpp>
 #include <userver/engine/semaphore.hpp>
 #include <userver/engine/task/task_processor_fwd.hpp>
 #include <userver/error_injection/settings_fwd.hpp>
@@ -178,6 +180,8 @@ class ConnectionImpl {
 
   void Cancel();
 
+  void ReportStatement(const std::string& name);
+
   const std::string uuid_;
   Connection::Statistics stats_;
   PGConnectionWrapper conn_wrapper_;
@@ -195,6 +199,9 @@ class ConnectionImpl {
   OptionalCommandControl transaction_cmd_ctl_;
   TimeoutDuration current_statement_timeout_{};
   const error_injection::Settings ei_settings_;
+
+  std::unordered_set<std::string> statements_reported_;
+  engine::Mutex statements_mutex_;
 };
 
 }  // namespace storages::postgres::detail