Skip to content

fix(multipooler): reset role/session_authorization on connection recycle#8

Open
haritabh17 wants to merge 1 commit into
mainfrom
hex-reset-role-on-recycle
Open

fix(multipooler): reset role/session_authorization on connection recycle#8
haritabh17 wants to merge 1 commit into
mainfrom
hex-reset-role-on-recycle

Conversation

@haritabh17

Copy link
Copy Markdown
Owner

Problem

When a pooled connection has SET ROLE or SET SESSION AUTHORIZATION applied via ApplySettings, these GUCs persist on the PG backend after the connection is returned to the pool. PostgreSQL marks both with GUC_NO_RESET_ALL, so RESET ALL alone doesn't clear them.

If the role is subsequently dropped by another concurrent session, the next user of that pooled connection gets ERROR: role NNN was concurrently dropped from PG's shdepLockAndCheckObject during DDL operations.

Root Cause

  1. SET ROLE / SET SESSION AUTHORIZATION tracked in multigateway SessionSettings, applied lazily via set_config() in ApplySettings
  2. When recycled, PG backend retains the role/session_auth (GUC_NO_RESET_ALL)
  3. ResetAllSettings only fires when connSettings != settings — if the next checkout matches the same *Settings pointer, no reset occurs
  4. A concurrent DROP ROLE causes the stale connection to fail

Fix

Add ResetRoleState() on regular.Conn — called before conn.Recycle() in all executor paths:

  • Checks if tracked settings include role or session_authorization
  • If so: sends RESET ROLE; RESET SESSION AUTHORIZATION; RESET ALL to PG
  • Clears tracked settings to nil (routes connection to pool.clean)
  • Next checkout re-applies full desired settings via ApplySettings

Important: Settings objects are cached and shared via SettingsCache — we must not mutate Vars map. Instead we clear the connection's tracked settings pointer to nil.

Testing

pgregress results:

  • Role-dropped errors: 464 → 0
  • Passing tests: 122 → 123 (no regressions, gained maintain_every)

When a pooled connection has SET ROLE or SET SESSION AUTHORIZATION
applied via ApplySettings, these GUCs persist on the PG backend even
after the connection is returned to the pool. PostgreSQL marks both
with GUC_NO_RESET_ALL, so RESET ALL alone doesn't clear them.

If the role is subsequently dropped by another concurrent session
(e.g., during pgregress parallel tests), the next user of that pooled
connection gets 'role NNN was concurrently dropped' from PG's
shdepLockAndCheckObject during DDL operations.

Fix: Add ResetRoleState() on regular.Conn that checks if tracked
settings include 'role' or 'session_authorization'. If so, sends
RESET ROLE; RESET SESSION AUTHORIZATION; RESET ALL to the PG backend
and clears tracked settings to nil (routing the connection to
pool.clean). The next checkout re-applies full desired settings via
ApplySettings.

The Settings objects in the pool are cached and shared via
SettingsCache, so we must not mutate their Vars map — instead we
clear the connection's tracked settings pointer to nil.

Called from all executor paths (StreamExecute, PortalStreamExecute,
DescribeStatement) before conn.Recycle().

Tested: pgregress role-dropped errors 464 -> 0, no regressions
(122 -> 123 passing tests).
haritabh17 added a commit that referenced this pull request Mar 27, 2026
…ener

Migrate PubSubListener from direct client.Config connection to pool-managed
reserved connection via PoolManager.NewReservedConn(), following the COPY
pattern. This ensures the pool accurately tracks all backend connections
and enables lazy connection acquisition/release.

Changes:
- Add RESERVATION_REASON_LISTEN (16) to proto enum and protoutil constants
- Add SendQuery, ReadRawMessage, ParseNotification to reserved.Conn
- Replace client.Config with PoolManager in Listener struct and constructor
- Use reserved.Conn directly (no regular.Conn extraction)
- Release connection with ReleaseError on failure/shutdown
- Release connection when all channels are unsubscribed (lazy lifecycle)
- Fix deadlock when reqUnsubscribeAll partial failure nils readerCh
- Remove ListenerClientConfig() from PoolManager interface (dead code)
- Tone down HTTP/2 overhead claim in notification_adapter.go
- Update design doc to reflect pool-managed architecture

