diff --git a/crates/sui-indexer-alt-graphql/src/api/subscription/events.rs b/crates/sui-indexer-alt-graphql/src/api/subscription/events.rs index d35745e5ea41..a990a9ddc995 100644 --- a/crates/sui-indexer-alt-graphql/src/api/subscription/events.rs +++ b/crates/sui-indexer-alt-graphql/src/api/subscription/events.rs @@ -32,8 +32,6 @@ use crate::task::streaming::StreamedCaches; use super::scan_then_live::Subscribable; impl Subscribable for Event { - const SUBSCRIPTION_TYPE: &'static str = "events"; - type Item = Self; type Cursor = CEvent; type Filter = EventFilter; diff --git a/crates/sui-indexer-alt-graphql/src/api/subscription/mod.rs b/crates/sui-indexer-alt-graphql/src/api/subscription/mod.rs index 069c8c5839b3..2af8ee505267 100644 --- a/crates/sui-indexer-alt-graphql/src/api/subscription/mod.rs +++ b/crates/sui-indexer-alt-graphql/src/api/subscription/mod.rs @@ -26,9 +26,12 @@ use crate::config::Limits; use crate::config::SubscriptionConfig; use crate::error::RpcError; use crate::error::bad_user_input; +use crate::error::upcast; use crate::scope::Scope; use crate::task::streaming::StreamedCaches; +use crate::task::streaming::SubscriberLimit; use crate::task::streaming::SubscriptionBroadcast; +use crate::task::streaming::SubscriptionLifecycleGuard; use crate::task::watermark::Watermarks; mod events; @@ -79,13 +82,14 @@ impl Subscription { (a, b) => a.or(b), }; reject_if_start_too_far_ahead(start_from, broadcast, config)?; + let guard = admit_subscription(ctx, broadcast, "checkpoints")?; let caches = caches.clone(); let resolver_limits = limits.package_resolver(); let stream = broadcast .clone() - .subscribe(start_from, fetcher.clone(), config); + .subscribe(start_from, fetcher.clone(), config, guard); Ok(stream.map(move |item| { item.map(|processed| { @@ -148,6 +152,7 @@ impl Subscription { (a, b) => a.or(b), }; reject_if_start_too_far_ahead(start_from, broadcast, config)?; + let guard = admit_subscription(ctx, broadcast, "transactions")?; Ok(subscribe::( reader.clone(), @@ -159,6 +164,7 @@ impl Subscription { after, after_checkpoint, config.clone(), + guard, )) } @@ -202,6 +208,7 @@ impl Subscription { (a, b) => a.or(b), }; reject_if_start_too_far_ahead(start_from, broadcast, config)?; + let guard = admit_subscription(ctx, broadcast, "events")?; Ok(subscribe::( reader.clone(), @@ -213,10 +220,24 @@ impl Subscription { after, after_checkpoint, config.clone(), + guard, )) } } +/// Admit a new subscription of `subscription_type` by claiming a concurrency slot, or return an +/// at-capacity error refusing it. The returned guard holds the slot and the per-subscriber metric +/// handles for the subscription's lifetime; it is moved into the subscription's stream driver. +fn admit_subscription( + ctx: &Context<'_>, + broadcast: &SubscriptionBroadcast, + subscription_type: &'static str, +) -> Result> { + let subscriber_limit: &SubscriberLimit = ctx.data()?; + SubscriptionLifecycleGuard::new(subscription_type, broadcast.metrics(), subscriber_limit) + .map_err(upcast) +} + /// Reject a start point sitting more than `max_ahead` checkpoints past the chain tip. There is /// nothing to backfill ahead of the tip, so such a request would only wait for the chain to reach /// it, and a far-future one would hold the connection open indefinitely. diff --git a/crates/sui-indexer-alt-graphql/src/api/subscription/scan_then_live.rs b/crates/sui-indexer-alt-graphql/src/api/subscription/scan_then_live.rs index ef4857a1ed25..928c23081e93 100644 --- a/crates/sui-indexer-alt-graphql/src/api/subscription/scan_then_live.rs +++ b/crates/sui-indexer-alt-graphql/src/api/subscription/scan_then_live.rs @@ -57,9 +57,6 @@ const BACKFILL_POLL_INTERVAL: Duration = Duration::from_millis(100); /// GraphQL type (transactions, events) supplies its cursor, filter, and matching, and the free /// functions here drive the shared machinery over it. pub(super) trait Subscribable { - /// The subscription kind, used to label this feed's aggregate subscriber metrics. - const SUBSCRIPTION_TYPE: &'static str; - /// The GraphQL node delivered in each edge. type Item: OutputType + Send + 'static; /// The opaque cursor minted for resumption. @@ -130,6 +127,7 @@ pub(super) fn subscribe( after: Option, after_checkpoint: Option, config: SubscriptionConfig, + guard: SubscriptionLifecycleGuard, ) -> impl Stream, RpcError>> { // Size the backfill scan page to the resolve concurrency. Scans are sequential (each needs the // previous page's cursor), so feeding one window of `n` concurrent resolutions takes ceil(n / @@ -142,7 +140,6 @@ pub(super) fn subscribe( let handoff_threshold = config.broadcast_buffer as u64 / 2; stream! { - let guard = SubscriptionLifecycleGuard::new(S::SUBSCRIPTION_TYPE, broadcast.metrics()); let mut pending_receiver = None; let mut handoff: Option = None; let mut last_checkpoint: Option = None; diff --git a/crates/sui-indexer-alt-graphql/src/api/subscription/transactions.rs b/crates/sui-indexer-alt-graphql/src/api/subscription/transactions.rs index d09ad89502b6..d95c2fd3b6d6 100644 --- a/crates/sui-indexer-alt-graphql/src/api/subscription/transactions.rs +++ b/crates/sui-indexer-alt-graphql/src/api/subscription/transactions.rs @@ -30,8 +30,6 @@ use crate::task::streaming::StreamedCaches; use super::scan_then_live::Subscribable; impl Subscribable for Transaction { - const SUBSCRIPTION_TYPE: &'static str = "transactions"; - type Item = Self; type Cursor = CTransaction; type Filter = TransactionFilter; diff --git a/crates/sui-indexer-alt-graphql/src/config.rs b/crates/sui-indexer-alt-graphql/src/config.rs index d266e0139f88..fc599165d707 100644 --- a/crates/sui-indexer-alt-graphql/src/config.rs +++ b/crates/sui-indexer-alt-graphql/src/config.rs @@ -299,6 +299,10 @@ pub struct SubscriptionConfig { /// is rejected instead. Raising this admits starts further past the tip, at the cost of /// connections parked waiting longer. pub max_start_checkpoints_ahead_of_tip: u64, + + /// Maximum number of concurrent subscriptions the server admits at once. A subscription opened + /// while this many are already active is rejected so it can retry later. + pub max_subscribers: usize, } impl Default for SubscriptionConfig { @@ -313,6 +317,7 @@ impl Default for SubscriptionConfig { per_subscriber_max_output_nodes_per_second: 1_000_000, // About a minute at the average checkpoint rate. max_start_checkpoints_ahead_of_tip: 300, + max_subscribers: 1024, } } } @@ -328,6 +333,7 @@ pub struct SubscriptionLayer { pub max_concurrent_resolutions: Option, pub per_subscriber_max_output_nodes_per_second: Option, pub max_start_checkpoints_ahead_of_tip: Option, + pub max_subscribers: Option, } impl SubscriptionLayer { @@ -355,6 +361,7 @@ impl SubscriptionLayer { max_start_checkpoints_ahead_of_tip: self .max_start_checkpoints_ahead_of_tip .unwrap_or(base.max_start_checkpoints_ahead_of_tip), + max_subscribers: self.max_subscribers.unwrap_or(base.max_subscribers), } } } @@ -680,6 +687,7 @@ impl From for SubscriptionLayer { value.per_subscriber_max_output_nodes_per_second, ), max_start_checkpoints_ahead_of_tip: Some(value.max_start_checkpoints_ahead_of_tip), + max_subscribers: Some(value.max_subscribers), } } } diff --git a/crates/sui-indexer-alt-graphql/src/lib.rs b/crates/sui-indexer-alt-graphql/src/lib.rs index 6cf421268fba..b7ffdd091996 100644 --- a/crates/sui-indexer-alt-graphql/src/lib.rs +++ b/crates/sui-indexer-alt-graphql/src/lib.rs @@ -61,6 +61,8 @@ use task::streaming::StreamedObjectStore; use task::streaming::StreamedTransactionStore; use task::streaming::StreamingPackageStore; #[cfg(feature = "staging")] +use task::streaming::SubscriberLimit; +#[cfg(feature = "staging")] use task::streaming::SubscriptionBroadcast; use task::streaming::SubscriptionReadiness; use task::watermark::WatermarkTask; @@ -506,6 +508,8 @@ pub async fn start_rpc( readiness, )) = streaming_setup { + #[cfg(feature = "staging")] + let max_subscribers = config.subscription.max_subscribers; rpc = rpc.data(caches).data(config.subscription); let s_stream = stream_task.run(); let s_eviction = eviction_task.run(); @@ -525,7 +529,8 @@ pub async fn start_rpc( )); rpc = rpc .data(subscription_broadcast) - .data(subscription_watermarks_rx); + .data(subscription_watermarks_rx) + .data(SubscriberLimit::new(max_subscribers)); } Some((s_stream, s_eviction)) } else { diff --git a/crates/sui-indexer-alt-graphql/src/metrics/subscription.rs b/crates/sui-indexer-alt-graphql/src/metrics/subscription.rs index 65ab8c4ae361..440112d4e54c 100644 --- a/crates/sui-indexer-alt-graphql/src/metrics/subscription.rs +++ b/crates/sui-indexer-alt-graphql/src/metrics/subscription.rs @@ -53,6 +53,7 @@ pub struct SubscriptionMetrics { // Metrics aggregated across all subscribers. pub active_subscriptions: IntGaugeVec, pub subscriptions_opened: IntCounterVec, + pub subscriptions_rejected: IntCounterVec, pub subscription_terminations: IntCounterVec, pub subscription_duration: HistogramVec, pub payloads_delivered: IntCounterVec, @@ -138,6 +139,13 @@ impl SubscriptionMetrics { registry, ) .unwrap(), + subscriptions_rejected: register_int_counter_vec_with_registry!( + "graphql_subscription_rejected", + "Total subscriptions refused before opening, by type and reason (e.g. the server was at its concurrent-subscription capacity)", + &["type", "reason"], + registry, + ) + .unwrap(), subscription_terminations: register_int_counter_vec_with_registry!( "graphql_subscription_terminations", "Total subscriptions that have ended, by type and reason", diff --git a/crates/sui-indexer-alt-graphql/src/task/streaming/checkpoint_stream_task.rs b/crates/sui-indexer-alt-graphql/src/task/streaming/checkpoint_stream_task.rs index f91d7745ff16..2701c17e96d6 100644 --- a/crates/sui-indexer-alt-graphql/src/task/streaming/checkpoint_stream_task.rs +++ b/crates/sui-indexer-alt-graphql/src/task/streaming/checkpoint_stream_task.rs @@ -110,6 +110,8 @@ use super::checkpoint_resume::scan_checkpoints; #[cfg(any(feature = "staging", test))] use super::gap_recovery::CheckpointFetcher; use super::gap_recovery::recover_gap; +#[cfg(test)] +use super::lifecycle::SubscriberLimit; #[cfg(any(feature = "staging", test))] use super::lifecycle::SubscriptionLifecycleGuard; #[cfg(any(feature = "staging", test))] @@ -242,6 +244,7 @@ impl SubscriptionBroadcast { resume_from: Option, fetcher: F, config: &SubscriptionConfig, + guard: SubscriptionLifecycleGuard, ) -> impl Stream, RpcError>> + 'static { // Resubscribe and pin the handoff once the scan is within this many checkpoints of the tip. // Half the buffer leaves room for checkpoints arriving during the handoff, so it won't lag. @@ -250,7 +253,6 @@ impl SubscriptionBroadcast { let config = config.clone(); stream! { - let guard = SubscriptionLifecycleGuard::new("checkpoints", &self.metrics); let mut last_yielded: Option = resume_from; let mut receiver = self.broadcaster.resubscribe(); // `resubscribe` is future-only (delivers `handoff + 1`), so Phase 1 stops exactly at @@ -854,6 +856,15 @@ mod tests { tx.send(Arc::new(processed)).ok(); } + fn test_guard() -> SubscriptionLifecycleGuard { + SubscriptionLifecycleGuard::new( + "checkpoints", + &SubscriptionMetrics::new_for_test(), + &SubscriberLimit::new(10), + ) + .unwrap() + } + #[tokio::test] async fn subscribe_no_resume_yields_live_only() { use futures::FutureExt; @@ -862,7 +873,8 @@ mod tests { // Fetcher is unused since resume_from is None. let fetcher = MockFetcher::success_for_range(0..=0); - let stream = broadcast.subscribe(None, fetcher, &SubscriptionConfig::default()); + let stream = + broadcast.subscribe(None, fetcher, &SubscriptionConfig::default(), test_guard()); tokio::pin!(stream); // Poll once so the receiver gets pinned at tail=0 before any sends. @@ -886,7 +898,12 @@ mod tests { // resume_from = 2 → Phase 1 yields 3, 4, 5; then Phase 2 picks up live items. let fetcher = MockFetcher::success_for_range(3..=5); - let stream = broadcast.subscribe(Some(2), fetcher, &SubscriptionConfig::default()); + let stream = broadcast.subscribe( + Some(2), + fetcher, + &SubscriptionConfig::default(), + test_guard(), + ); tokio::pin!(stream); // Phase 1 catches up via scan. @@ -902,7 +919,8 @@ mod tests { async fn subscribe_yields_error_when_channel_closes() { let (tx, broadcast) = test_broadcast(/* first_live_checkpoint */ 1); let fetcher = MockFetcher::success_for_range(0..=0); - let stream = broadcast.subscribe(None, fetcher, &SubscriptionConfig::default()); + let stream = + broadcast.subscribe(None, fetcher, &SubscriptionConfig::default(), test_guard()); tokio::pin!(stream); // Dropping the sender closes the channel; subscriber should yield an error and end. diff --git a/crates/sui-indexer-alt-graphql/src/task/streaming/lifecycle.rs b/crates/sui-indexer-alt-graphql/src/task/streaming/lifecycle.rs index 2d2ac1fff2c1..35a165e5eb8b 100644 --- a/crates/sui-indexer-alt-graphql/src/task/streaming/lifecycle.rs +++ b/crates/sui-indexer-alt-graphql/src/task/streaming/lifecycle.rs @@ -1,15 +1,39 @@ // Copyright (c) Mysten Labs, Inc. // SPDX-License-Identifier: Apache-2.0 +use std::sync::Arc; + use prometheus::Histogram; use prometheus::IntCounter; use prometheus::IntCounterVec; use prometheus::IntGauge; +use tokio::sync::OwnedSemaphorePermit; +use tokio::sync::Semaphore; use tokio::sync::broadcast; use tokio::time::Instant; +use crate::error::RpcError; +use crate::error::resource_exhausted; use crate::metrics::SubscriptionMetrics; +/// Server-wide cap on the number of concurrent subscriptions, as a pool of admission slots. +/// Registered once in the schema data; each admitted subscription's guard holds one slot for its +/// lifetime, and a subscription opened while all slots are taken is rejected. +#[derive(Clone)] +pub(crate) struct SubscriberLimit(Arc); + +impl SubscriberLimit { + pub(crate) fn new(max_subscribers: usize) -> Self { + Self(Arc::new(Semaphore::new(max_subscribers))) + } + + /// Claim a slot, or `None` when all are taken. The slot is released when the returned permit is + /// dropped. + fn try_acquire(&self) -> Option { + self.0.clone().try_acquire_owned().ok() + } +} + /// Why a subscription ended, used to label `subscription_terminations`. #[derive(Clone, Copy)] pub(crate) enum SubscriptionTerminationReason { @@ -57,18 +81,36 @@ pub(crate) struct SubscriptionLifecycleGuard { subscription_type: &'static str, termination_reason: Option, started_at: Instant, + /// The concurrent-subscription slot this subscription holds; released when the guard drops. + _permit: OwnedSemaphorePermit, } impl SubscriptionLifecycleGuard { - pub(crate) fn new(subscription_type: &'static str, metrics: &SubscriptionMetrics) -> Self { + /// Admit a subscription: claim a concurrency slot from `subscriber_limit`, or return an + /// at-capacity error (counting a rejection) when the server is already full. The returned guard + /// holds the slot and all per-subscriber metric handles for the subscription's lifetime. + pub(crate) fn new( + subscription_type: &'static str, + metrics: &SubscriptionMetrics, + subscriber_limit: &SubscriberLimit, + ) -> Result { let label = subscription_type; + let Some(permit) = subscriber_limit.try_acquire() else { + metrics + .subscriptions_rejected + .with_label_values(&[label, "at_capacity"]) + .inc(); + return Err(resource_exhausted(anyhow::anyhow!( + "The server is at its concurrent-subscription capacity; retry later" + ))); + }; let active_subscriptions = metrics.active_subscriptions.with_label_values(&[label]); active_subscriptions.inc(); metrics .subscriptions_opened .with_label_values(&[label]) .inc(); - Self { + Ok(Self { active_subscriptions, subscription_terminations: metrics.subscription_terminations.clone(), subscription_duration: metrics.subscription_duration.with_label_values(&[label]), @@ -84,7 +126,8 @@ impl SubscriptionLifecycleGuard { subscription_type, termination_reason: None, started_at: Instant::now(), - } + _permit: permit, + }) } /// Record why the subscription ended, consuming the guard so it cannot be used afterward. @@ -140,7 +183,12 @@ mod tests { let label = "transactions"; { - let _guard = SubscriptionLifecycleGuard::new("transactions", &metrics); + let _guard = SubscriptionLifecycleGuard::new( + "transactions", + &metrics, + &SubscriberLimit::new(10), + ) + .unwrap(); assert_eq!(active(&metrics, label), 1); assert_eq!(opened(&metrics, label), 1); } @@ -157,7 +205,9 @@ mod tests { let label = "events"; { - let guard = SubscriptionLifecycleGuard::new("events", &metrics); + let guard = + SubscriptionLifecycleGuard::new("events", &metrics, &SubscriberLimit::new(10)) + .unwrap(); guard.terminate(SubscriptionTerminationReason::Error); } @@ -172,7 +222,9 @@ mod tests { async fn record_delivered_counts_by_phase() { let metrics = SubscriptionMetrics::new_for_test(); let label = "transactions"; - let guard = SubscriptionLifecycleGuard::new("transactions", &metrics); + let guard = + SubscriptionLifecycleGuard::new("transactions", &metrics, &SubscriberLimit::new(10)) + .unwrap(); guard.record_delivered(0); guard.record_delivered(0); @@ -183,6 +235,30 @@ mod tests { assert_eq!(delivery_lag_samples(&metrics, label), 2); } + /// A subscription opened at capacity is refused (`None`) and counted as rejected; freeing the + /// held slot lets the next one in. + #[tokio::test] + async fn rejects_at_capacity() { + let metrics = SubscriptionMetrics::new_for_test(); + let label = "transactions"; + let subscriber_limit = SubscriberLimit::new(1); + + let first = SubscriptionLifecycleGuard::new("transactions", &metrics, &subscriber_limit); + assert!(first.is_ok(), "first admitted"); + assert!( + SubscriptionLifecycleGuard::new("transactions", &metrics, &subscriber_limit).is_err(), + "second refused at capacity", + ); + assert_eq!(rejected(&metrics, label, "at_capacity"), 1); + + // Dropping the first frees its slot, admitting the next. + drop(first); + assert!( + SubscriptionLifecycleGuard::new("transactions", &metrics, &subscriber_limit).is_ok(), + "slot freed on drop", + ); + } + fn active(m: &SubscriptionMetrics, label: &str) -> i64 { m.active_subscriptions.with_label_values(&[label]).get() } @@ -191,6 +267,12 @@ mod tests { m.subscriptions_opened.with_label_values(&[label]).get() } + fn rejected(m: &SubscriptionMetrics, label: &str, reason: &str) -> u64 { + m.subscriptions_rejected + .with_label_values(&[label, reason]) + .get() + } + fn terminations(m: &SubscriptionMetrics, label: &str, reason: &str) -> u64 { m.subscription_terminations .with_label_values(&[label, reason]) diff --git a/crates/sui-indexer-alt-graphql/src/task/streaming/mod.rs b/crates/sui-indexer-alt-graphql/src/task/streaming/mod.rs index 545d4d0c0ba6..e7322236d786 100644 --- a/crates/sui-indexer-alt-graphql/src/task/streaming/mod.rs +++ b/crates/sui-indexer-alt-graphql/src/task/streaming/mod.rs @@ -74,6 +74,8 @@ pub(crate) use checkpoint_stream_task::reconnect_error; #[cfg(feature = "staging")] pub(crate) use gap_recovery::wait_for_pipelines_catching_up_at; #[cfg(feature = "staging")] +pub(crate) use lifecycle::SubscriberLimit; +#[cfg(feature = "staging")] pub(crate) use lifecycle::SubscriptionLifecycleGuard; #[cfg(feature = "staging")] pub(crate) use lifecycle::SubscriptionTerminationReason;