[yugabyte#19137][yugabyte#20022] YSQL, ASH: ASH Infrastructure and instrumentation for PGGate

Summary:
Postgres has the infrastructure to start and end wait events.
This diff adds a callback to use a postgres function to start and end
wait events in pggate layer

A new class `Flusher` is introduced which flushes the bufferable operations and
returns a `FlushFuture` object which contains the `PerformFuture` object and a
struct `MetricContext`. The `MetricContext` struct can update the wait event
and the duration of the flush request while waiting for the future object.

This diff also instruments the reads and flushes in the pggate layer
using RAII resource management.

This diff also includes some other minor changes.
Jira: DB-7935, DB-8989

Test Plan: Jenkins

Reviewers: jason, dmitry

Reviewed By: dmitry

Subscribers: hbhanawat, amitanand, yql

Differential Revision: https://phorge.dev.yugabyte.com/D30454

Author: abhinab-yb

src/postgres/src/backend/utils/misc/pg_yb_utils.c

@@ -863,6 +863,7 @@ YBInitPostgresBackend(
 		callbacks.UnixEpochToPostgresEpoch = &YbUnixEpochToPostgresEpoch;
 		callbacks.ConstructArrayDatum = &YbConstructArrayDatum;
 		callbacks.CheckUserMap = &check_usermap;
+		callbacks.PgstatReportWaitStart = &yb_pgstat_report_wait_start;
 		YBCInitPgGate(type_table, count, callbacks, session_id, &MyProc->yb_ash_metadata,
 					  &MyProc->yb_is_ash_metadata_set);
 		YBCInstallTxnDdlHook();

src/postgres/src/include/pgstat.h

@@ -1433,6 +1433,34 @@ pgstat_report_wait_end(void)
 	return pgstat_report_wait_end_for_proc(MyProc);
 }
 
+/* ----------
+ * yb_pgstat_report_wait_start() -
+ *
+ *	Called to get the current wait event info and set a new wait
+ *  event info.
+ *
+ * NB: this *must* be able to survive being called before MyProc has been
+ * initialized.
+ * ----------
+ */
+static inline uint32
+yb_pgstat_report_wait_start(uint32 wait_event_info)
+{
+	uint32 prev_wait_event_info = 0;
+	volatile PGPROC *proc = MyProc;
+
+	if (pgstat_track_activities && proc)
+	{
+		/*
+		 * Since this is a four-byte field which is always read and written as
+		 * four-bytes, updates are atomic.
+		 */
+		prev_wait_event_info = proc->wait_event_info;
+		proc->wait_event_info = wait_event_info;
+	}
+	return prev_wait_event_info;
+}
+
 /* nontransactional event counts are simple enough to inline */
 
 #define pgstat_count_heap_scan(rel)									\

src/yb/ash/wait_state.h

@@ -105,6 +105,14 @@ YB_DEFINE_TYPED_ENUM(WaitStateCode, uint32_t,
     // Don't change the value of kUnused
     ((kUnused, 0xFFFFFFFFU))
 
+    // Wait states related to postgres
+    // Don't change the position of kPostgresReserved
+    ((kPostgresReserved, YB_ASH_MAKE_EVENT(TServerWait)))
+    (kCatalogRead)
+    (kIndexRead)
+    (kStorageRead)
+    (kStorageFlush)
+
     // Common wait states
     ((kOnCpu_Active, YB_ASH_MAKE_EVENT(Common)))
     (kOnCpu_Passive)

src/yb/yql/pggate/pg_doc_op.cc

@@ -17,6 +17,8 @@
 #include <algorithm>
 #include <utility>
 
+#include "yb/ash/wait_state.h"
+
 #include "yb/common/pg_system_attr.h"
 #include "yb/common/row_mark.h"
 
@@ -36,10 +38,7 @@
 #include "yb/yql/pggate/util/pg_doc_data.h"
 #include "yb/yql/pggate/util/ybc_util.h"
 
-using std::string;
-
-namespace yb {
-namespace pggate {
+namespace yb::pggate {
 namespace {
 
 struct PgDocReadOpCachedHelper {
@@ -128,7 +127,7 @@ auto BuildRowOrders(const LWPgsqlResponsePB& response,
 // Helper function to determine the type of relation that the given Pgsql operation is being
 // performed on. This function classifies the operation into one of three buckets: system catalog,
 // secondary index or user table requests.
-TableType ResolveRelationType(const PgsqlOp& op, const PgTable& table) {
+[[nodiscard]] TableType ResolveRelationType(const PgsqlOp& op, const PgTable& table) {
   if (table->schema().table_properties().is_ysql_catalog_table()) {
     // We don't distinguish between table reads and index reads for a catalog table.
     return TableType::SYSTEM;
@@ -150,6 +149,35 @@ TableType ResolveRelationType(const PgsqlOp& op, const PgTable& table) {
   return TableType::USER;
 }
 
+[[nodiscard]] inline ash::WaitStateCode ResolveWaitEventCode(TableType table_type) {
+  switch (table_type) {
+    case TableType::SYSTEM: return ash::WaitStateCode::kCatalogRead;
+    case TableType::INDEX: return ash::WaitStateCode::kIndexRead;
+    case TableType::USER: return ash::WaitStateCode::kStorageRead;
+  }
+  FATAL_INVALID_ENUM_VALUE(TableType, table_type);
+}
+
+Result<PgDocResponse::Data> FetchResponseData(
+    PgDocResponse* response, PgSession* session, TableType table_type, bool is_write) {
+  if (is_write) {
+    return response->Get();
+  }
+  // Update session stats instrumentation only for read requests. Reads are executed
+  // synchronously with respect to Postgres query execution, and thus it is possible to
+  // correlate wait/execution times directly with the request. We update instrumentation for
+  // reads exactly once, upon receiving a success response from the underlying storage layer.
+  auto& metrics = session->metrics();
+  uint64_t wait_time = 0;
+  const auto result = VERIFY_RESULT(metrics.CallWithDuration(
+      [response,
+       event_watcher = session->StartWaitEvent(ResolveWaitEventCode(table_type))] {
+        return response->Get(); },
+      &wait_time));
+  metrics.ReadRequest(table_type, wait_time);
+  return result;
+}
+
 } // namespace
 
 PgDocResult::PgDocResult(rpc::SidecarHolder data, std::vector<int64_t>&& row_orders)
@@ -283,18 +311,8 @@ Result<std::list<PgDocResult>> PgDocOp::GetResult() {
       RETURN_NOT_OK(SendRequest());
     }
 
-    uint64_t wait_time = 0;
-    auto result_data = VERIFY_RESULT(pg_session_->metrics().CallWithDuration(
-        [&response = response_] { return response.Get(); }, &wait_time));
-
-    // Update session stats instrumentation only for read requests. Reads are executed
-    // synchronously with respect to Postgres query execution, and thus it is possible to
-    // correlate wait/execution times directly with the request. We update instrumentation for
-    // reads exactly once, upon receiving a success response from the underlying storage layer.
-    if (!IsWrite()) {
-      pg_session_->metrics().ReadRequest(
-          ResolveRelationType(*pgsql_ops_.front(), table_), wait_time);
-    }
+    const auto result_data = VERIFY_RESULT(FetchResponseData(
+        &response_, &*pg_session_, ResolveRelationType(*pgsql_ops_.front(), table_), IsWrite()));
 
     result = VERIFY_RESULT(ProcessResponse(result_data));
 
@@ -343,7 +361,7 @@ Status PgDocOp::SendRequestImpl(ForceNonBufferable force_non_bufferable) {
 
   // Update session stats instrumentation for write requests only. Writes are buffered and flushed
   // asynchronously, and thus it is not possible to correlate wait/execution times directly with
-  // the request. We update instrumentation for writes sexactly once, after successfully sending an
+  // the request. We update instrumentation for writes exactly once, after successfully sending an
   // RPC request to the underlying storage layer.
   if (IsWrite()) {
     pg_session_->metrics().WriteRequest(ResolveRelationType(*pgsql_ops_.front(), table_));
@@ -1175,7 +1193,7 @@ Result<bool> PgDocReadOp::SetScanPartitionBoundary() {
       partition_key != partition_keys.end(), InvalidArgument, "invalid partition key given");
 
   // Seek upper bound (Beginning of next tablet).
-  string upper_bound;
+  std::string upper_bound;
   const auto& next_partition_key = std::next(partition_key, 1);
   if (next_partition_key != partition_keys.end()) {
     upper_bound = *next_partition_key;
@@ -1493,5 +1511,4 @@ PgDocOp::SharedPtr MakeDocReadOpWithData(
   return std::make_shared<PgDocReadOpCached>(pg_session, std::move(data));
 }
 
-}  // namespace pggate
-}  // namespace yb
+}  // namespace yb::pggate

src/yb/yql/pggate/pg_operation_buffer.cc

@@ -42,8 +42,7 @@
 #include "yb/yql/pggate/pg_op.h"
 #include "yb/yql/pggate/pg_tabledesc.h"
 
-namespace yb {
-namespace pggate {
+namespace yb::pggate {
 
 namespace {
 
@@ -142,11 +141,60 @@ inline bool IsTableUsedByRequest(const LWPgsqlReadRequestPB& request, const Slic
 
 using RowKeys = std::unordered_set<RowIdentifier, boost::hash<RowIdentifier>>;
 
+struct MetricContext {
+  PgDocMetrics& metrics;
+  PgWaitEventWatcher::Starter wait_starter;
+};
+
+class FlushFuture {
+ public:
+  FlushFuture(PerformFuture&& future, MetricContext* context)
+      : future_(std::move(future)), context_(context) {}
+
+  bool Ready() const {
+    return future_.Ready();
+  }
+
+  Status EnsureCompleted() {
+    uint64_t duration = 0;
+    auto& metrics = context_->metrics;
+    {
+      PgWaitEventWatcher watcher(context_->wait_starter,
+                                 ash::WaitStateCode::kStorageFlush);
+      RETURN_NOT_OK(metrics.CallWithDuration(
+          [&future = future_] { return future.Get(); }, &duration));
+    }
+    metrics.FlushRequest(duration);
+    return Status::OK();
+  }
+
+ private:
+  PerformFuture future_;
+  MetricContext* context_;
+};
+
+class Flusher {
+ public:
+  Flusher(
+      PgOperationBuffer::OperationsFlusher&& ops_flusher, PgDocMetrics* metrics,
+      PgWaitEventWatcher::Starter wait_starter)
+      : ops_flusher_(std::move(ops_flusher)),
+        context_{*metrics, wait_starter} {}
+
+  Result<FlushFuture> Flush(BufferableOperations&& ops, bool transactional) {
+    return FlushFuture{VERIFY_RESULT(ops_flusher_(std::move(ops), transactional)), &context_};
+  }
+
+ private:
+  PgOperationBuffer::OperationsFlusher ops_flusher_;
+  MetricContext context_;
+};
+
 struct InFlightOperation {
   RowKeys keys;
-  PerformFuture future;
+  FlushFuture future;
 
-  explicit InFlightOperation(PerformFuture future_)
+  explicit InFlightOperation(FlushFuture&& future_)
       : future(std::move(future_)) {}
 };
 
@@ -199,12 +247,10 @@ size_t BufferableOperations::size() const {
 class PgOperationBuffer::Impl {
  public:
   Impl(
-    const Flusher& flusher,
-    const BufferingSettings& buffering_settings,
-    PgDocMetrics* metrics)
-    : flusher_(flusher),
-      buffering_settings_(buffering_settings),
-      metrics_(*metrics) {}
+      OperationsFlusher&& ops_flusher, PgDocMetrics* metrics,
+      PgWaitEventWatcher::Starter wait_starter, const BufferingSettings& buffering_settings)
+      : flusher_(std::move(ops_flusher), metrics, wait_starter),
+        buffering_settings_(buffering_settings) {}
 
   Status Add(const PgTableDesc& table, PgsqlWriteOpPtr op, bool transactional) {
     return ClearOnError(DoAdd(table, std::move(op), transactional));
@@ -317,10 +363,7 @@ class PgOperationBuffer::Impl {
 
   Status EnsureCompleted(size_t count) {
     for (; count && !in_flight_ops_.empty(); --count) {
-      uint64_t duration = 0;
-      auto result = VERIFY_RESULT(metrics_.CallWithDuration(
-          [&future = in_flight_ops_.front().future] { return future.Get(); }, &duration));
-      metrics_.FlushRequest(duration);
+      RETURN_NOT_OK(in_flight_ops_.front().future.EnsureCompleted());
       in_flight_ops_.pop_front();
     }
     return Status::OK();
@@ -378,7 +421,7 @@ class PgOperationBuffer::Impl {
         RETURN_NOT_OK(EnsureCompleted(1));
       }
       in_flight_ops_.push_back(
-        InFlightOperation(VERIFY_RESULT(flusher_(std::move(ops), transactional))));
+        InFlightOperation(VERIFY_RESULT(flusher_.Flush(std::move(ops), transactional))));
       return true;
     }
     return false;
@@ -401,20 +444,18 @@ class PgOperationBuffer::Impl {
     return false;
   }
 
-  const Flusher flusher_;
+  Flusher flusher_;
   const BufferingSettings& buffering_settings_;
   BufferableOperations ops_;
   BufferableOperations txn_ops_;
   RowKeys keys_;
   InFlightOps in_flight_ops_;
-  PgDocMetrics& metrics_;
 };
 
-PgOperationBuffer::PgOperationBuffer(const Flusher& flusher,
-                                     const BufferingSettings& buffering_settings,
-                                     PgDocMetrics* metrics)
-    : impl_(new Impl(flusher, buffering_settings, metrics)) {
-}
+PgOperationBuffer::PgOperationBuffer(
+    OperationsFlusher&& ops_flusher, PgDocMetrics* metrics,
+    PgWaitEventWatcher::Starter wait_starter, const BufferingSettings& buffering_settings)
+    : impl_(new Impl(std::move(ops_flusher), metrics, wait_starter, buffering_settings)) {}
 
 PgOperationBuffer::~PgOperationBuffer() = default;
 
@@ -439,5 +480,4 @@ void PgOperationBuffer::Clear() {
     impl_->Clear();
 }
 
-} // namespace pggate
-} // namespace yb
+} // namespace yb::pggate

src/yb/yql/pggate/pg_operation_buffer.h

@@ -16,6 +16,8 @@
 #include <functional>
 #include <memory>
 
+#include "yb/ash/wait_state.h"
+
 #include "yb/common/common_fwd.h"
 #include "yb/common/pg_types.h"
 
@@ -24,9 +26,9 @@
 
 #include "yb/yql/pggate/pg_gate_fwd.h"
 #include "yb/yql/pggate/pg_perform_future.h"
+#include "yb/yql/pggate/pg_tools.h"
 
-namespace yb {
-namespace pggate {
+namespace yb::pggate {
 
 class PgDocMetrics;
 
@@ -49,12 +51,11 @@ struct BufferableOperations {
 
 class PgOperationBuffer {
  public:
-  using Flusher = std::function<Result<PerformFuture>(BufferableOperations, bool)>;
+  using OperationsFlusher = std::function<Result<PerformFuture>(BufferableOperations&&, bool)>;
 
   PgOperationBuffer(
-    const Flusher& flusher,
-    const BufferingSettings& buffering_settings,
-    PgDocMetrics* metrics);
+      OperationsFlusher&& ops_flusher, PgDocMetrics* metrics,
+      PgWaitEventWatcher::Starter wait_starter, const BufferingSettings& buffering_settings);
   ~PgOperationBuffer();
   Status Add(const PgTableDesc& table, PgsqlWriteOpPtr op, bool transactional);
   Status Flush();
@@ -68,5 +69,4 @@ class PgOperationBuffer {
   std::unique_ptr<Impl> impl_;
 };
 
-} // namespace pggate
-} // namespace yb
+} // namespace yb::pggate