Skip to main content

pgroles_operator/
context.rs

1//! Shared operator context — database pool cache, metrics, and configuration.
2
3use std::collections::HashMap;
4use std::fmt::Write;
5use std::str::FromStr;
6use std::sync::Arc;
7use std::time::{Duration, SystemTime};
8
9use futures::future::{BoxFuture, FutureExt};
10use kube::runtime::events::Recorder;
11use serde::{Deserialize, Serialize};
12
13use sha2::{Digest, Sha256};
14use sqlx::postgres::{PgConnectOptions, PgPool, PgPoolOptions};
15use tokio::sync::{Mutex, RwLock};
16
17use crate::crd::{ConnectionAuth, ConnectionSpec, SecretKeySelector};
18use crate::observability::OperatorObservability;
19use crate::request_index::RequestIndex;
20
21/// Minimum pool size required for reconciliation.
22///
23/// One connection is held for the session-scoped advisory lock while the
24/// reconcile loop performs inspection and apply work on the pool.
25const POOL_MAX_CONNECTIONS: u32 = 5;
26
27/// Bound how long a reconcile waits for a pooled connection before surfacing
28/// a transient database connectivity failure.
29const POOL_ACQUIRE_TIMEOUT_SECS: u64 = 10;
30
31const _: () = assert!(POOL_MAX_CONNECTIONS >= 2);
32
33/// How long a pooled connection may sit unused before it is closed.
34///
35/// Pools are cached for the operator's lifetime, so without this the operator
36/// holds connections open against every managed database forever. sqlx's
37/// default is 10 minutes — longer than the 5-minute default requeue interval,
38/// so a reconcile always re-touches the pool before the reaper drains it and
39/// the connections are never released. Reconnecting once per reconcile is
40/// cheap next to occupying a backend slot on an idle database around the clock.
41const POOL_IDLE_TIMEOUT_SECS: u64 = 60;
42
43/// Hard ceiling on connection age, independent of activity.
44///
45/// Bounds how long a reconcile can keep using a connection whose server-side
46/// state has drifted (failover, restarted pooler, rotated credentials).
47const POOL_MAX_LIFETIME_SECS: u64 = 30 * 60;
48
49/// Keep no connections open when nothing is reconciling.
50const POOL_MIN_CONNECTIONS: u32 = 0;
51
52const GCP_METADATA_TOKEN_ENDPOINT: &str =
53    "http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/token";
54const GCP_IAM_CREDENTIALS_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform";
55const GCP_TOKEN_CACHE_SKEW_SECS: u64 = 300;
56const GCP_IMPERSONATED_TOKEN_LIFETIME_SECS: u64 = 3600;
57const GCP_AUTH_HTTP_TIMEOUT_SECS: u64 = 10;
58
59#[derive(Clone)]
60struct CachedPool {
61    resource_version: Option<String>,
62    /// Fingerprint of all referenced secrets' resourceVersions (params mode).
63    secret_fingerprint: Option<String>,
64    /// Expiry for token-backed connection passwords.
65    token_expires_at: Option<SystemTime>,
66    pool: PgPool,
67}
68
69struct ResolvedConnectionUrl {
70    database_url: String,
71    token_expires_at: Option<SystemTime>,
72    /// Optional role to `SET ROLE` to on every pooled connection. The value
73    /// has already passed CRD-level identifier validation; the after-connect
74    /// hook re-quotes defensively before interpolating it into SQL.
75    set_role: Option<String>,
76}
77
78/// Build the `SET ROLE` SQL statement for the given identifier.
79///
80/// `SET ROLE` does not accept bind parameters, so the value is interpolated.
81/// CRD admission already restricts the identifier to the
82/// [`crate::crd::SET_ROLE_PATTERN`] regex; the embedded `"` doubling here is
83/// defence in depth for the connection-pool path.
84pub fn build_set_role_stmt(role: &str) -> String {
85    format!("SET ROLE \"{}\"", role.replace('"', "\"\""))
86}
87
88/// Marker prefix used to flag SET-ROLE failures from the `after_connect`
89/// hook so they can be distinguished from other connect-time errors.
90const SET_ROLE_FAILURE_MARKER: &str = "pgroles:set-role-failed:";
91
92/// Wrap a SET-ROLE failure in [`sqlx::Error::Protocol`] with a marker so
93/// the outer `connect()` error can be classified as
94/// [`ContextError::SetRoleFailed`] instead of a generic database connect
95/// failure.
96fn wrap_set_role_failure(role: &str, source: sqlx::Error) -> sqlx::Error {
97    sqlx::Error::Protocol(format!("{SET_ROLE_FAILURE_MARKER}{role}: {source}"))
98}
99
100/// Classify a pool-level connect error, surfacing SET-ROLE hook failures
101/// distinctly from genuine database-connect failures.
102fn classify_pool_connect_error(set_role: Option<&str>, err: sqlx::Error) -> ContextError {
103    if let (Some(role), sqlx::Error::Protocol(msg)) = (set_role, &err)
104        && msg.starts_with(SET_ROLE_FAILURE_MARKER)
105    {
106        return ContextError::SetRoleFailed {
107            role: role.to_string(),
108            source: err,
109        };
110    }
111    ContextError::DatabaseConnect { source: err }
112}
113
114#[derive(Clone)]
115struct GcpAccessToken {
116    token: String,
117    expires_at: SystemTime,
118}
119
120trait GcpAccessTokenProvider: Send + Sync {
121    fn fetch_token<'a>(
122        &'a self,
123        auth: &'a ConnectionAuth,
124    ) -> BoxFuture<'a, Result<GcpAccessToken, ContextError>>;
125}
126
127#[derive(Clone)]
128struct MetadataGcpAccessTokenProvider {
129    client: reqwest::Client,
130}
131
132impl Default for MetadataGcpAccessTokenProvider {
133    fn default() -> Self {
134        Self {
135            client: reqwest::Client::builder()
136                .no_proxy()
137                .timeout(Duration::from_secs(GCP_AUTH_HTTP_TIMEOUT_SECS))
138                .build()
139                .expect("GCP auth HTTP client should build"),
140        }
141    }
142}
143
144impl GcpAccessTokenProvider for MetadataGcpAccessTokenProvider {
145    fn fetch_token<'a>(
146        &'a self,
147        auth: &'a ConnectionAuth,
148    ) -> BoxFuture<'a, Result<GcpAccessToken, ContextError>> {
149        async move {
150            let scope = auth.gcp_scope();
151            if let Some(target) = auth.gcp_impersonate_service_account() {
152                self.fetch_impersonated_access_token(target, scope).await
153            } else {
154                self.fetch_metadata_access_token(scope).await
155            }
156        }
157        .boxed()
158    }
159}
160
161impl MetadataGcpAccessTokenProvider {
162    async fn fetch_metadata_access_token(
163        &self,
164        scope: &str,
165    ) -> Result<GcpAccessToken, ContextError> {
166        let response = self
167            .client
168            .get(GCP_METADATA_TOKEN_ENDPOINT)
169            .header("Metadata-Flavor", "Google")
170            .query(&[("scopes", scope)])
171            .send()
172            .await
173            .map_err(|source| ContextError::GcpAuthHttp {
174                endpoint: "metadata",
175                source,
176            })?;
177
178        let status = response.status();
179        if !status.is_success() {
180            let body = response_body_for_error(response).await;
181            return Err(ContextError::GcpAuthRejected {
182                endpoint: "metadata".to_string(),
183                status: status.as_u16(),
184                body,
185            });
186        }
187
188        let body: MetadataTokenResponse =
189            response
190                .json()
191                .await
192                .map_err(|source| ContextError::GcpAuthHttp {
193                    endpoint: "metadata",
194                    source,
195                })?;
196
197        if body.access_token.trim().is_empty() {
198            return Err(ContextError::GcpAuthInvalidResponse {
199                detail: "metadata token response omitted access_token".to_string(),
200            });
201        }
202        if body.expires_in == 0 {
203            return Err(ContextError::GcpAuthInvalidResponse {
204                detail: "metadata token response had zero expires_in".to_string(),
205            });
206        }
207
208        Ok(GcpAccessToken {
209            token: body.access_token,
210            expires_at: SystemTime::now() + Duration::from_secs(body.expires_in),
211        })
212    }
213
214    async fn fetch_impersonated_access_token(
215        &self,
216        target_service_account: &str,
217        scope: &str,
218    ) -> Result<GcpAccessToken, ContextError> {
219        let source = self
220            .fetch_metadata_access_token(GCP_IAM_CREDENTIALS_SCOPE)
221            .await?;
222        let encoded_target = percent_encoding::utf8_percent_encode(
223            target_service_account,
224            percent_encoding::NON_ALPHANUMERIC,
225        )
226        .to_string();
227        let endpoint = format!(
228            "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/{encoded_target}:generateAccessToken"
229        );
230        let request = GenerateAccessTokenRequest {
231            scope: vec![scope.to_string()],
232            lifetime: format!("{GCP_IMPERSONATED_TOKEN_LIFETIME_SECS}s"),
233        };
234
235        let response = self
236            .client
237            .post(&endpoint)
238            .bearer_auth(&source.token)
239            .json(&request)
240            .send()
241            .await
242            .map_err(|source| ContextError::GcpAuthHttp {
243                endpoint: "iamcredentials",
244                source,
245            })?;
246
247        let status = response.status();
248        if !status.is_success() {
249            let body = response_body_for_error(response).await;
250            return Err(ContextError::GcpAuthRejected {
251                endpoint: "iamcredentials".to_string(),
252                status: status.as_u16(),
253                body,
254            });
255        }
256
257        let body: GenerateAccessTokenResponse =
258            response
259                .json()
260                .await
261                .map_err(|source| ContextError::GcpAuthHttp {
262                    endpoint: "iamcredentials",
263                    source,
264                })?;
265
266        if body.access_token.trim().is_empty() {
267            return Err(ContextError::GcpAuthInvalidResponse {
268                detail: "IAMCredentials response omitted accessToken".to_string(),
269            });
270        }
271        let expires_at = parse_google_expire_time(&body.expire_time).ok_or_else(|| {
272            ContextError::GcpAuthInvalidResponse {
273                detail: format!(
274                    "IAMCredentials response had invalid expireTime {:?}",
275                    body.expire_time
276                ),
277            }
278        })?;
279
280        Ok(GcpAccessToken {
281            token: body.access_token,
282            expires_at,
283        })
284    }
285}
286
287#[derive(Deserialize)]
288struct MetadataTokenResponse {
289    access_token: String,
290    expires_in: u64,
291}
292
293#[derive(Serialize)]
294struct GenerateAccessTokenRequest {
295    scope: Vec<String>,
296    lifetime: String,
297}
298
299#[derive(Deserialize)]
300struct GenerateAccessTokenResponse {
301    #[serde(rename = "accessToken")]
302    access_token: String,
303    #[serde(rename = "expireTime")]
304    expire_time: String,
305}
306
307async fn response_body_for_error(response: reqwest::Response) -> String {
308    match response.text().await {
309        Ok(body) => truncate_for_error(body),
310        Err(error) => format!("failed to read error body: {error}"),
311    }
312}
313
314fn truncate_for_error(mut body: String) -> String {
315    const MAX_ERROR_BODY_BYTES: usize = 512;
316    if body.len() <= MAX_ERROR_BODY_BYTES {
317        return body;
318    }
319    let mut end = MAX_ERROR_BODY_BYTES;
320    while !body.is_char_boundary(end) {
321        end -= 1;
322    }
323    body.truncate(end);
324    body.push_str("...");
325    body
326}
327
328fn parse_google_expire_time(expire_time: &str) -> Option<SystemTime> {
329    expire_time
330        .parse::<jiff::Timestamp>()
331        .ok()
332        .map(SystemTime::from)
333}
334
335fn token_expires_after_skew(expires_at: Option<SystemTime>, now: SystemTime) -> bool {
336    let Some(expires_at) = expires_at else {
337        return true;
338    };
339    let Some(refresh_at) = now.checked_add(Duration::from_secs(GCP_TOKEN_CACHE_SKEW_SECS)) else {
340        return false;
341    };
342    expires_at > refresh_at
343}
344
345/// Guard returned by [`OperatorContext::try_lock_database`].
346///
347/// Holding this guard prevents other reconcile loops (within the same process)
348/// from starting work on the same database target. The lock is released when
349/// the guard is dropped.
350pub struct DatabaseLockGuard {
351    key: String,
352    locks: Arc<Mutex<HashMap<String, ()>>>,
353}
354
355impl Drop for DatabaseLockGuard {
356    fn drop(&mut self) {
357        // Best-effort removal — `try_lock` avoids blocking the drop.
358        if let Ok(mut map) = self.locks.try_lock() {
359            map.remove(&self.key);
360            tracing::debug!(database = %self.key, "released in-memory database lock");
361        } else {
362            // Spawn a task to clean up if the mutex is currently held.
363            // Use Handle::try_current() so we don't panic when dropped
364            // outside an active Tokio runtime (e.g. during shutdown).
365            let key = self.key.clone();
366            let locks = Arc::clone(&self.locks);
367            if let Ok(handle) = tokio::runtime::Handle::try_current() {
368                handle.spawn(async move {
369                    locks.lock().await.remove(&key);
370                    tracing::debug!(database = %key, "released in-memory database lock (deferred)");
371                });
372                tracing::debug!(
373                    database = %self.key,
374                    "deferred in-memory database lock release to background task"
375                );
376            } else {
377                // No runtime available — fall back to synchronous cleanup
378                // via blocking_lock so the entry is still removed.
379                let mut map = self.locks.blocking_lock();
380                map.remove(&key);
381                tracing::debug!(
382                    database = %key,
383                    "released in-memory database lock (fallback sync)"
384                );
385            }
386        }
387    }
388}
389
390/// Shared state for the operator, passed to every reconciliation.
391#[derive(Clone)]
392pub struct OperatorContext {
393    /// Kubernetes client for API calls.
394    pub kube_client: kube::Client,
395
396    /// Kubernetes Event recorder for transition-based policy Events.
397    pub event_recorder: Recorder,
398
399    /// Cached database connection pools keyed by `"namespace/secret-name/secret-key"`.
400    pool_cache: Arc<RwLock<HashMap<String, CachedPool>>>,
401    /// In-process per-database reconciliation locks.
402    ///
403    /// Prevents concurrent reconcile loops from operating on the same database
404    /// within a single operator replica. Cross-replica safety is provided by
405    /// PostgreSQL advisory locks (see [`crate::advisory`]).
406    database_locks: Arc<Mutex<HashMap<String, ()>>>,
407
408    /// Shared health/metrics state.
409    pub observability: OperatorObservability,
410
411    /// Watch-fed indexes for relevant ephemeral requests.
412    pub request_index: RequestIndex,
413
414    /// Optional namespace which bounds every operator watch and list.
415    pub watch_namespace: Option<String>,
416
417    /// Fetches short-lived provider-backed database passwords.
418    gcp_token_provider: Arc<dyn GcpAccessTokenProvider>,
419}
420
421impl OperatorContext {
422    /// Create a context with the request index and watch scope shared by the
423    /// controller runtime.
424    pub fn new_with_runtime_config(
425        kube_client: kube::Client,
426        observability: OperatorObservability,
427        event_recorder: Recorder,
428        request_index: RequestIndex,
429        watch_namespace: Option<String>,
430    ) -> Self {
431        Self {
432            kube_client,
433            event_recorder,
434            pool_cache: Arc::new(RwLock::new(HashMap::new())),
435            observability,
436            request_index,
437            watch_namespace,
438            database_locks: Arc::new(Mutex::new(HashMap::new())),
439            gcp_token_provider: Arc::new(MetadataGcpAccessTokenProvider::default()),
440        }
441    }
442
443    /// Try to acquire the in-process lock for the given database identity.
444    ///
445    /// Returns `Some(guard)` if no other reconcile is in progress for this
446    /// database, `None` if one is already running. The lock is released when
447    /// the guard is dropped.
448    pub async fn try_lock_database(&self, database_identity: &str) -> Option<DatabaseLockGuard> {
449        let mut locks = self.database_locks.lock().await;
450        if locks.contains_key(database_identity) {
451            tracing::info!(
452                database = %database_identity,
453                "in-memory database lock contention — another reconcile is in progress"
454            );
455            return None;
456        }
457        locks.insert(database_identity.to_string(), ());
458        tracing::debug!(database = %database_identity, "acquired in-memory database lock");
459        Some(DatabaseLockGuard {
460            key: database_identity.to_string(),
461            locks: Arc::clone(&self.database_locks),
462        })
463    }
464
465    /// Resolve a param from either its literal value or a Secret reference.
466    ///
467    /// Returns `Ok(Some(value))` if one is set, `Ok(None)` if neither is set.
468    async fn resolve_param(
469        &self,
470        namespace: &str,
471        literal: &Option<String>,
472        secret: &Option<SecretKeySelector>,
473    ) -> Result<Option<String>, ContextError> {
474        if let Some(val) = literal {
475            return Ok(Some(val.clone()));
476        }
477        if let Some(sel) = secret {
478            return Ok(Some(
479                self.fetch_secret_value(namespace, &sel.name, &sel.key)
480                    .await?,
481            ));
482        }
483        Ok(None)
484    }
485
486    /// Resolve a [`ConnectionSpec`] into a PostgreSQL connection URL string.
487    ///
488    /// - **URL mode** (`secret_ref` is Some): reads the Secret key as a connection URL.
489    /// - **Params mode** (`params` is Some): resolves each field and constructs a URL.
490    pub async fn resolve_connection_url(
491        &self,
492        namespace: &str,
493        connection: &ConnectionSpec,
494    ) -> Result<String, ContextError> {
495        Ok(self
496            .resolve_connection_url_with_metadata(namespace, connection)
497            .await?
498            .database_url)
499    }
500
501    /// Resolve a stable, non-secret fingerprint of the database target.
502    ///
503    /// Credentials and connection options are deliberately excluded so that
504    /// password and token rotation do not invalidate active access. Host,
505    /// port, or database changes do invalidate it, preventing cleanup from
506    /// revoking an identically named role in a different database.
507    pub async fn resolve_database_target_fingerprint(
508        &self,
509        namespace: &str,
510        connection: &ConnectionSpec,
511    ) -> Result<String, ContextError> {
512        let (host, port, database) = if let Some(ref secret_ref) = connection.secret_ref {
513            let database_url = self
514                .fetch_secret_value(
515                    namespace,
516                    &secret_ref.name,
517                    connection.effective_secret_key(),
518                )
519                .await?;
520            database_target_from_url(&database_url)
521                .map_err(|detail| ContextError::InvalidDatabaseUrl { detail })?
522        } else if let Some(ref params) = connection.params {
523            let host = self
524                .resolve_param(namespace, &params.host, &params.host_secret)
525                .await?
526                .ok_or_else(|| ContextError::EmptyResolvedValue {
527                    field: "host".to_string(),
528                })?;
529            let port = match self
530                .resolve_param(
531                    namespace,
532                    &params.port.map(|port| port.to_string()),
533                    &params.port_secret,
534                )
535                .await?
536            {
537                Some(value) => {
538                    value
539                        .parse::<u16>()
540                        .map_err(|_| ContextError::InvalidResolvedPort {
541                            value: value.clone(),
542                        })?
543                }
544                None => 5432,
545            };
546            let database = self
547                .resolve_param(namespace, &params.dbname, &params.dbname_secret)
548                .await?
549                .ok_or_else(|| ContextError::EmptyResolvedValue {
550                    field: "dbname".to_string(),
551                })?;
552            (host.to_ascii_lowercase(), port, database)
553        } else {
554            return Err(ContextError::SecretMissing {
555                name: "connection".to_string(),
556                key: "neither secretRef nor params is set".to_string(),
557            });
558        };
559
560        if host.trim().is_empty() {
561            return Err(ContextError::EmptyResolvedValue {
562                field: "host".to_string(),
563            });
564        }
565        if database.trim().is_empty() {
566            return Err(ContextError::EmptyResolvedValue {
567                field: "dbname".to_string(),
568            });
569        }
570        Ok(database_target_fingerprint(&host, port, &database))
571    }
572
573    async fn resolve_connection_url_with_metadata(
574        &self,
575        namespace: &str,
576        connection: &ConnectionSpec,
577    ) -> Result<ResolvedConnectionUrl, ContextError> {
578        if let Some(ref secret_ref) = connection.secret_ref {
579            // URL mode — read the full connection URL from the Secret.
580            let database_url = self
581                .fetch_secret_value(
582                    namespace,
583                    &secret_ref.name,
584                    connection.effective_secret_key(),
585                )
586                .await?;
587            Ok(ResolvedConnectionUrl {
588                database_url,
589                token_expires_at: None,
590                set_role: None,
591            })
592        } else if let Some(ref params) = connection.params {
593            // Params mode — resolve each field and build the URL.
594            let host = self
595                .resolve_param(namespace, &params.host, &params.host_secret)
596                .await?
597                .ok_or_else(|| ContextError::EmptyResolvedValue {
598                    field: "host".to_string(),
599                })?;
600            if host.trim().is_empty() {
601                return Err(ContextError::EmptyResolvedValue {
602                    field: "host".to_string(),
603                });
604            }
605
606            let port_str = params.port.map(|p| p.to_string());
607            let port = self
608                .resolve_param(namespace, &port_str, &params.port_secret)
609                .await?
610                .unwrap_or_else(|| "5432".to_string());
611            if port.trim().is_empty() {
612                return Err(ContextError::EmptyResolvedValue {
613                    field: "port".to_string(),
614                });
615            }
616
617            let dbname = self
618                .resolve_param(namespace, &params.dbname, &params.dbname_secret)
619                .await?
620                .ok_or_else(|| ContextError::EmptyResolvedValue {
621                    field: "dbname".to_string(),
622                })?;
623            if dbname.trim().is_empty() {
624                return Err(ContextError::EmptyResolvedValue {
625                    field: "dbname".to_string(),
626                });
627            }
628
629            let username = self
630                .resolve_param(namespace, &params.username, &params.username_secret)
631                .await?
632                .ok_or_else(|| ContextError::EmptyResolvedValue {
633                    field: "username".to_string(),
634                })?;
635            if username.trim().is_empty() {
636                return Err(ContextError::EmptyResolvedValue {
637                    field: "username".to_string(),
638                });
639            }
640
641            let (password, token_expires_at) = if let Some(auth) = &params.auth {
642                let token = self.gcp_token_provider.fetch_token(auth).await?;
643                (token.token, Some(token.expires_at))
644            } else {
645                let password = self
646                    .resolve_param(namespace, &params.password, &params.password_secret)
647                    .await?
648                    .ok_or_else(|| ContextError::EmptyResolvedValue {
649                        field: "password".to_string(),
650                    })?;
651                (password, None)
652            };
653            if password.trim().is_empty() {
654                return Err(ContextError::EmptyResolvedValue {
655                    field: "password".to_string(),
656                });
657            }
658
659            use percent_encoding::{NON_ALPHANUMERIC, utf8_percent_encode};
660            let encoded_username = utf8_percent_encode(&username, NON_ALPHANUMERIC).to_string();
661            let encoded_password = utf8_percent_encode(&password, NON_ALPHANUMERIC).to_string();
662
663            let mut url = format!(
664                "postgresql://{encoded_username}:{encoded_password}@{host}:{port}/{dbname}"
665            );
666
667            let ssl_mode = self
668                .resolve_param(namespace, &params.ssl_mode, &params.ssl_mode_secret)
669                .await?
670                .or_else(|| params.auth.as_ref().map(|_| "require".to_string()));
671            if let Some(ssl_mode) = ssl_mode {
672                // Validate sslMode at runtime — CRD validation only catches
673                // literal values; a secret ref could resolve to anything.
674                if !crate::crd::VALID_SSL_MODES.contains(&ssl_mode.as_str()) {
675                    return Err(ContextError::InvalidResolvedSslMode { value: ssl_mode });
676                }
677                url.push_str("?sslmode=");
678                url.push_str(&ssl_mode);
679            }
680
681            Ok(ResolvedConnectionUrl {
682                database_url: url,
683                token_expires_at,
684                set_role: params.set_role.clone(),
685            })
686        } else {
687            Err(ContextError::SecretMissing {
688                name: "connection".to_string(),
689                key: "neither secretRef nor params is set".to_string(),
690            })
691        }
692    }
693
694    /// Get or create a PgPool for the given connection spec.
695    ///
696    /// Resolves the connection URL from the referenced Secret(s),
697    /// and caches the resulting pool for reuse.
698    pub async fn get_or_create_pool(
699        &self,
700        namespace: &str,
701        connection: &ConnectionSpec,
702    ) -> Result<PgPool, ContextError> {
703        let cache_key = connection.cache_key(namespace);
704
705        // For URL mode, we can do resource-version-based cache invalidation.
706        // For params mode, compute a fingerprint from all referenced secrets'
707        // resourceVersions so that secret rotations invalidate the cache.
708        let (resource_version, secret_fingerprint) =
709            if let Some(ref secret_ref) = connection.secret_ref {
710                let secrets_api: kube::Api<k8s_openapi::api::core::v1::Secret> =
711                    kube::Api::namespaced(self.kube_client.clone(), namespace);
712                let secret = secrets_api.get(&secret_ref.name).await.map_err(|err| {
713                    ContextError::SecretFetch {
714                        name: secret_ref.name.clone(),
715                        namespace: namespace.to_string(),
716                        source: err,
717                    }
718                })?;
719                (secret.metadata.resource_version, None)
720            } else if connection.params.is_some() {
721                // Params mode — collect all referenced secret names and fetch their
722                // resourceVersions to build a fingerprint.
723                let mut secret_names = std::collections::BTreeSet::new();
724                connection.collect_secret_names(&mut secret_names);
725
726                if secret_names.is_empty() {
727                    // All values are literals — no secrets to watch.
728                    (None, Some(String::new()))
729                } else {
730                    let secrets_api: kube::Api<k8s_openapi::api::core::v1::Secret> =
731                        kube::Api::namespaced(self.kube_client.clone(), namespace);
732                    let mut fingerprint_parts = Vec::new();
733                    for name in &secret_names {
734                        let secret = secrets_api.get(name).await.map_err(|err| {
735                            ContextError::SecretFetch {
736                                name: name.clone(),
737                                namespace: namespace.to_string(),
738                                source: err,
739                            }
740                        })?;
741                        let rv = secret
742                            .metadata
743                            .resource_version
744                            .unwrap_or_else(|| "unknown".to_string());
745                        fingerprint_parts.push(format!("{name}={rv}"));
746                    }
747                    (None, Some(fingerprint_parts.join(",")))
748                }
749            } else {
750                (None, None)
751            };
752
753        // Check cache.
754        {
755            let cache = self.pool_cache.read().await;
756            if let Some(cached) = cache.get(&cache_key) {
757                // URL mode: reuse if the Secret's resource_version matches.
758                // Params mode: reuse if the secret fingerprint matches.
759                let version_matches = match (&resource_version, &cached.resource_version) {
760                    (Some(current), Some(cached_rv)) => current == cached_rv,
761                    _ => true,
762                };
763                let fingerprint_matches = match (&secret_fingerprint, &cached.secret_fingerprint) {
764                    (Some(current), Some(cached_fp)) => current == cached_fp,
765                    (None, None) => true,
766                    _ => false,
767                };
768                let token_fresh =
769                    token_expires_after_skew(cached.token_expires_at, SystemTime::now());
770                if version_matches && fingerprint_matches && token_fresh {
771                    return Ok(cached.pool.clone());
772                }
773            }
774        }
775
776        let resolved = self
777            .resolve_connection_url_with_metadata(namespace, connection)
778            .await?;
779
780        // Create pool with explicit sizing. Reconciliation holds one dedicated
781        // connection for PostgreSQL advisory locking and needs additional pool
782        // capacity for inspection/apply queries.
783        let set_role = resolved.set_role.clone();
784        let pool = pool_options()
785            .after_connect(move |conn, _meta| {
786                let set_role = set_role.clone();
787                Box::pin(async move {
788                    if let Some(role) = set_role {
789                        let stmt = build_set_role_stmt(&role);
790                        sqlx::Executor::execute(&mut *conn, stmt.as_str())
791                            .await
792                            .map_err(|err| wrap_set_role_failure(&role, err))?;
793                    }
794                    Ok(())
795                })
796            })
797            .connect(&resolved.database_url)
798            .await
799            .map_err(|err| classify_pool_connect_error(resolved.set_role.as_deref(), err))?;
800
801        // Cache it (write lock).
802        let superseded = {
803            let mut cache = self.pool_cache.write().await;
804            cache.insert(
805                cache_key,
806                CachedPool {
807                    resource_version,
808                    secret_fingerprint,
809                    token_expires_at: resolved.token_expires_at,
810                    pool: pool.clone(),
811                },
812            )
813        };
814        close_unreachable_pool(superseded);
815
816        Ok(pool)
817    }
818
819    /// Fetch a single string value from a Kubernetes Secret.
820    ///
821    /// Used to resolve role passwords from Secret references at reconcile time.
822    pub async fn fetch_secret_value(
823        &self,
824        namespace: &str,
825        secret_name: &str,
826        secret_key: &str,
827    ) -> Result<String, ContextError> {
828        let secrets_api: kube::Api<k8s_openapi::api::core::v1::Secret> =
829            kube::Api::namespaced(self.kube_client.clone(), namespace);
830
831        let secret =
832            secrets_api
833                .get(secret_name)
834                .await
835                .map_err(|err| ContextError::SecretFetch {
836                    name: secret_name.to_string(),
837                    namespace: namespace.to_string(),
838                    source: err,
839                })?;
840
841        let data = secret.data.ok_or_else(|| ContextError::SecretMissing {
842            name: secret_name.to_string(),
843            key: secret_key.to_string(),
844        })?;
845
846        let value_bytes = data
847            .get(secret_key)
848            .ok_or_else(|| ContextError::SecretMissing {
849                name: secret_name.to_string(),
850                key: secret_key.to_string(),
851            })?;
852
853        String::from_utf8(value_bytes.0.clone()).map_err(|_| ContextError::SecretMissing {
854            name: secret_name.to_string(),
855            key: secret_key.to_string(),
856        })
857    }
858
859    /// Remove a cached pool (e.g. when secret changes or CR is deleted).
860    pub async fn evict_pool(&self, namespace: &str, connection: &ConnectionSpec) {
861        let cache_key = connection.cache_key(namespace);
862        let evicted = {
863            let mut cache = self.pool_cache.write().await;
864            cache.remove(&cache_key)
865        };
866        close_unreachable_pool(evicted);
867    }
868}
869
870/// Connection-pool sizing and lifetime shared by every managed database.
871fn pool_options() -> PgPoolOptions {
872    PgPoolOptions::new()
873        .max_connections(POOL_MAX_CONNECTIONS)
874        .min_connections(POOL_MIN_CONNECTIONS)
875        .acquire_timeout(Duration::from_secs(POOL_ACQUIRE_TIMEOUT_SECS))
876        .idle_timeout(Duration::from_secs(POOL_IDLE_TIMEOUT_SECS))
877        .max_lifetime(Duration::from_secs(POOL_MAX_LIFETIME_SECS))
878}
879
880/// Close a pool that is no longer reachable from the cache.
881///
882/// Dropping a `PgPool` only marks it closed: the backend connections are torn
883/// down when the last handle drops the socket, which PostgreSQL logs as an
884/// unexpected EOF. Closing explicitly sends a proper termination once in-flight
885/// reconciles hand their connections back. The close is spawned because it
886/// waits for those connections, and no caller should block on an eviction.
887fn close_unreachable_pool(cached: Option<CachedPool>) {
888    let Some(cached) = cached else {
889        return;
890    };
891    tokio::spawn(async move {
892        cached.pool.close().await;
893    });
894}
895
896fn database_target_from_url(database_url: &str) -> Result<(String, u16, String), String> {
897    let options = PgConnectOptions::from_str(database_url).map_err(|source| source.to_string())?;
898    Ok((
899        options.get_host().to_ascii_lowercase(),
900        options.get_port(),
901        options
902            .get_database()
903            .unwrap_or_else(|| options.get_username())
904            .to_string(),
905    ))
906}
907
908fn database_target_fingerprint(host: &str, port: u16, database: &str) -> String {
909    let digest = Sha256::digest(format!("{}\0{port}\0{database}", host.to_ascii_lowercase()));
910    let mut fingerprint = String::with_capacity(7 + digest.len() * 2);
911    fingerprint.push_str("sha256:");
912    for byte in digest {
913        write!(&mut fingerprint, "{byte:02x}")
914            .expect("writing a database target fingerprint cannot fail");
915    }
916    fingerprint
917}
918
919/// Errors from operator context operations.
920#[derive(Debug, thiserror::Error)]
921pub enum ContextError {
922    #[error("failed to fetch Secret {namespace}/{name}: {source}")]
923    SecretFetch {
924        name: String,
925        namespace: String,
926        source: kube::Error,
927    },
928
929    #[error("Secret \"{name}\" does not contain key \"{key}\"")]
930    SecretMissing { name: String, key: String },
931
932    #[error("failed to connect to database: {source}")]
933    DatabaseConnect { source: sqlx::Error },
934
935    #[error("failed to apply SET ROLE \"{role}\" on pooled connection: {source}")]
936    SetRoleFailed { role: String, source: sqlx::Error },
937
938    #[error("connection param \"{field}\" resolved to an empty or whitespace-only value")]
939    EmptyResolvedValue { field: String },
940
941    #[error("connection URL is invalid: {detail}")]
942    InvalidDatabaseUrl { detail: String },
943
944    #[error("connection port resolved to invalid value \"{value}\"")]
945    InvalidResolvedPort { value: String },
946
947    #[error(
948        "connection param sslMode resolved to invalid value \"{value}\" (expected one of: disable, allow, prefer, require, verify-ca, verify-full)"
949    )]
950    InvalidResolvedSslMode { value: String },
951
952    #[error("failed to fetch GCP auth token from {endpoint}: {source}")]
953    GcpAuthHttp {
954        endpoint: &'static str,
955        source: reqwest::Error,
956    },
957
958    #[error("GCP auth token endpoint {endpoint} returned HTTP {status}: {body}")]
959    GcpAuthRejected {
960        endpoint: String,
961        status: u16,
962        body: String,
963    },
964
965    #[error("GCP auth token response was invalid: {detail}")]
966    GcpAuthInvalidResponse { detail: String },
967}
968
969impl ContextError {
970    /// Returns true when a Secret fetch failed due to a non-transient client-side API error.
971    pub fn is_secret_fetch_non_transient(&self) -> bool {
972        matches!(
973            self,
974            ContextError::SecretFetch {
975                source: kube::Error::Api(response),
976                ..
977            } if (400..500).contains(&response.code) && response.code != 429
978        )
979    }
980
981    pub fn is_gcp_auth_non_transient(&self) -> bool {
982        matches!(
983            self,
984            ContextError::GcpAuthRejected { status, .. }
985                if (400..500).contains(status) && *status != 429
986        ) || matches!(self, ContextError::GcpAuthInvalidResponse { .. })
987    }
988}
989
990#[cfg(test)]
991mod tests {
992    use super::*;
993
994    /// Pools are cached for the operator's lifetime, so the only thing that
995    /// ever closes an unused connection is the idle reaper. If the reap window
996    /// reaches the requeue interval, every reconcile re-touches the pool before
997    /// the reaper runs and the operator sits on open backends against every
998    /// managed database forever — which is what sqlx's 10-minute default did
999    /// against the 5-minute default requeue.
1000    #[test]
1001    fn idle_connections_are_reaped_between_reconciles() {
1002        let options = pool_options();
1003
1004        assert_eq!(
1005            options.get_min_connections(),
1006            0,
1007            "a floor above zero would hold connections open on an idle database"
1008        );
1009
1010        let idle_timeout = options
1011            .get_idle_timeout()
1012            .expect("an unset idle timeout never reaps");
1013        assert!(
1014            idle_timeout < Duration::from_secs(crate::reconciler::DEFAULT_REQUEUE_SECS),
1015            "idle timeout {idle_timeout:?} must drain within the {}s requeue interval",
1016            crate::reconciler::DEFAULT_REQUEUE_SECS
1017        );
1018
1019        assert!(
1020            options.get_max_lifetime().is_some(),
1021            "connections need an age ceiling independent of activity"
1022        );
1023    }
1024
1025    #[test]
1026    fn database_target_fingerprint_excludes_credentials_and_options() {
1027        let first = database_target_from_url(
1028            "postgresql://alice:first@DB.EXAMPLE:6432/inventory?sslmode=require",
1029        )
1030        .expect("first URL");
1031        let second = database_target_from_url(
1032            "postgresql://bob:second@db.example:6432/inventory?application_name=test",
1033        )
1034        .expect("second URL");
1035
1036        assert_eq!(
1037            database_target_fingerprint(&first.0, first.1, &first.2),
1038            database_target_fingerprint(&second.0, second.1, &second.2)
1039        );
1040    }
1041
1042    #[test]
1043    fn database_target_fingerprint_changes_on_retarget() {
1044        let original = database_target_fingerprint("db.example", 5432, "inventory");
1045        assert_ne!(
1046            original,
1047            database_target_fingerprint("db.example", 5432, "billing")
1048        );
1049        assert_ne!(
1050            original,
1051            database_target_fingerprint("other.example", 5432, "inventory")
1052        );
1053    }
1054
1055    #[test]
1056    fn build_set_role_stmt_quotes_identifier() {
1057        assert_eq!(
1058            build_set_role_stmt("cloudsqlsuperuser"),
1059            "SET ROLE \"cloudsqlsuperuser\"",
1060        );
1061    }
1062
1063    #[test]
1064    fn classify_pool_connect_error_surfaces_set_role_failure_with_role() {
1065        let raw = wrap_set_role_failure(
1066            "cloudsqlsuperuser",
1067            sqlx::Error::Protocol("permission denied".to_string()),
1068        );
1069        let classified = classify_pool_connect_error(Some("cloudsqlsuperuser"), raw);
1070        assert!(matches!(
1071            classified,
1072            ContextError::SetRoleFailed { ref role, .. } if role == "cloudsqlsuperuser"
1073        ));
1074    }
1075
1076    #[test]
1077    fn classify_pool_connect_error_passes_through_unrelated_errors() {
1078        let err = sqlx::Error::PoolTimedOut;
1079        let classified = classify_pool_connect_error(Some("any_role"), err);
1080        assert!(matches!(classified, ContextError::DatabaseConnect { .. }));
1081    }
1082
1083    #[test]
1084    fn classify_pool_connect_error_without_set_role_is_database_connect() {
1085        // Even a Protocol error with our marker shouldn't be classified as
1086        // SetRoleFailed when no role was configured for this pool.
1087        let raw = wrap_set_role_failure("ghost", sqlx::Error::Protocol("oops".to_string()));
1088        let classified = classify_pool_connect_error(None, raw);
1089        assert!(matches!(classified, ContextError::DatabaseConnect { .. }));
1090    }
1091
1092    #[test]
1093    fn build_set_role_stmt_doubles_embedded_quote() {
1094        // CRD validation rejects identifiers containing `"`. This test pins
1095        // the defensive quoting in the connection-pool path anyway, so a
1096        // future relaxation of the validator can't silently allow injection.
1097        assert_eq!(build_set_role_stmt("a\"b"), "SET ROLE \"a\"\"b\"",);
1098    }
1099
1100    #[test]
1101    fn pool_cache_key_format() {
1102        // Verify the cache key format is "namespace/secret-name/secret-key"
1103        let key = format!("{}/{}/{}", "prod", "pg-credentials", "DATABASE_URL");
1104        assert_eq!(key, "prod/pg-credentials/DATABASE_URL");
1105    }
1106
1107    #[test]
1108    fn secret_fetch_not_found_is_non_transient() {
1109        let error = ContextError::SecretFetch {
1110            name: "db-credentials".into(),
1111            namespace: "default".into(),
1112            source: kube::Error::Api(
1113                kube::core::Status::failure("secrets \"db-credentials\" not found", "NotFound")
1114                    .with_code(404)
1115                    .boxed(),
1116            ),
1117        };
1118
1119        assert!(error.is_secret_fetch_non_transient());
1120    }
1121
1122    #[test]
1123    fn secret_fetch_forbidden_is_non_transient() {
1124        let error = ContextError::SecretFetch {
1125            name: "db-credentials".into(),
1126            namespace: "default".into(),
1127            source: kube::Error::Api(
1128                kube::core::Status::failure("forbidden", "Forbidden")
1129                    .with_code(403)
1130                    .boxed(),
1131            ),
1132        };
1133
1134        assert!(error.is_secret_fetch_non_transient());
1135    }
1136
1137    #[test]
1138    fn secret_fetch_server_error_remains_transient() {
1139        let error = ContextError::SecretFetch {
1140            name: "db-credentials".into(),
1141            namespace: "default".into(),
1142            source: kube::Error::Api(
1143                kube::core::Status::failure("internal error", "InternalError")
1144                    .with_code(500)
1145                    .boxed(),
1146            ),
1147        };
1148
1149        assert!(!error.is_secret_fetch_non_transient());
1150    }
1151
1152    #[test]
1153    fn gcp_auth_client_error_is_non_transient() {
1154        let error = ContextError::GcpAuthRejected {
1155            endpoint: "metadata".into(),
1156            status: 403,
1157            body: "forbidden".into(),
1158        };
1159
1160        assert!(error.is_gcp_auth_non_transient());
1161    }
1162
1163    #[test]
1164    fn gcp_auth_rate_limit_remains_transient() {
1165        let error = ContextError::GcpAuthRejected {
1166            endpoint: "metadata".into(),
1167            status: 429,
1168            body: "rate limited".into(),
1169        };
1170
1171        assert!(!error.is_gcp_auth_non_transient());
1172    }
1173
1174    #[test]
1175    fn token_expiry_uses_five_minute_refresh_skew() {
1176        let now = SystemTime::UNIX_EPOCH + Duration::from_secs(1_000);
1177        assert!(token_expires_after_skew(
1178            Some(now + Duration::from_secs(GCP_TOKEN_CACHE_SKEW_SECS + 1)),
1179            now
1180        ));
1181        assert!(!token_expires_after_skew(
1182            Some(now + Duration::from_secs(GCP_TOKEN_CACHE_SKEW_SECS)),
1183            now
1184        ));
1185    }
1186
1187    #[test]
1188    fn parse_google_expire_time_accepts_rfc3339() {
1189        let parsed =
1190            parse_google_expire_time("2026-05-14T02:30:00Z").expect("expireTime should parse");
1191        assert_eq!(
1192            parsed
1193                .duration_since(SystemTime::UNIX_EPOCH)
1194                .unwrap()
1195                .as_secs(),
1196            1_778_725_800
1197        );
1198    }
1199
1200    #[test]
1201    fn truncate_for_error_keeps_utf8_boundary() {
1202        let body = "é".repeat(300);
1203        let truncated = truncate_for_error(body);
1204
1205        assert!(truncated.ends_with("..."));
1206        assert!(truncated.is_char_boundary(truncated.len() - 3));
1207    }
1208
1209    #[tokio::test]
1210    async fn try_lock_database_acquires_when_free() {
1211        let locks: Arc<Mutex<HashMap<String, ()>>> = Arc::new(Mutex::new(HashMap::new()));
1212        let ctx = OperatorContextLockHelper {
1213            database_locks: locks,
1214        };
1215        let guard = ctx.try_lock("db-a").await;
1216        assert!(guard.is_some(), "should acquire lock on free database");
1217    }
1218
1219    #[tokio::test]
1220    async fn try_lock_database_contention_returns_none() {
1221        let locks: Arc<Mutex<HashMap<String, ()>>> = Arc::new(Mutex::new(HashMap::new()));
1222        let ctx = OperatorContextLockHelper {
1223            database_locks: locks,
1224        };
1225
1226        let _guard1 = ctx
1227            .try_lock("db-a")
1228            .await
1229            .expect("first lock should succeed");
1230        let guard2 = ctx.try_lock("db-a").await;
1231        assert!(guard2.is_none(), "second lock on same database should fail");
1232    }
1233
1234    #[tokio::test]
1235    async fn try_lock_database_different_databases_independent() {
1236        let locks: Arc<Mutex<HashMap<String, ()>>> = Arc::new(Mutex::new(HashMap::new()));
1237        let ctx = OperatorContextLockHelper {
1238            database_locks: locks,
1239        };
1240
1241        let guard_a = ctx.try_lock("db-a").await;
1242        let guard_b = ctx.try_lock("db-b").await;
1243        assert!(guard_a.is_some(), "lock on db-a should succeed");
1244        assert!(
1245            guard_b.is_some(),
1246            "lock on db-b should succeed (different database)"
1247        );
1248    }
1249
1250    #[tokio::test]
1251    async fn try_lock_database_released_after_drop() {
1252        let locks: Arc<Mutex<HashMap<String, ()>>> = Arc::new(Mutex::new(HashMap::new()));
1253        let ctx = OperatorContextLockHelper {
1254            database_locks: Arc::clone(&locks),
1255        };
1256
1257        {
1258            let _guard = ctx.try_lock("db-a").await.expect("should acquire");
1259            // guard is dropped here
1260        }
1261
1262        // After drop, should be able to acquire again.
1263        let guard2 = ctx.try_lock("db-a").await;
1264        assert!(
1265            guard2.is_some(),
1266            "should re-acquire after previous guard dropped"
1267        );
1268    }
1269
1270    #[tokio::test]
1271    async fn try_lock_database_concurrent_contention() {
1272        let locks: Arc<Mutex<HashMap<String, ()>>> = Arc::new(Mutex::new(HashMap::new()));
1273
1274        // Simulate two concurrent reconciles for the same database.
1275        let locks1 = Arc::clone(&locks);
1276        let locks2 = Arc::clone(&locks);
1277
1278        let handle1 = tokio::spawn(async move {
1279            let ctx = OperatorContextLockHelper {
1280                database_locks: locks1,
1281            };
1282            let guard = ctx.try_lock("shared-db").await;
1283            if guard.is_some() {
1284                // Hold the lock briefly.
1285                tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1286            }
1287            guard.is_some()
1288        });
1289
1290        // Small delay so handle1 is likely first.
1291        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
1292
1293        let handle2 = tokio::spawn(async move {
1294            let ctx = OperatorContextLockHelper {
1295                database_locks: locks2,
1296            };
1297            let guard = ctx.try_lock("shared-db").await;
1298            guard.is_some()
1299        });
1300
1301        let (r1, r2) = tokio::join!(handle1, handle2);
1302        let acquired1 = r1.unwrap();
1303        let acquired2 = r2.unwrap();
1304
1305        // Exactly one should succeed.
1306        assert!(
1307            acquired1 ^ acquired2,
1308            "exactly one of two concurrent locks should succeed: got ({acquired1}, {acquired2})"
1309        );
1310    }
1311
1312    /// Helper to test locking without a real kube client.
1313    struct OperatorContextLockHelper {
1314        database_locks: Arc<Mutex<HashMap<String, ()>>>,
1315    }
1316
1317    impl OperatorContextLockHelper {
1318        async fn try_lock(&self, database_identity: &str) -> Option<DatabaseLockGuard> {
1319            let mut locks = self.database_locks.lock().await;
1320            if locks.contains_key(database_identity) {
1321                return None;
1322            }
1323            locks.insert(database_identity.to_string(), ());
1324            Some(DatabaseLockGuard {
1325                key: database_identity.to_string(),
1326                locks: Arc::clone(&self.database_locks),
1327            })
1328        }
1329    }
1330
1331    #[tokio::test]
1332    async fn try_lock_database_high_concurrency_same_db() {
1333        // Spawn many tasks all racing to lock the same database.
1334        let locks: Arc<Mutex<HashMap<String, ()>>> = Arc::new(Mutex::new(HashMap::new()));
1335        let concurrency = 50;
1336        let acquired_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
1337        let barrier = Arc::new(tokio::sync::Barrier::new(concurrency));
1338
1339        let mut handles = Vec::with_capacity(concurrency);
1340        for _ in 0..concurrency {
1341            let locks_clone = Arc::clone(&locks);
1342            let count = Arc::clone(&acquired_count);
1343            let bar = Arc::clone(&barrier);
1344            handles.push(tokio::spawn(async move {
1345                // Synchronize start so all tasks race at the same instant.
1346                bar.wait().await;
1347                let ctx = OperatorContextLockHelper {
1348                    database_locks: locks_clone,
1349                };
1350                let guard = ctx.try_lock("contested-db").await;
1351                if guard.is_some() {
1352                    count.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1353                    // Hold lock briefly to let other tasks observe contention.
1354                    tokio::time::sleep(std::time::Duration::from_millis(10)).await;
1355                }
1356            }));
1357        }
1358
1359        for h in handles {
1360            h.await.unwrap();
1361        }
1362
1363        // Exactly one task should have acquired the lock.
1364        let total = acquired_count.load(std::sync::atomic::Ordering::SeqCst);
1365        assert_eq!(
1366            total, 1,
1367            "exactly one of {concurrency} concurrent tasks should acquire the lock, got {total}"
1368        );
1369    }
1370
1371    #[tokio::test]
1372    async fn try_lock_database_high_concurrency_different_dbs() {
1373        // Many tasks each locking a different database — all should succeed.
1374        let locks: Arc<Mutex<HashMap<String, ()>>> = Arc::new(Mutex::new(HashMap::new()));
1375        let concurrency = 50;
1376        let acquired_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
1377        let barrier = Arc::new(tokio::sync::Barrier::new(concurrency));
1378
1379        let mut handles = Vec::with_capacity(concurrency);
1380        for i in 0..concurrency {
1381            let locks_clone = Arc::clone(&locks);
1382            let count = Arc::clone(&acquired_count);
1383            let bar = Arc::clone(&barrier);
1384            handles.push(tokio::spawn(async move {
1385                bar.wait().await;
1386                let ctx = OperatorContextLockHelper {
1387                    database_locks: locks_clone,
1388                };
1389                let db_name = format!("db-{i}");
1390                let guard = ctx.try_lock(&db_name).await;
1391                if guard.is_some() {
1392                    count.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1393                    tokio::time::sleep(std::time::Duration::from_millis(5)).await;
1394                }
1395            }));
1396        }
1397
1398        for h in handles {
1399            h.await.unwrap();
1400        }
1401
1402        let total = acquired_count.load(std::sync::atomic::Ordering::SeqCst);
1403        assert_eq!(
1404            total, concurrency,
1405            "all {concurrency} tasks locking different dbs should succeed, got {total}"
1406        );
1407    }
1408
1409    #[tokio::test]
1410    async fn try_lock_database_acquire_release_cycle_under_contention() {
1411        // Repeatedly acquire and release the same database lock from many tasks.
1412        // Each task attempts the lock in a loop until it succeeds, simulating
1413        // the requeue-after-contention pattern used in the reconciler.
1414        let locks: Arc<Mutex<HashMap<String, ()>>> = Arc::new(Mutex::new(HashMap::new()));
1415        let concurrency = 20;
1416        let success_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
1417        let barrier = Arc::new(tokio::sync::Barrier::new(concurrency));
1418
1419        let mut handles = Vec::with_capacity(concurrency);
1420        for _ in 0..concurrency {
1421            let locks_clone = Arc::clone(&locks);
1422            let count = Arc::clone(&success_count);
1423            let bar = Arc::clone(&barrier);
1424            handles.push(tokio::spawn(async move {
1425                bar.wait().await;
1426                // Retry up to 100 times with a small sleep between attempts,
1427                // simulating the jittered requeue pattern.
1428                for _ in 0..100 {
1429                    let ctx = OperatorContextLockHelper {
1430                        database_locks: Arc::clone(&locks_clone),
1431                    };
1432                    if let Some(_guard) = ctx.try_lock("shared-db").await {
1433                        count.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1434                        // Brief simulated work, then guard drops (releasing lock).
1435                        tokio::time::sleep(std::time::Duration::from_millis(1)).await;
1436                        return;
1437                    }
1438                    tokio::time::sleep(std::time::Duration::from_millis(1)).await;
1439                }
1440                // Should not reach here in practice — fail the test if we do.
1441                panic!("task failed to acquire lock after 100 retries");
1442            }));
1443        }
1444
1445        for h in handles {
1446            h.await.unwrap();
1447        }
1448
1449        let total = success_count.load(std::sync::atomic::Ordering::SeqCst);
1450        assert_eq!(
1451            total, concurrency,
1452            "all {concurrency} tasks should eventually acquire the lock"
1453        );
1454    }
1455}