Addresses PR multigres#750 review comments #1, #8, multigres#10.

Signed-off-by: Haritabh Gupta <20576107+haritabh17@users.noreply.github.com>
haritabh17 added a commit that referenced this pull request Mar 27, 2026
…ener

Migrate PubSubListener from direct client.Config connection to pool-managed
reserved connection via PoolManager.NewReservedConn(), following the COPY
pattern. This ensures the pool accurately tracks all backend connections
and enables lazy connection acquisition/release.

Changes:
- Add RESERVATION_REASON_LISTEN (16) to proto enum and protoutil constants
- Add SendQuery, ReadRawMessage, ParseNotification to reserved.Conn
- Replace client.Config with PoolManager in Listener struct and constructor
- Use reserved.Conn directly (no regular.Conn extraction)
- Release connection with ReleaseError on failure/shutdown
- Release connection when all channels are unsubscribed (lazy lifecycle)
- Fix deadlock when reqUnsubscribeAll partial failure nils readerCh
- Remove ListenerClientConfig() from PoolManager interface (dead code)
- Tone down HTTP/2 overhead claim in notification_adapter.go
- Update design doc to reflect pool-managed architecture

Addresses PR multigres#750 review comments #1, #8, multigres#10.

Signed-off-by: Haritabh Gupta <20576107+haritabh17@users.noreply.github.com>
haritabh17 added a commit that referenced this pull request Mar 30, 2026
…ener

Migrate PubSubListener from direct client.Config connection to pool-managed
reserved connection via PoolManager.NewReservedConn(), following the COPY
pattern. This ensures the pool accurately tracks all backend connections
and enables lazy connection acquisition/release.

Changes:
- Add RESERVATION_REASON_LISTEN (16) to proto enum and protoutil constants
- Add SendQuery, ReadRawMessage, ParseNotification to reserved.Conn
- Replace client.Config with PoolManager in Listener struct and constructor
- Use reserved.Conn directly (no regular.Conn extraction)
- Release connection with ReleaseError on failure/shutdown
- Release connection when all channels are unsubscribed (lazy lifecycle)
- Fix deadlock when reqUnsubscribeAll partial failure nils readerCh
- Remove ListenerClientConfig() from PoolManager interface (dead code)
- Tone down HTTP/2 overhead claim in notification_adapter.go
- Update design doc to reflect pool-managed architecture

Addresses PR multigres#750 review comments #1, #8, multigres#10.

Signed-off-by: Haritabh Gupta <20576107+haritabh17@users.noreply.github.com>
haritabh17 added a commit that referenced this pull request Mar 30, 2026
…gres#750)

* feat(proto): add LISTEN/NOTIFY proto definitions and async notification support

Add StreamNotifications RPC to multipooler service proto for streaming
PG notifications from multipooler to gateway. Add Notification type to
sqltypes for internal notification handling. Add async notification
pusher to server.Conn for delivering NotificationResponse messages to
clients between queries.

Signed-off-by: Haritabh Gupta <20576107+haritabh17@users.noreply.github.com>

* feat(multipooler): add PubSubListener with split read/write and StreamNotifications gRPC

Add PubSubListener that manages a single dedicated PG connection for all
LISTEN/NOTIFY operations. Uses split read/write on the TCP connection:
reader goroutine reads all PG messages via ReadRawMessage(), event loop
writes LISTEN/UNLISTEN via SendQuery(). This avoids reconnecting on
subscribe/unsubscribe, preventing notification loss for existing channels.

Key design:
- Channel refcounting: LISTEN on first subscriber, UNLISTEN on last
- pendingCmds counter (not boolean) to correctly drain multiple
  ReadyForQuery messages from batch UNLISTEN operations
- Reconnect only on actual connection failure, re-LISTENs all channels

Also adds SendQuery(), ReadRawMessage(), ParseNotification() to
client.Conn to support the split read/write pattern, and exposes
PubSubListener via ConnPoolManager and StreamNotifications gRPC.

Signed-off-by: Haritabh Gupta <20576107+haritabh17@users.noreply.github.com>

* feat(multigateway): add LISTEN/NOTIFY support with transaction semantics

Implement gateway-side LISTEN/NOTIFY handling:

