connect with timeout with buffers

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "geyser-grpc-connector"
-version = "0.10.1+yellowstone.1.12"
+version = "0.11.0+yellowstone.1.12"
 edition = "2021"
 
 description = "Multiplexing and Reconnection on Yellowstone gRPC Geyser client streaming"
@@ -27,6 +27,8 @@ tracing = "0.1.37"
 itertools = "0.10.5"
 derive_more = "0.99.17"
 
+tonic-health = "0.10.2"
+
 base64 = "0.21.5"
 bincode = "1.3.3"
 

examples/stream_blocks_mainnet_task.rs

@@ -1,9 +1,7 @@
-use futures::Stream;
 use log::{info, warn};
 use solana_sdk::clock::Slot;
 use solana_sdk::commitment_config::CommitmentConfig;
 use std::env;
-use std::pin::pin;
 
 use base64::Engine;
 use itertools::Itertools;
@@ -22,12 +20,8 @@ use solana_sdk::transaction::TransactionError;
 use tokio::sync::mpsc::Receiver;
 use yellowstone_grpc_proto::geyser::SubscribeUpdateBlock;
 
-use geyser_grpc_connector::grpc_subscription_autoreconnect_tasks::{
-    create_geyser_autoconnection_task, create_geyser_autoconnection_task_with_mpsc,
-};
-use geyser_grpc_connector::grpcmultiplex_fastestwins::{
-    create_multiplexed_stream, FromYellowstoneExtractor,
-};
+use geyser_grpc_connector::grpc_subscription_autoreconnect_tasks::create_geyser_autoconnection_task_with_mpsc;
+use geyser_grpc_connector::grpcmultiplex_fastestwins::FromYellowstoneExtractor;
 use geyser_grpc_connector::{GeyserFilter, GrpcConnectionTimeouts, GrpcSourceConfig, Message};
 use tokio::time::{sleep, Duration};
 use yellowstone_grpc_proto::geyser::subscribe_update::UpdateOneof;

src/grpc_subscription_autoreconnect_streams.rs

@@ -1,11 +1,11 @@
-use crate::{Attempt, GrpcSourceConfig, Message};
+use crate::{yellowstone_grpc_util, Attempt, GrpcSourceConfig, Message};
 use async_stream::stream;
 use futures::{Stream, StreamExt};
 use log::{debug, info, log, trace, warn, Level};
 use std::time::Duration;
 use tokio::task::JoinHandle;
 use tokio::time::{sleep, timeout};
-use yellowstone_grpc_client::{GeyserGrpcClient, GeyserGrpcClientResult};
+use yellowstone_grpc_client::GeyserGrpcClientResult;
 use yellowstone_grpc_proto::geyser::{SubscribeRequest, SubscribeUpdate};
 use yellowstone_grpc_proto::tonic::Status;
 
