Harden exporter state and add follow-up TAP tests

src/export/stats_exporter.cc

@@ -335,6 +335,8 @@ int PschGetRetryDelayMs(void) {
 }
 
 int PschGetConsecutiveFailures(void) {
+  if (!g_exporter.exporter)
+    return 0;
   return g_exporter.exporter->NumConsecutiveFailures();
 }
 

src/queue/shmem.cc

@@ -145,7 +145,8 @@ static bool TryEnqueueLocked(const PschEvent* event, uint32 capacity) {
 
 static void PschShmemShutdown([[maybe_unused]] int code, [[maybe_unused]] Datum arg) {
   if (psch_shared_state != nullptr) {
-    elog(LOG, "pg_stat_ch: shutdown (enqueued=%lu, dropped=%lu, exported=%lu)",
+    elog(LOG, "pg_stat_ch: shutdown (enqueued=" UINT64_FORMAT ", dropped=" UINT64_FORMAT
+             ", exported=" UINT64_FORMAT ")",
          pg_atomic_read_u64(&psch_shared_state->enqueued),
          pg_atomic_read_u64(&psch_shared_state->dropped),
          pg_atomic_read_u64(&psch_shared_state->exported));

src/worker/bgworker.cc

@@ -38,6 +38,7 @@ extern "C" {
 
 #include "queue/shmem.h"
 
+#include <cerrno>
 #include <csignal>
 
 #include "config/guc.h"
@@ -87,6 +88,7 @@ static void HandleConfigReload() {
 
 // Callback for bgworker process exit (registered via on_proc_exit)
 static void PschBgworkerShutdown([[maybe_unused]] int code, [[maybe_unused]] Datum arg) {
+  PschSetBgworkerPid(0);
   PschExporterShutdown();
 }
 
@@ -224,7 +226,12 @@ void PschSignalFlush(void) {
   }
 
   if (kill(bgworker_pid, SIGUSR2) != 0) {
-    ereport(WARNING, (errmsg("pg_stat_ch: failed to signal background worker")));
+    if (errno == ESRCH) {
+      PschSetBgworkerPid(0);
+      ereport(WARNING, (errmsg("pg_stat_ch: background worker not running")));
+    } else {
+      ereport(WARNING, (errmsg("pg_stat_ch: failed to signal background worker: %m")));
+    }
   }
 }
 

t/009_bgworker.pl

@@ -6,8 +6,10 @@
 use lib 't';
 
 use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::BackgroundPsql;
 use PostgreSQL::Test::Utils;
 use Test::More;
+use Time::HiRes qw(sleep);
 
 use psch;
 
@@ -66,6 +68,47 @@
 });
 is($capacity, '65536', 'Queue capacity matches default GUC value');
 
+# Test 7: Flush should not signal a stale PID after the worker exits
+my $bgworker_pid = $node->safe_psql('postgres', q{
+    SELECT pid
+    FROM pg_stat_activity
+    WHERE backend_type = 'pg_stat_ch exporter'
+});
+ok($bgworker_pid =~ /^\d+$/, 'Found background worker PID');
+
+my $terminated = $node->safe_psql('postgres', qq{
+    SELECT pg_terminate_backend($bgworker_pid)
+});
+is($terminated, 't', 'Background worker terminated');
+
+my $bgworker_count_after_terminate = 1;
+for my $i (1 .. 50) {
+    $bgworker_count_after_terminate = $node->safe_psql('postgres', q{
+        SELECT count(*)
+        FROM pg_stat_activity
+        WHERE backend_type = 'pg_stat_ch exporter'
+    });
+    last if $bgworker_count_after_terminate eq '0';
+    sleep(0.1);
+}
+is($bgworker_count_after_terminate, '0', 'Background worker exited before restart');
+
+my $session = $node->background_psql('postgres', on_error_stop => 1);
+my ($stdout, $ret) = $session->query('SET client_min_messages = warning');
+is($ret, 0, 'Can lower message threshold for flush warning');
+
+$session->{stderr} = '';
+($stdout, $ret) = $session->query('SELECT pg_stat_ch_flush()');
+ok(defined $ret, 'Flush returned control while worker is down');
+like($session->{stderr} // '', qr/background worker not running/,
+    'Flush reports that the background worker is not running');
+unlike($session->{stderr} // '', qr/failed to signal background worker/,
+    'Flush does not report a stale-PID signal failure');
+
+($stdout, $ret) = $session->query('SELECT 1');
+like($stdout, qr/^1$/m, 'Session remains usable after flush warning');
+$session->quit();
+
 diag("Background worker test results:");
 diag("  Enqueued: $stats_post_reload->{enqueued}");
 diag("  Dropped:  $stats_post_reload->{dropped}");

t/031_squashed_in_list_normalization.pl

@@ -0,0 +1,168 @@
+#!/usr/bin/env perl
+# Test: Squashed IN-list normalization (PG18+ only)
+#
+# PG18 introduced "squashed" constant lists: SELECT * FROM t WHERE id IN (1,2,3)
+# gets a single LocationLen entry with squashed=true that covers the entire list.
+# The normalized output should be: ... IN ($1 /*, ... */)
+#
+# This test verifies that FillInConstantLengths correctly computes the token
+# length for squashed constants, producing clean normalized output.
+
+use strict;
+use warnings;
+use lib 't';
+
+use PostgreSQL::Test::BackgroundPsql;
+use PostgreSQL::Test::Cluster;
+use PostgreSQL::Test::Utils;
+use Test::More;
+
+use psch;
+
+# Skip if Docker/ClickHouse not available
+if (!psch_clickhouse_available()) {
+    plan skip_all => 'Docker not available';
+}
+
+my $ch_check = `curl -s 'http://localhost:18123/' --data 'SELECT 1' 2>/dev/null`;
+if ($ch_check !~ /^1/) {
+    plan skip_all => 'ClickHouse container not running';
+}
+
+psch_query_clickhouse("TRUNCATE TABLE IF EXISTS pg_stat_ch.events_raw");
+
+my $node = psch_init_node_with_clickhouse('squashed_in',
+    flush_interval_ms => 100,
+    batch_max => 100
+);
+
+# Detect PG version — squashed IN lists are PG18+ only
+my $pg_version = $node->safe_psql('postgres',
+    'SHOW server_version_num');
+chomp($pg_version);
+if ($pg_version < 180000) {
+    plan skip_all => "Squashed IN-list normalization requires PG18+ (have $pg_version)";
+}
+
+$node->safe_psql('postgres', 'CREATE TABLE test_squash(id int, val text)');
+$node->safe_psql('postgres',
+    "INSERT INTO test_squash SELECT g, 'v' || g FROM generate_series(1,20) g");
+
+# Helper: run a query, flush, and return the captured query text.
+sub get_captured_query {
+    my ($node, $sql) = @_;
+
+    psch_query_clickhouse("TRUNCATE TABLE pg_stat_ch.events_raw");
+    psch_reset_stats($node);
+
+    my $session = $node->background_psql('postgres', on_error_stop => 1);
+    my ($stdout, $ret) = $session->query('SELECT pg_backend_pid()');
+    die "failed to get backend pid" unless $ret == 0;
+    my ($pid) = $stdout =~ /(\d+)/;
+    die "backend pid missing" unless defined $pid;
+
+    ($stdout, $ret) = $session->query($sql);
+    die "query failed: $sql" unless $ret == 0;
+
+    ($stdout, $ret) = $session->query('SELECT pg_stat_ch_flush()');
+    die "flush failed" unless $ret == 0;
+    $session->quit();
+
+    psch_wait_for_export($node, 1, 10);
+
+    return psch_wait_for_clickhouse_query(
+        "SELECT query FROM pg_stat_ch.events_raw " .
+        "WHERE pid = $pid " .
+        "AND query NOT LIKE '%pg_stat_ch%' " .
+        "AND query NOT LIKE '%pg_extension%' " .
+        "AND query != '' " .
+        "ORDER BY ts_start DESC LIMIT 1",
+        sub { $_[0] ne '' },
+        10
+    );
+}
+
+# ---------------------------------------------------------------------------
+# Test 1: Basic IN list with numeric constants gets squashed
+# ---------------------------------------------------------------------------
+subtest 'IN list with numerics is squashed' => sub {
+    my $q = get_captured_query($node,
+        "SELECT * FROM test_squash WHERE id IN (1, 2, 3, 4, 5)");
+
+    like($q, qr/\$\d/, 'IN list has placeholder');
+    # The squashed form should collapse the entire list into one placeholder
+    like($q, qr{/\*, \.\.\. \*/}, 'Squashed comment marker present');
+    # The original comma-separated list should be gone
+    unlike($q, qr/\d\s*,\s*\d/, 'No comma-separated raw numbers');
+    like($q, qr/IN\s*\(/, 'IN keyword and opening paren preserved');
+    # The full squashed form: IN ($N /*, ... */) with nothing else inside parens
+    like($q, qr/IN\s*\(\s*\$\d+\s+\/\*,\s*\.\.\.\s*\*\/\s*\)/,
+        'IN list fully squashed to single placeholder with comment');
+};
+
+# ---------------------------------------------------------------------------
+# Test 2: IN list with string constants gets squashed
+# ---------------------------------------------------------------------------
+subtest 'IN list with strings is squashed' => sub {
+    my $q = get_captured_query($node,
+        "SELECT * FROM test_squash WHERE val IN ('alpha', 'beta', 'gamma')");
+
+    like($q, qr/\$\d/, 'IN list has placeholder');
+    like($q, qr{/\*, \.\.\. \*/}, 'Squashed comment marker present');
+    unlike($q, qr/alpha/, 'String constant alpha removed');
+    unlike($q, qr/beta/, 'String constant beta removed');
+    unlike($q, qr/gamma/, 'String constant gamma removed');
+};
+
+# ---------------------------------------------------------------------------
+# Test 3: Mixed query — IN list squashed AND other constants normalized
+# ---------------------------------------------------------------------------
+subtest 'IN list squashed alongside regular constant' => sub {
+    my $q = get_captured_query($node,
+        "SELECT * FROM test_squash WHERE id IN (10, 20, 30) AND val = 'secret'");
+
+    like($q, qr{/\*, \.\.\. \*/}, 'Squashed comment marker present for IN list');
+    unlike($q, qr/\b10\b/, 'IN constant 10 removed');
+    unlike($q, qr/\b20\b/, 'IN constant 20 removed');
+    unlike($q, qr/\b30\b/, 'IN constant 30 removed');
+    unlike($q, qr/secret/, 'Standalone string constant removed');
+    # Should have at least two distinct placeholders ($N for IN, $M for val)
+    like($q, qr/\$1/, 'First placeholder present');
+    like($q, qr/\$2/, 'Second placeholder present');
+};
+
+# ---------------------------------------------------------------------------
+# Test 4: Large IN list to ensure full list coverage
+# ---------------------------------------------------------------------------
+subtest 'large IN list fully squashed' => sub {
+    my @nums = (100 .. 119);
+    my $in_list = join(', ', @nums);
+    my $q = get_captured_query($node,
+        "SELECT * FROM test_squash WHERE id IN ($in_list)");
+
+    like($q, qr{/\*, \.\.\. \*/}, 'Squashed comment marker present');
+    # Verify no raw numbers from the list survive
+    for my $n (@nums) {
+        unlike($q, qr/\b$n\b/, "IN constant $n removed");
+    }
+    # The normalized form should be compact, not contain 20 separate placeholders
+    unlike($q, qr/\$5/, 'No excessive placeholders (list is squashed, not expanded)');
+};
+
+# ---------------------------------------------------------------------------
+# Test 5: NOT IN list also squashed
+# ---------------------------------------------------------------------------
+subtest 'NOT IN list squashed' => sub {
+    my $q = get_captured_query($node,
+        "SELECT * FROM test_squash WHERE id NOT IN (7, 8, 9)");
+
+    like($q, qr{/\*, \.\.\. \*/}, 'Squashed comment marker present for NOT IN');
+    unlike($q, qr/\b[7-9]\b/, 'NOT IN constants removed');
+};
+
+# ---------------------------------------------------------------------------
+# Cleanup
+# ---------------------------------------------------------------------------
+$node->safe_psql('postgres', 'DROP TABLE IF EXISTS test_squash');
+$node->stop();
+done_testing();