- Planner intercepts LISTEN/UNLISTEN in both simple (Plan) and extended
  (PlanPortal) query protocols, creating ListenNotifyPrimitive. NOTIFY
  is routed to PG as a regular query.
- ListenNotifyPrimitive updates per-connection state, buffering changes
  inside transactions and applying immediately in autocommit mode.
- Channel names truncated to 63 chars (NAMEDATALEN-1) matching PG.
- handleListenStateChanges syncs subscriptions after COMMIT/ROLLBACK.
- GRPCNotificationManager manages per-channel gRPC streams to the
  multipooler with automatic retry on stream failure. Subscribe waits
  outside the lock to avoid the unlock/relock race window.
- CommitPendingListens avoids redundant individual unsubscribes when
  UNLISTEN * is pending (UnsubscribeAll handles everything).
- Async notification delivery between queries via forwardNotifications
  goroutine; synchronous flush after query via flushNotifications.

Signed-off-by: Haritabh Gupta <20576107+haritabh17@users.noreply.github.com>

* test(queryserving): add LISTEN/NOTIFY integration tests

Add comprehensive end-to-end tests covering:
- Basic LISTEN/NOTIFY with payload
- UNLISTEN specific channel and UNLISTEN *
- Multiple listeners on same channel
- Transaction semantics (LISTEN/UNLISTEN deferred to COMMIT)
- ROLLBACK discards pending LISTEN/UNLISTEN
- pg_notify() function
- Channel name case sensitivity and special characters
- Empty and NULL payload handling
- Extended query protocol (Parse/Bind/Execute) for LISTEN

Signed-off-by: Haritabh Gupta <20576107+haritabh17@users.noreply.github.com>

* docs(listen-notify): add design doc with architecture and semantics

Document the shared listener architecture, notification delivery flow,
transaction semantics, PubSubListener split read/write design, gateway
subscription management, cleanup lifecycle, and resilience model.
Includes future optimization note about single gRPC stream consolidation.

Signed-off-by: Haritabh Gupta <20576107+haritabh17@users.noreply.github.com>

* fix(pgregresstest): build pg_isolation_regress binary and use source dir for inputs

Build the pg_isolation_regress binary if it doesn't exist before running
selective isolation tests. Use the source directory (not build directory)
for --inputdir so test spec files are found correctly.

Signed-off-by: Haritabh Gupta <20576107+haritabh17@users.noreply.github.com>

* refactor(listen-notify): address PR review comments #2-7, #9, multigres#11-13

