graph, store: Add periodic probe for unavailable database recovery

When a database shard is marked unavailable due to connection timeouts,
StateTracker now allows one probe request through every few seconds
(configurable via GRAPH_STORE_CONNECTION_UNAVAILABLE_RETRY, default 2s)
to check if the database has recovered. This prevents the recovery
deadlock where mark_available() only fires on successful connection,
but no connections are attempted when unavailable.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: lutter

docs/environment-variables.md

@@ -216,6 +216,11 @@ those.
   decisions. Set to `true` to turn simulation on, defaults to `false`
 - `GRAPH_STORE_CONNECTION_TIMEOUT`: How long to wait to connect to a
   database before assuming the database is down in ms. Defaults to 5000ms.
+- `GRAPH_STORE_CONNECTION_UNAVAILABLE_RETRY`: When a database shard is marked
+  unavailable due to connection timeouts, this controls how often to allow a
+  single probe request through to check if the database has recovered. Only one
+  request per interval will attempt a connection; all others fail instantly.
+  Value is in seconds and defaults to 2s.
 - `EXPERIMENTAL_SUBGRAPH_VERSION_SWITCHING_MODE`: default is `instant`, set
   to `synced` to only switch a named subgraph to a new deployment once it
   has synced, making the new deployment the "Pending" version.

graph/src/env/store.rs

@@ -171,6 +171,12 @@ pub struct EnvVarsStore {
     /// Set to 0 to always validate (previous behavior).
     /// Set by `GRAPH_STORE_CONNECTION_VALIDATION_IDLE_SECS`. Default is 30 seconds.
     pub connection_validation_idle_secs: Duration,
+    /// When a database shard is marked unavailable due to connection timeouts,
+    /// this controls how often to allow a single probe request through to check
+    /// if the database has recovered. Only one request per interval will attempt
+    /// a connection; all others fail instantly with DatabaseUnavailable.
+    /// Set by `GRAPH_STORE_CONNECTION_UNAVAILABLE_RETRY`. Default is 2 seconds.
+    pub connection_unavailable_retry: Duration,
 }
 
 // This does not print any values avoid accidentally leaking any sensitive env vars
@@ -234,6 +240,9 @@ impl TryFrom<InnerStore> for EnvVarsStore {
             disable_call_cache: x.disable_call_cache,
             disable_chain_head_ptr_cache: x.disable_chain_head_ptr_cache,
             connection_validation_idle_secs: Duration::from_secs(x.connection_validation_idle_secs),
+            connection_unavailable_retry: Duration::from_secs(
+                x.connection_unavailable_retry_in_secs,
+            ),
         };
         if let Some(timeout) = vars.batch_timeout {
             if timeout < 2 * vars.batch_target_duration {
@@ -345,6 +354,8 @@ pub struct InnerStore {
     disable_chain_head_ptr_cache: bool,
     #[envconfig(from = "GRAPH_STORE_CONNECTION_VALIDATION_IDLE_SECS", default = "30")]
     connection_validation_idle_secs: u64,
+    #[envconfig(from = "GRAPH_STORE_CONNECTION_UNAVAILABLE_RETRY", default = "2")]
+    connection_unavailable_retry_in_secs: u64,
 }
 
 #[derive(Clone, Copy, Debug)]

store/postgres/src/pool/manager.rs

@@ -20,9 +20,10 @@ use graph::slog::Logger;
 
 use std::collections::HashMap;
 use std::sync::atomic::AtomicBool;
+use std::sync::atomic::AtomicU64;
 use std::sync::atomic::Ordering;
 use std::sync::Arc;
-use std::time::{Duration, Instant};
+use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
 
 use crate::pool::AsyncPool;
 
@@ -141,6 +142,9 @@ pub(super) struct StateTracker {
     logger: Logger,
     available: Arc<AtomicBool>,
     ignore_timeout: Arc<AtomicBool>,
+    /// Timestamp (as millis since UNIX_EPOCH) when we can next probe.
+    /// 0 means available/no limit.
+    next_probe_at: Arc<AtomicU64>,
 }
 
 impl StateTracker {
@@ -149,14 +153,16 @@ impl StateTracker {
             logger,
             available: Arc::new(AtomicBool::new(true)),
             ignore_timeout: Arc::new(AtomicBool::new(false)),
+            next_probe_at: Arc::new(AtomicU64::new(0)),
         }
     }
 
     pub(super) fn mark_available(&self) {
         if !self.is_available() {
-            info!(self.logger, "Conection checkout"; "event" => "available");
+            info!(self.logger, "Connection checkout"; "event" => "available");
         }
         self.available.store(true, Ordering::Relaxed);
+        self.next_probe_at.store(0, Ordering::Relaxed);
     }
 
     pub(super) fn mark_unavailable(&self, waited: Duration) {
@@ -171,10 +177,47 @@ impl StateTracker {
             }
         }
         self.available.store(false, Ordering::Relaxed);
+
+        // Set next probe time
+        let retry_interval = ENV_VARS.store.connection_unavailable_retry.as_millis() as u64;
+        let next_probe = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap()
+            .as_millis() as u64
+            + retry_interval;
+        self.next_probe_at.store(next_probe, Ordering::Relaxed);
     }
 
     pub(super) fn is_available(&self) -> bool {
-        AtomicBool::load(&self.available, Ordering::Relaxed)
+        if AtomicBool::load(&self.available, Ordering::Relaxed) {
+            return true;
+        }
+
+        // Allow one probe through every `connection_unavailable_retry` interval
+        let next_probe = AtomicU64::load(&self.next_probe_at, Ordering::Relaxed);
+        let now_millis = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap()
+            .as_millis() as u64;
+
+        if now_millis >= next_probe {
+            // Try to claim this probe slot with CAS
+            let retry_interval = ENV_VARS.store.connection_unavailable_retry.as_millis() as u64;
+            if self
+                .next_probe_at
+                .compare_exchange(
+                    next_probe,
+                    now_millis + retry_interval,
+                    Ordering::Relaxed,
+                    Ordering::Relaxed,
+                )
+                .is_ok()
+            {
+                // We claimed the probe - allow this request through
+                return true;
+            }
+        }
+        false
     }
 
     pub(super) fn timeout_is_ignored(&self) -> bool {