Let the regex gate see set_config()

[Bougerous]

Follow-up to the discussion on #1298, opened separately at @IgorOhrimenko's suggestion since regex_parser.rs isn't a file that PR touches.

The problem

SELECT set_config(...) changes session state exactly like SET, but it's a function call inside a SELECT rather than a statement-start keyword, so it matches none of the patterns in CMD_BASE. The regex fast path never hands the statement to the parser, so the interception in QueryParser::set_config can't run — no matter how well it handles the statement once it gets there.

Where that bites:

• At session_control / session_control_and_locks, unconditionally — those levels consult only the gate.
• At auto, whenever the cluster doesn't already force the parser on. Cluster::router_needed() is false for a single shard that is primary-only or replica-only, so a plain unsharded deployment with no read/write split falls through to the gate.

On those, a session-scoped set_config() survives checkin and the next client inherits it. Adding a replica to an otherwise identical config makes it go away, which is what makes this easy to miss.

The change

One unanchored pattern. Chaining it into cmd_base_patterns() rather than CMD_ADVISORY puts it in both CMD_RE and CMD_RE_ADVISORY — set_config() is session control, not a lock, so it shouldn't sit behind the locks-only level. The \b anchors keep it from matching identifiers that merely contain the word; SELECT offset_configuration FROM t is in the test as the negative case.

Verification

test_set_config covers the literal form, the $1 bind form, pg_catalog.set_config, and a comment-prefixed variant, at SessionControl, SessionControlAndLocks and Auto. The regex_parser module is 17/17.

I measured the underlying behaviour against a pool_size = 1 cluster with pg_backend_pid() compared between clients and log_statement=all to see what actually reached the server; details are in the comment on #1298.

This should also turn the pgdog_leak_auto cases in #1298 green, which is where it came from.

One note on running the tests locally

I ran the unit tests at 9918e963 rather than at bafec81f. regex_parser.rs is byte-identical between the two, but current main doesn't link on macOS/arm64: since #1324 made pg_raw_parse non-optional, the build fails with

Undefined symbols for architecture arm64:
  "_wrapped_raw_expression_tree_walker_impl", referenced from:
      pg_raw_parse::walk::walk_node_cb::...

The generated wrap_static_fns.c in the build directory is 10 lines and doesn't mention raw_expression_tree_walker_impl, so bindgen isn't emitting a wrapper for it on that platform — presumably a static inline that clang handles differently there than on the Linux CI images. Unrelated to this change, and cargo build --lib is fine; only linking an executable fails. Happy to open a separate issue with the details if that's useful.

[CLAassistant]

All committers have signed the CLA.

[levkk]

I would be more comfortable if this regex was more like:

SELECT set_config

because the \bset_config\b scan is relatively expensive on long queries and will return nothing 99.99% of the time, while still running in O(n) time (I think).

[Bougerous]

You're right that it's very expensive, so i did a benchmark to learn and experiment .

Every pattern in CMD_BASE is ^-anchored through COMMENT_PREFIX, so CMD_RE rejects a non-matching query at offset 0 without reading it. Adding any unanchored alternative removes that and the whole set starts scanning. I measured it in a standalone crate with regex pinned to =1.12.4, the patterns copied verbatim and the same 1000-byte limit, over three non-matching queries (673B, 56B, and a offset_configuration near-miss):

base: anchored patterns only                8.0 ns/query
set_config folded into the set            295.5 ns/query   <- this PR as written
"SELECT set_config" folded into the set   235.3 ns/query   <- your suggestion

So your version is about 20% cheaper than mine, but both are ~30x the baseline. The literal doesn't get you out of the scan, because it's unanchored too.

It also drops four forms that reach set_config:

	\bset_config\b	SELECT set_config
SELECT set_config('a','b',false)	match	match
SELECT pg_catalog.set_config(...)	match	no
select + 2 spaces + set_config(...)	match	no
SELECT\n set_config(...)	match	no
SELECT 1, set_config(...)	match	no
SELECT offset_configuration FROM t	no	no

The pg_catalog. one is the form the Python test in #1298 uses, so that combination would go red on rebase.

The guard

Answering what you actually asked: yes, and it's most of the cost back. Keep both sets anchored-only and test set_config as its own Regex, after the set has already said no:

base set, then standalone Regex            49.1 ns/query

Same results on all 12 inputs I tried, positive and negative. The reason it's 6x cheaper than folding the identical pattern into the set is that RegexSet unions everything into one automaton and loses the literal prefilter; as its own Regex it keeps the SIMD substring scan for set_config. I also tried an explicit aho-corasick pre-check in front of it and it came out at 51.7 ns — no better, so there's no reason to take the dependency.

Two things on the remaining 41 ns. It's bounded by truncate_utf8(query.query(), self.limit), so it's O(min(n, regex_parser_limit)) — 1000 bytes by default, not O(n), and it doesn't degrade on long queries. And CMD_RE_ADVISORY measures 292.6 ns on the same inputs today, so at session_control_and_locks this is about a sixth of what the gate already costs.

Shape:

static SET_CONFIG_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"(?i)\bset_config\b").unwrap());

let prefix = truncate_utf8(query.query(), self.limit);
let cmds = if with_locks { &*CMD_RE_ADVISORY } else { &*CMD_RE };
return cmds.is_match(prefix) || SET_CONFIG_RE.is_match(prefix);

Happy to push that, and to post the benchmark source if you want to rerun it — it's about 100 lines and has no dependency on the pgdog tree.

[levkk]

Can you anchor it? We should make sure it doesn't scan the whole statement if it doesn't match the start of the query.

[IgorOhrimenko]

