Expose the SubscribeRequest sink for runtime filter updates

crates/runtime/src/builder.rs

@@ -3,6 +3,7 @@ use shipstern_core::{
     instruction::InstructionUpdate, AccountUpdate, BlockMetaUpdate, BlockUpdate, SlotUpdate,
     TransactionUpdate,
 };
+use tokio::sync::mpsc;
 
 use crate::{
     config::ShipsternConfig,
@@ -12,6 +13,16 @@ use crate::{
     util, Runtime,
 };
 
+/// Depth of the filter update channel handed to callers by
+/// [`Runtime::filter_updates`]. Updates are rare, so a shallow queue is enough
+/// to keep a caller from blocking on a short burst.
+///
+/// Queued sets are forwarded in order and the server ends on the newest. The
+/// exception is a set the sink rejected: the source holds that one and a newer
+/// arrival replaces it rather than queueing behind it, so an intermediate set
+/// can be skipped after a rejection.
+const FILTER_UPDATE_CHANNEL_SIZE: usize = 8;
+
 /// Helper trait for defining the intended use for a builder.
 pub trait BuilderKind: Default {
     /// The type of error returned by the builder.
@@ -265,10 +276,14 @@ impl<S: SourceTrait> RuntimeBuilder<S> {
             return Err(BuilderError::SlotPipelineCollision);
         }
 
+        let (filter_updates_tx, filter_updates_rx) = mpsc::channel(FILTER_UPDATE_CHANNEL_SIZE);
+
         Ok(Runtime {
             buffer: buffer_cfg,
             source: source_cfg,
             pipelines,
+            filter_updates_tx: Some(filter_updates_tx),
+            filter_updates_rx,
             _source: std::marker::PhantomData,
             #[cfg(feature = "prometheus")]
             metrics_registry,

crates/runtime/src/lib.rs

@@ -15,6 +15,7 @@
 use std::marker::PhantomData;
 
 use config::BufferConfig;
+use shipstern_core::Filters;
 use tokio::sync::{mpsc, oneshot};
 use yellowstone_grpc_proto::tonic::Status;
 
@@ -83,6 +84,8 @@ pub struct Runtime<S: SourceTrait> {
     buffer: BufferConfig,
     source: S::Config,
     pipelines: handler::PipelineSets,
+    filter_updates_tx: Option<mpsc::Sender<Filters>>,
+    filter_updates_rx: mpsc::Receiver<Filters>,
     #[cfg(feature = "prometheus")]
     metrics_registry: prometheus::Registry,
     _source: PhantomData<S>,
@@ -91,6 +94,55 @@ pub struct Runtime<S: SourceTrait> {
 impl<S: SourceTrait> Runtime<S> {
     /// Create a new runtime builder.
     pub fn builder() -> RuntimeBuilder<S> { RuntimeBuilder::<S>::default() }
+
+    /// Take the sending half of the filter update channel.
+    ///
+    /// Each [`Filters`] sent replaces the whole subscription rather than
+    /// adding to it, because that is what the gRPC servers do with a
+    /// mid-stream request, so send the complete set every time.
+    ///
+    /// Keys are parser IDs. A key matching no registered pipeline still
+    /// changes the wire subscription, so the server streams that data and the
+    /// runtime then drops every update of it with only a trace-level line, so
+    /// take the keys from the parsers actually registered on this runtime.
+    ///
+    /// The server applies the new set promptly, but a consumer sees it only
+    /// once whatever it has already queued drains, so the delay is however far
+    /// behind the pipeline already was rather than a property of the update.
+    /// Measured against a live endpoint, a consumer running about 15 seconds
+    /// behind kept receiving the old set for roughly that long after the send,
+    /// and the first updates matching the new set arrived stale by the same
+    /// margin before catching up to real time. A pipeline keeping pace sees
+    /// the change almost at once.
+    ///
+    /// Treat a returned `send` as the request having been handed off, not as
+    /// the subscription having changed, and keep handlers able to cope with
+    /// updates matching the old set until the backlog clears.
+    ///
+    /// Delivery is best effort. A set rejected while the source is between
+    /// connections is retried once the stream recovers, but nothing reports
+    /// back to the sender either way.
+    ///
+    /// A set the server itself refuses, by exceeding its configured filter
+    /// limits for example, is answered on the stream with a code the client
+    /// does not treat as recoverable, and the run ends with that error rather
+    /// than the previous subscription staying in place. Under [`Self::run`]
+    /// and [`Self::run_async`] that error is fatal and exits the process, so a
+    /// set the provider will not accept takes the whole indexer down. Use
+    /// [`Self::try_run_async`] if a caller needs to survive one.
+    ///
+    /// Returns `None` when the source does not support filter updates, or
+    /// when the sender has already been taken. Call this before running the
+    /// runtime, since [`Self::run`], [`Self::try_run`], [`Self::run_async`]
+    /// and [`Self::try_run_async`] all consume it.
+    ///
+    pub fn filter_updates(&mut self) -> Option<mpsc::Sender<Filters>> {
+        if !S::supports_filter_updates() {
+            return None;
+        }
+
+        self.filter_updates_tx.take()
+    }
 }
 impl<S: SourceTrait> Runtime<S> {
     /// Create a new Tokio runtime and run the Shipstern runtime within it,
@@ -252,9 +304,17 @@ impl<S: SourceTrait> Runtime<S> {
         let filters = self.pipelines.filters();
 
         let source = S::new(self.source, filters);
+        let filter_updates = self.filter_updates_rx;
+
+        // Close the channel when nobody asked for the sending half, so a source
+        // that waits on updates is not left waiting on a sender that can never
+        // produce one.
+        drop(self.filter_updates_tx);
 
         tokio::spawn(async move {
-            let _ = source.connect(tx, status_tx).await;
+            let _ = source
+                .connect_with_filter_updates(tx, status_tx, filter_updates)
+                .await;
         });
 
         let signal;

crates/runtime/src/runtime_tests.rs

@@ -1,12 +1,21 @@
 use std::{
     borrow::Cow,
-    sync::atomic::{AtomicUsize, Ordering},
+    collections::{HashMap, HashSet},
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Mutex,
+    },
     time::Duration,
 };
 
 use async_trait::async_trait;
-use shipstern_core::{ParseResult, Parser, Prefilter, SlotUpdate};
-use tokio::sync::{mpsc::Sender, oneshot};
+use shipstern_core::{
+    AccountPrefilter, Filters, ParseResult, Parser, Prefilter, Pubkey, SlotUpdate,
+};
+use tokio::sync::{
+    mpsc::{Receiver, Sender},
+    oneshot,
+};
 use yellowstone_grpc_proto::{
     geyser::{subscribe_update::UpdateOneof, SlotStatus, SubscribeUpdate, SubscribeUpdateSlot},
     tonic,
@@ -373,6 +382,131 @@ impl Handler<SlotUpdate, SlotUpdate> for SlowSlotHandler {
     }
 }
 
+/// Parser IDs carried by the filter sets a source received through the
+/// runtime filter update channel. Tests run in parallel and all share this,
+/// so each one records under its own ID and asserts on that ID alone.
+static RECEIVED_FILTER_IDS: Mutex<Vec<String>> = Mutex::new(Vec::new());
+
+/// Parser ID used only by [`test_filter_updates_reach_a_supporting_source`].
+const SWAPPED_PARSER_ID: &str = "test::SwappedParser";
+
+#[derive(Debug)]
+struct MockFilterUpdateSource;
+
+#[async_trait]
+impl SourceTrait for MockFilterUpdateSource {
+    type Config = NullConfig;
+
+    fn new(_: NullConfig, _: Filters) -> Self { Self }
+
+    fn supports_filter_updates() -> bool { true }
+
+    async fn connect(
+        &self,
+        tx: Sender<Result<SubscribeUpdate, tonic::Status>>,
+        status_tx: oneshot::Sender<SourceExitStatus>,
+    ) -> Result<(), Error> {
+        wait_for_runtime_ready().await;
+        signal_stream_ended(status_tx);
+        hold_channel_open_briefly().await;
+        drop(tx);
+        Ok(())
+    }
+
+    async fn connect_with_filter_updates(
+        &self,
+        tx: Sender<Result<SubscribeUpdate, tonic::Status>>,
+        status_tx: oneshot::Sender<SourceExitStatus>,
+        mut filter_updates: Receiver<Filters>,
+    ) -> Result<(), Error> {
+        wait_for_runtime_ready().await;
+
+        if let Some(filters) = filter_updates.recv().await {
+            let mut ids = filters.parsers_filters.keys().cloned().collect::<Vec<_>>();
+            ids.sort();
+
+            RECEIVED_FILTER_IDS.lock().unwrap().extend(ids);
+        }
+
+        signal_stream_ended(status_tx);
+        hold_channel_open_briefly().await;
+        drop(tx);
+        Ok(())
+    }
+}
+
+/// A filter set carrying an actual account prefilter. `Prefilter::default()`
+/// has every sub-filter unset and converts to an entirely empty
+/// `SubscribeRequest`, so a test built on it would pass on a payload that says
+/// nothing on the wire.
+fn filters_for(ids: &[&str]) -> Filters {
+    let prefilter = Prefilter {
+        account: Some(AccountPrefilter {
+            accounts: HashSet::new(),
+            owners: HashSet::from([Pubkey::default()]),
+        }),
+        ..Default::default()
+    };
+
+    Filters::new(
+        ids.iter()
+            .map(|id| ((*id).to_owned(), prefilter.clone()))
+            .collect::<HashMap<_, _>>(),
+    )
+}
+
+#[tokio::test]
+async fn test_filter_updates_reach_a_supporting_source() {
+    let mut runtime = Runtime::<MockFilterUpdateSource>::builder()
+        .try_build(default_test_config())
+        .unwrap();
+
+    let updates = runtime
+        .filter_updates()
+        .expect("source advertises filter update support");
+
+    updates
+        .send(filters_for(&[SWAPPED_PARSER_ID]))
+        .await
+        .unwrap();
+
+    assert_server_hangup(runtime.try_run_async().await);
+
+    let received = RECEIVED_FILTER_IDS.lock().unwrap().clone();
+    assert!(
+        received.contains(&SWAPPED_PARSER_ID.to_owned()),
+        "source never saw the updated filter set, recorded {received:?}"
+    );
+}
+
+#[tokio::test]
+async fn test_filter_updates_unavailable_when_source_does_not_support_them() {
+    let mut runtime = Runtime::<MockStreamEndSource>::builder()
+        .try_build(default_test_config())
+        .unwrap();
+
+    assert!(runtime.filter_updates().is_none());
+}
+
+#[tokio::test]
+async fn test_filter_updates_handed_out_only_once() {
+    let mut runtime = Runtime::<MockFilterUpdateSource>::builder()
+        .try_build(default_test_config())
+        .unwrap();
+
+    assert!(runtime.filter_updates().is_some());
+    assert!(runtime.filter_updates().is_none());
+}
+
+#[tokio::test]
+async fn test_source_runs_when_the_filter_update_handle_is_never_taken() {
+    let runtime = Runtime::<MockFilterUpdateSource>::builder()
+        .try_build(default_test_config())
+        .unwrap();
+
+    assert_server_hangup(runtime.try_run_async().await);
+}
+
 #[tokio::test]
 async fn test_stream_end_returns_error() {
     let runtime = Runtime::<MockStreamEndSource>::builder()

crates/runtime/src/sources.rs

@@ -5,7 +5,10 @@
 
 use async_trait::async_trait;
 use shipstern_core::Filters;
-use tokio::sync::{mpsc::Sender, oneshot};
+use tokio::sync::{
+    mpsc::{Receiver, Sender},
+    oneshot,
+};
 use yellowstone_grpc_proto::{geyser::SubscribeUpdate, tonic};
 
 /// How a source exited.
@@ -45,4 +48,33 @@ pub trait SourceTrait: std::fmt::Debug + Send + Sync + 'static {
         tx: Sender<Result<SubscribeUpdate, tonic::Status>>,
         status_tx: oneshot::Sender<SourceExitStatus>,
     ) -> Result<(), crate::Error>;
+
+    /// Whether this source applies filter updates to a live subscription.
+    ///
+    /// The runtime checks this before handing a caller the sending half of the
+    /// filter update channel, so a caller wiring updates to a source that
+    /// ignores them finds out at the call site instead of silently sending
+    /// into a void.
+    ///
+    #[must_use]
+    fn supports_filter_updates() -> bool { false }
+
+    /// Connect and stream updates, applying filter sets received on
+    /// `filter_updates` to the live subscription.
+    ///
+    /// The default ignores `filter_updates` and defers to [`Self::connect`],
+    /// so a source that cannot change its subscription mid-stream needs no
+    /// implementation. Override this together with
+    /// [`Self::supports_filter_updates`].
+    ///
+    async fn connect_with_filter_updates(
+        &self,
+        tx: Sender<Result<SubscribeUpdate, tonic::Status>>,
+        status_tx: oneshot::Sender<SourceExitStatus>,
+        filter_updates: Receiver<Filters>,
+    ) -> Result<(), crate::Error> {
+        drop(filter_updates);
+
+        self.connect(tx, status_tx).await
+    }
 }

crates/yellowstone-grpc-source/Cargo.toml

@@ -10,7 +10,7 @@ readme = "./README.md"
 
 [dependencies]
 async-trait = { workspace = true }
-tokio = { workspace = true, features = ["rt-multi-thread", "signal"] }
+tokio = { workspace = true, features = ["rt-multi-thread", "signal", "macros", "time"] }
 tracing = { workspace = true }
 futures-util = { workspace = true, features = ["sink"] }
 shipstern = { workspace = true }