feat: add multiple program subscriptions (#160)

Author: sonicfromnewyoke

.gitignore

@@ -1,4 +1,6 @@
 target/
 crates/mock/fixtures
 *.env*
+.DS_Store
+**/.DS_Store
 Vixen.toml

crates/proto/proto/stream.proto

@@ -8,6 +8,7 @@ package vixen.stream;
 // updates for a given program ID.
 service ProgramStreams {
   rpc Subscribe(SubscribeRequest) returns (stream SubscribeUpdate) {}
+  rpc SubscribeMany(SubscribeManyRequest) returns (stream SubscribeUpdate) {}
 }
 
 // Request to subscribe to a stream of updates for a given program ID.
@@ -16,6 +17,12 @@ message SubscribeRequest {
   string program = 1;
 }
 
+// Request to subscribe to a stream of updates for a given program IDs.
+message SubscribeManyRequest {
+  // The program IDs to subscribe to.
+  repeated string programs = 1;
+}
+
 // Update from the requested program containing a parsed value.
 message SubscribeUpdate {
   // The parsed value.

crates/runtime/src/handler.rs

@@ -5,7 +5,7 @@ use std::{borrow::Cow, collections::HashMap, pin::Pin};
 
 use futures_util::{Future, FutureExt, StreamExt};
 use smallvec::SmallVec;
-use tracing::{warn, Instrument, Span};
+use tracing::{trace, Instrument, Span};
 use vixen_core::{
     AccountUpdate, BlockMetaUpdate, BlockUpdate, GetPrefilter, ParserId, SlotUpdate,
     TransactionUpdate,
@@ -95,9 +95,9 @@ mod pipeline_error {
 
     #[derive(Debug, thiserror::Error)]
     pub enum Error {
-        #[error("Error parsing input value")]
+        #[error("Error parsing input value: ({0})")]
         Parser(#[source] BoxedError),
-        #[error("Handler returned an error on parsed value")]
+        #[error("Handler returned an error on parsed value: ({0})")]
         Handler(#[source] BoxedError),
     }
 
@@ -320,7 +320,7 @@ where I::Item: AsRef<str> + Send + 'm
             let pipeline = pipelines.0.get(filter);
 
             if pipeline.is_none() {
-                warn!(filter, "No pipeline matched filter on incoming update");
+                trace!(filter, "No pipeline matched filter on incoming update");
             }
 
             pipeline.map(|p| (f, p))

crates/stream/src/grpc.rs

@@ -13,7 +13,7 @@ use yellowstone_vixen_proto::{
     stream::{
         self,
         program_streams_server::{ProgramStreams, ProgramStreamsServer},
-        SubscribeRequest, SubscribeUpdate,
+        SubscribeManyRequest, SubscribeRequest, SubscribeUpdate,
     },
     tonic::{self, transport, Request, Response, Status},
     tonic_reflection,
@@ -46,6 +46,7 @@ pub(super) struct Service(Channels);
 
 #[tonic::async_trait]
 impl ProgramStreams for Service {
+    type SubscribeManyStream = futures_util::stream::SelectAll<ReceiverStream>;
     type SubscribeStream = futures_util::stream::SelectAll<ReceiverStream>;
 
     async fn subscribe(
@@ -67,6 +68,35 @@ impl ProgramStreams for Service {
 
         Ok(Response::new(stream))
     }
+
+    async fn subscribe_many(
+        &self,
+        request: Request<SubscribeManyRequest>,
+    ) -> Result<Response<Self::SubscribeManyStream>, Status> {
+        let pubkeys: Vec<Pubkey> = request
+            .into_inner()
+            .programs
+            .iter()
+            .map(|p| {
+                p.parse()
+                    .map_err(|e: yellowstone_vixen_core::KeyFromStrError| {
+                        Status::new(tonic::Code::InvalidArgument, e.to_string())
+                    })
+            })
+            .collect::<Result<Vec<_>, _>>()?;
+
+        let rxs: Vec<_> = pubkeys
+            .iter()
+            .filter_map(|key| self.0.get(key))
+            .flatten()
+            .collect();
+
+        // TODO: make max_tries configurable?
+        let stream =
+            futures_util::stream::select_all(rxs.iter().map(|rx| ReceiverStream::new(rx, 8)));
+
+        Ok(Response::new(stream))
+    }
 }
 
 type BoxedRx = Box<Receiver>;