Independent numbers on the anchoring question, since it's a performance argument
and those are cheap to settle.

Benchmark and corpus generator:
https://gist.github.com/IgorOhrimenko/5f88066cd664b572bd6ed8b311b9b1be — three
files, two dependencies, cargo run --release.

The corpus turned out to matter more than I expected. My first attempt used an
invented statement mix and understated the gap by about 4x, because the
statements were roughly a quarter of real length and an unanchored scan costs in
proportion to length. So the generator is shaped from 2089 distinct query texts
out of a production PgDog's AST cache, reduced to aggregates — kind mix, length
deciles, identifier density. None of the sampled text is in it.

On 2089 statements truncated at regex_parser_limit, none containing
set_config, which is what the gate spends effectively all of its time on:

base: anchored patterns only                        23.6 ns/query
A. \bset_config\b in the set                       606.3
B. "SELECT set_config" in the set                  593.0
C. anchored SELECT + set_config in the set         541.0
D. set, then \bset_config\b as its own Regex        86.0
E. set, then the anchored pattern on its own       496.2
for scale: CMD_RE_ADVISORY as it is today          522.9

Against the real sample rather than the generator: 17.8 / 446.6 / 443.4 / 427.0
/ 57.5 / 416.4 / 452.5. The generator is 15–30% pessimistic and preserves
the ordering. Laptop under WSL2, so the ratios are the result and the absolute
numbers are indicative.

Anchoring was worth trying and it doesn't pay. I built
^\s*SELECT\b[\s\S]*\bset_config\b and measured it both folded into the set and
standing alone. The anchor only buys a cheap rejection for statements that don't
begin with SELECT, and in this traffic 90% of them do — so nearly everything
gets scanned anyway, and the anchored pattern has no literal to prefilter on. On
non-SELECT traffic alone it does win, 29.0 against 51.5 ns, but that is the
minority of what arrives. It also drops
WITH x AS (...) SELECT set_config(...), so it ends up both slower and
narrower.

The literal is cheaper than the word-boundary version, but only just — 593
against 606 here. I'd add that this gap is not stable: on a smaller sample with
longer statements the literal came out ~4x cheaper, and on this one they're
within 2%. Either way it drops pg_catalog.set_config(...), select with extra
whitespace, a newline before the call, and SELECT 1, set_config(...). The
pg_catalog. form is the one libpq-based clients and pg_dump emit.

What actually moves the number is keeping the pattern out of the RegexSet.
Unioning it in loses the literal prefilter; as its own Regex the SIMD substring
scan survives. That is 606 → 86 ns, and 86 ns is a sixth of what
CMD_RE_ADVISORY already costs on the same corpus at
session_control_and_locks — the gate is already paying six times this to look
for advisory locks.

So the standalone-Regex shape proposed earlier in this thread looks right to
me, and I couldn't find a cheaper one that keeps the forms. Happy to rerun any
variant against the generator if there's another shape worth measuring.

[Bougerous]

Thank you for comment , the point about statement length dominating
an unanchored scan is one I'd missed, and my own numbers were measured on an
invented mix, so I'd take yours over mine.

Your harness measures five variants, and the shape currently on the PR isn't one
of them. It isn't D and it isn't E. E bridges to the call with [\s\S]*, which
leaves nothing to prefilter on. What's pushed is a short anchored bail, then
the standalone literal as a second step:

SET_CONFIG_LEAD  =  COMMENT_PREFIX + r"[(\s]*(?:SELECT|WITH)\b"
SET_CONFIG_RE    =  r"(?i)\bset_config\b"

cmds.is_match(prefix) || (LEAD.is_match(prefix) && RE.is_match(prefix))

Because the anchor never spans the statement, the SIMD scan on the second regex
survives. I added it to your generator as F and ran it three times:

base: anchored patterns only                       9.5 ns/query
A. \bset_config\b in the set                     417.5
D. set, then \bset_config\b as its own Regex      40.7
E. set, then the anchored pattern on its own     392.6
F. set, then anchored bail + own Regex            46.0
for scale: CMD_RE_ADVISORY as it is today        420.5

Apple silicon rather than WSL2, so these run about 2.4x faster than yours in
absolute terms; the ordering is identical.

Split by traffic shape, since that's the whole difference:

	D	F
non-SELECT only	40.1 ns	18.2 ns
SELECT only	40.0 ns	47.0 ns

D is flat. F is 2.2x cheaper when the anchor bails and ~17% dearer when it
can't, so they cross at about 76% SELECT — your corpus is 90%, and D wins there
by 12%.

On coverage, F doesn't give anything up:

form	C	D	E	F
SELECT pg_catalog.set_config($1, $2, false)	yes	yes	yes	yes
WITH x AS (SELECT 1) SELECT set_config(...)	no	yes	no	yes
(SELECT set_config(...))	no	yes	no	yes
-- audit\nSELECT set_config(...)	yes	yes	yes	yes
SELECT offset_configuration FROM t	no	no	no	no

So "both slower and narrower" is fair against the [\s\S]* form, but the
narrowing comes from the bridge rather than from anchoring as such. F keeps
WITH and the wrapping paren, and costs 12% on your corpus instead of 10x.

One correction to my own earlier table, which your list inherited:
SELECT 1, set_config(...) shouldn't be counted as a form we lose.
extract_set_config uses .exactly_one() on the target list, so the parser
declines it whatever the gate does.

Either shape is fine by me — 5 ns on a path already spending 420 on advisory
locks. @levkk asked for the anchor and F is the version that keeps it without
the coverage loss, so I've left it pushed, but I'll switch to D if you'd both
rather have the simpler code.