@@ -45,15 +45,18 @@ pub fn create_geyser_reconnecting_stream(
                         log!(if attempt > 1 { Level::Warn } else { Level::Debug }, "Connecting attempt #{} to {}", attempt, addr);
                         async move {
 
-                            let connect_result = GeyserGrpcClient::connect_with_timeout(
-                                    addr, token, config,
-                                    connect_timeout,
-                                    request_timeout,
-                                false)
-                                .await;
+                            let buffer_config = yellowstone_grpc_util::GeyserGrpcClientBufferConfig::optimize_for_subscription(&subscribe_filter);
+                            debug!("Using Grpc Buffer config {:?}", buffer_config);
+                            let connect_result = yellowstone_grpc_util::connect_with_timeout_with_buffers(
+                                addr,
+                                token,
+                                config,
+                                connect_timeout,
+                                request_timeout,
+                                buffer_config,
+                            );
                             let mut client = connect_result?;
 
-
                             debug!("Subscribe with filter {:?}", subscribe_filter);
 
                             let subscribe_result = timeout(subscribe_timeout.unwrap_or(Duration::MAX),

src/grpc_subscription_autoreconnect_tasks.rs

@@ -1,4 +1,4 @@
-use crate::{Attempt, GrpcSourceConfig, Message};
+use crate::{yellowstone_grpc_util, Attempt, GrpcSourceConfig, Message};
 use futures::{Stream, StreamExt};
 use log::{debug, error, info, log, trace, warn, Level};
 use std::time::Duration;
@@ -75,15 +75,17 @@ pub fn create_geyser_autoconnection_task_with_mpsc(
                         attempt,
                         addr
                     );
-                    let connect_result = GeyserGrpcClient::connect_with_timeout(
+
+                    let buffer_config = yellowstone_grpc_util::GeyserGrpcClientBufferConfig::optimize_for_subscription(&subscribe_filter);
+                    debug!("Using Grpc Buffer config {:?}", buffer_config);
+                    let connect_result = yellowstone_grpc_util::connect_with_timeout_with_buffers(
                         addr,
                         token,
                         config,
                         connect_timeout,
                         request_timeout,
-                        false,
-                    )
-                    .await;
+                        buffer_config,
+                    );
 
                     match connect_result {
                         Ok(client) => ConnectionState::Connecting(attempt, client),

src/lib.rs

@@ -1,5 +1,5 @@
 use solana_sdk::commitment_config::CommitmentConfig;
 use std::collections::HashMap;
 use std::fmt::{Debug, Display};
 use std::time::Duration;
 use yellowstone_grpc_proto::geyser::{
@@ -13,6 +13,7 @@
 pub mod grpc_subscription_autoreconnect_tasks;
 pub mod grpcmultiplex_fastestwins;
 mod obfuscate;
+mod yellowstone_grpc_util;
 
 // 1-based attempt counter
 type Attempt = u32;

src/yellowstone_grpc_util.rs

@@ -0,0 +1,93 @@
+use std::time::Duration;
+use tonic_health::pb::health_client::HealthClient;
+use yellowstone_grpc_client::{GeyserGrpcClient, GeyserGrpcClientResult, InterceptorXToken};
+use yellowstone_grpc_proto::geyser::geyser_client::GeyserClient;
+use yellowstone_grpc_proto::geyser::SubscribeRequest;
+use yellowstone_grpc_proto::prost::bytes::Bytes;
+use yellowstone_grpc_proto::tonic;
+use yellowstone_grpc_proto::tonic::metadata::errors::InvalidMetadataValue;
+use yellowstone_grpc_proto::tonic::metadata::AsciiMetadataValue;
+use yellowstone_grpc_proto::tonic::service::Interceptor;
+use yellowstone_grpc_proto::tonic::transport::ClientTlsConfig;
+
+// see https://github.com/hyperium/tonic/blob/v0.10.2/tonic/src/transport/channel/mod.rs
+const DEFAULT_BUFFER_SIZE: usize = 1024;
+// see https://github.com/hyperium/hyper/blob/v0.14.28/src/proto/h2/client.rs#L45
+const DEFAULT_CONN_WINDOW: u32 = 1024 * 1024 * 5; // 5mb
+const DEFAULT_STREAM_WINDOW: u32 = 1024 * 1024 * 2; // 2mb
+
+#[derive(Debug, Clone)]
+pub struct GeyserGrpcClientBufferConfig {
+    pub buffer_size: Option<usize>,
+    pub conn_window: Option<u32>,
+    pub stream_window: Option<u32>,
+}
+
+impl Default for GeyserGrpcClientBufferConfig {
+    fn default() -> Self {
+        GeyserGrpcClientBufferConfig {
+            buffer_size: Some(DEFAULT_BUFFER_SIZE),
+            conn_window: Some(DEFAULT_CONN_WINDOW),
+            stream_window: Some(DEFAULT_STREAM_WINDOW),
+        }
+    }
+}
+
+impl GeyserGrpcClientBufferConfig {
+    pub fn optimize_for_subscription(filter: &SubscribeRequest) -> GeyserGrpcClientBufferConfig {
+        if !filter.blocks.is_empty() {
+            GeyserGrpcClientBufferConfig {
+                buffer_size: Some(65536),     // 64kb (default: 1k)
+                conn_window: Some(5242880),   // 5mb (=default)
+                stream_window: Some(4194304), // 4mb (default: 2m)
+            }
+        } else {
+            GeyserGrpcClientBufferConfig::default()
+        }
+    }
+}
+
+pub fn connect_with_timeout_with_buffers<E, T>(
+    endpoint: E,
+    x_token: Option<T>,
+    tls_config: Option<ClientTlsConfig>,
+    connect_timeout: Option<Duration>,
+    request_timeout: Option<Duration>,
+    buffer_config: GeyserGrpcClientBufferConfig,
+) -> GeyserGrpcClientResult<GeyserGrpcClient<impl Interceptor>>
+where
+    E: Into<Bytes>,
+    T: TryInto<AsciiMetadataValue, Error = InvalidMetadataValue>,
+{
+    // see https://github.com/blockworks-foundation/geyser-grpc-connector/issues/10
+    let mut endpoint = tonic::transport::Endpoint::from_shared(endpoint)?
+        .buffer_size(buffer_config.buffer_size)
+        .initial_connection_window_size(buffer_config.conn_window)
+        .initial_stream_window_size(buffer_config.stream_window);
+
+    if let Some(tls_config) = tls_config {
+        endpoint = endpoint.tls_config(tls_config)?;
+    }
+
+    if let Some(connect_timeout) = connect_timeout {
+        endpoint = endpoint.timeout(connect_timeout);
+    }
+
+    if let Some(request_timeout) = request_timeout {
+        endpoint = endpoint.timeout(request_timeout);
+    }
+
+    let x_token: Option<AsciiMetadataValue> = match x_token {
+        Some(x_token) => Some(x_token.try_into()?),
+        None => None,
+    };
+    let interceptor = InterceptorXToken { x_token };
+
+    let channel = endpoint.connect_lazy();
+    let client = GeyserGrpcClient::new(
+        HealthClient::with_interceptor(channel.clone(), interceptor.clone()),
+        GeyserClient::with_interceptor(channel, interceptor)
+            .max_decoding_message_size(GeyserGrpcClient::max_decoding_message_size()),
+    );
+    Ok(client)
+}