- Use ast.QuoteIdentifier instead of local quoteIdent helper (#2)
- Remove duplicate UnsubscribeCh method, keep Unsubscribe (#3)
- Remove unused Subscribe method that allocates channels (#4)
- Remove is_warning/warning_message proto fields (dead code) (#5)
- Consolidate 4 duplicate Notification structs into sqltypes.Notification (#6, #7, multigres#12)
- Wire PubSubListener into state manager for PRIMARY transitions (#9)
- Rename ListenNotifyPrimitive.SQL to Query for consistency (multigres#11)
- Consolidate WriteNotificationResponse with writeNotificationResponseMsg (multigres#13)

Signed-off-by: Haritabh Gupta <20576107+haritabh17@users.noreply.github.com>

* refactor(listen-notify): move subscription sync from handler into engine primitives

Move LISTEN/NOTIFY subscription management out of handler post-execution
hooks and into the engine primitives that own the logic:

- ListenNotifyPrimitive: syncs subscriptions for autocommit LISTEN/UNLISTEN
- TransactionPrimitive: syncs on COMMIT, discards on ROLLBACK

Introduce SubscriptionSync interface on connection state so primitives can
call into the handler's notification infrastructure without a direct
dependency. This eliminates handleListenStateChanges() and the associated
post-execution hooks in HandleQuery, HandleExecute, and
executeWithImplicitTransaction.

Also fixes a goroutine leak where forwardNotifications outlived its
useful lifetime by giving it a cancellable context, and cleans up
NotifCh when all channels are unlistened.

Signed-off-by: Haritabh Gupta <20576107+haritabh17@users.noreply.github.com>

* refactor(pubsub): use pool-managed reserved connection for PubSubListener

Migrate PubSubListener from direct client.Config connection to pool-managed
reserved connection via PoolManager.NewReservedConn(), following the COPY
pattern. This ensures the pool accurately tracks all backend connections
and enables lazy connection acquisition/release.

Changes:
- Add RESERVATION_REASON_LISTEN (16) to proto enum and protoutil constants
- Add SendQuery, ReadRawMessage, ParseNotification to reserved.Conn
- Replace client.Config with PoolManager in Listener struct and constructor
- Use reserved.Conn directly (no regular.Conn extraction)
- Release connection with ReleaseError on failure/shutdown
- Release connection when all channels are unsubscribed (lazy lifecycle)
- Fix deadlock when reqUnsubscribeAll partial failure nils readerCh
- Remove ListenerClientConfig() from PoolManager interface (dead code)
- Tone down HTTP/2 overhead claim in notification_adapter.go
- Update design doc to reflect pool-managed architecture

Addresses PR multigres#750 review comments #1, #8, multigres#10.

Signed-off-by: Haritabh Gupta <20576107+haritabh17@users.noreply.github.com>

* fix(listen-notify): fix notification delivery race, deadlock, and refcount bugs

Fix three pre-existing bugs in the LISTEN/NOTIFY implementation:

1. flushNotifications race condition: flushNotifications and
   forwardNotifications both read from the same NotifCh, causing
   notifications to go to either reader nondeterministically.
   Additionally, WriteNotificationResponse writes to the client socket
   without holding bufMu, racing with the async pusher goroutine.
   Fix: add FlushPendingNotifications() to server.Conn that drains the
   async pusher channel with proper bufMu locking. Remove the now-unused
   WriteNotificationResponse method.

2. Subscribe/Unsubscribe deadlock on Listener stop: when the Listener
   stops (PRIMARY → REPLICA transition), the event loop exits and nobody
   reads from the requests channel. Pending SubscribeCh/Unsubscribe calls
   from gRPC stream defer blocks would block forever.
   Fix: add a stopped channel that is closed when the event loop exits.
   Public methods select on it to return immediately when stopped.

3. removeSub unconditional refcount decrement: removeSub always
   decrements the channel refcount and removes from allSubs, even if
   the subscriber was not found. This corrupts refcounts on duplicate
   unsubscribe calls and prematurely removes multi-channel subscribers
   from allSubs (preventing cleanup on shutdown).
   Fix: check if the subscriber was actually found before decrementing,
   and only remove from allSubs when no subscriptions remain.

Signed-off-by: Haritabh Gupta <20576107+haritabh17@users.noreply.github.com>

* fix(listen-notify): skip duplicate LISTEN in autocommit mode

Repeated LISTEN for the same channel outside a transaction always called
Subscribe on the notification manager, adding duplicate subscribers.
This caused notifications to be delivered multiple times and left stale
gRPC streams after a single UNLISTEN (which only removes one copy).

Add IsListening() check to skip the subscribe when the channel is
already active. The transaction path (CommitPendingListens) already had
this dedup check; this aligns the autocommit path to match.

Signed-off-by: Haritabh Gupta <20576107+haritabh17@users.noreply.github.com>

* fix(listen-notify): preserve LISTEN/UNLISTEN ordering in transactions

Replace three separate pending fields (PendingListens, PendingUnlistens,
PendingUnlistenAll) with a single ordered action list that processes
operations in the order they were issued, matching PostgreSQL's
AtCommit_Notify behavior. A net diff against the pre-transaction state
produces the subscribe/unsubscribe lists for the notification manager.

Previously, BEGIN; LISTEN x; UNLISTEN x; COMMIT would incorrectly
subscribe because all unlistens were processed before all listens.

Signed-off-by: Haritabh Gupta <20576107+haritabh17@users.noreply.github.com>

* docs(listen-notify): fix markdown table alignment and update design doc

Signed-off-by: Haritabh Gupta <20576107+haritabh17@users.noreply.github.com>

* refactor(server): remove unused WriteNotice method

Signed-off-by: Haritabh Gupta <20576107+haritabh17@users.noreply.github.com>

* fix(planner): pass nil txnMetrics to NewPlanner in prepared_stmt_test

Signed-off-by: Haritabh Gupta <20576107+haritabh17@users.noreply.github.com>

---------

Signed-off-by: Haritabh Gupta <20576107+haritabh17@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant