From f5c4d988a87237625724498a5ad3f09579792d2f Mon Sep 17 00:00:00 2001
From: Tristan Partin <tristan@neon.tech>
Date: Mon, 25 Nov 2024 10:46:06 -0600
Subject: [PATCH] Add compute_logical_snapshot_files_bytes metric

This metric exposes the size of all non-temporary logical snapshot
files.

Signed-off-by: Tristan Partin <tristan@neon.tech>
---
 compute/etc/neon_collector.jsonnet              |  1 +
 .../compute_logical_snapshots_bytes.16.sql      |  7 +++++++
 .../compute_logical_snapshots_bytes.libsonnet   | 17 +++++++++++++++++
 .../compute_logical_snapshots_bytes.sql         |  9 +++++++++
 4 files changed, 34 insertions(+)
 create mode 100644 compute/etc/sql_exporter/compute_logical_snapshots_bytes.16.sql
 create mode 100644 compute/etc/sql_exporter/compute_logical_snapshots_bytes.libsonnet
 create mode 100644 compute/etc/sql_exporter/compute_logical_snapshots_bytes.sql

diff --git a/compute/etc/neon_collector.jsonnet b/compute/etc/neon_collector.jsonnet
index 75d69c7b6888b..148089f1349f8 100644
--- a/compute/etc/neon_collector.jsonnet
+++ b/compute/etc/neon_collector.jsonnet
@@ -6,6 +6,7 @@
     import 'sql_exporter/compute_backpressure_throttling_seconds.libsonnet',
     import 'sql_exporter/compute_current_lsn.libsonnet',
     import 'sql_exporter/compute_logical_snapshot_files.libsonnet',
+    import 'sql_exporter/compute_logical_snapshot_files_bytes.libsonnet',
     import 'sql_exporter/compute_max_connections.libsonnet',
     import 'sql_exporter/compute_receive_lsn.libsonnet',
     import 'sql_exporter/compute_subscriptions_count.libsonnet',
diff --git a/compute/etc/sql_exporter/compute_logical_snapshots_bytes.16.sql b/compute/etc/sql_exporter/compute_logical_snapshots_bytes.16.sql
new file mode 100644
index 0000000000000..73a9c114053d2
--- /dev/null
+++ b/compute/etc/sql_exporter/compute_logical_snapshots_bytes.16.sql
@@ -0,0 +1,7 @@
+SELECT
+  (SELECT current_setting('neon.timeline_id')) AS timeline_id,
+  -- Postgres creates temporary snapshot files of the form %X-%X.snap.%d.tmp.
+  -- These temporary snapshot files are renamed to the actual snapshot files
+  -- after they are completely built. We only WAL-log the completely built
+  -- snapshot files
+  (SELECT COALESCE(sum(size), 0) FROM pg_ls_logicalsnapdir() WHERE name LIKE '%.snap') AS logical_snapshots_bytes;
diff --git a/compute/etc/sql_exporter/compute_logical_snapshots_bytes.libsonnet b/compute/etc/sql_exporter/compute_logical_snapshots_bytes.libsonnet
new file mode 100644
index 0000000000000..cd10f599379b0
--- /dev/null
+++ b/compute/etc/sql_exporter/compute_logical_snapshots_bytes.libsonnet
@@ -0,0 +1,17 @@
+local neon = import 'neon.libsonnet';
+
+local pg_ls_logicalsnapdir = importstr 'sql_exporter/compute_logical_snapshots_bytes.16.sql';
+local pg_ls_dir = importstr 'sql_exporter/compute_logical_snapshots_bytes.sql';
+
+{
+  metric_name: 'compute_logical_snapshots_bytes',
+  type: 'gauge',
+  help: 'Size of the pg_logical/snapshots directory, not including temporary files',
+  key_labels: [
+    'timeline_id',
+  ],
+  values: [
+    'logical_snapshots_bytes',
+  ],
+  query: if neon.PG_MAJORVERSION_NUM < 16 then pg_ls_dir else pg_ls_logicalsnapdir,
+}
diff --git a/compute/etc/sql_exporter/compute_logical_snapshots_bytes.sql b/compute/etc/sql_exporter/compute_logical_snapshots_bytes.sql
new file mode 100644
index 0000000000000..9ba4eeebe8ab7
--- /dev/null
+++ b/compute/etc/sql_exporter/compute_logical_snapshots_bytes.sql
@@ -0,0 +1,9 @@
+SELECT
+  (SELECT setting FROM pg_settings WHERE name = 'neon.timeline_id') AS timeline_id,
+  -- Postgres creates temporary snapshot files of the form %X-%X.snap.%d.tmp.
+  -- These temporary snapshot files are renamed to the actual snapshot files
+  -- after they are completely built. We only WAL-log the completely built
+  -- snapshot files
+  (SELECT COALESCE(sum((pg_stat_file('pg_logical/snapshots/' || name, missing_ok => true)).size), 0)
+    FROM (SELECT pg_ls_dir('pg_logical/snapshots') AS name WHERE name LIKE '%.snap');
+  ) AS logical_snapshots_bytes;