use std::sync::Arc;
use futures::StreamExt;
use solana_sdk::pubkey::Pubkey;
use crate::common::AnyResult;
use crate::protos::shredstream::SubscribeEntriesRequest;
use crate::streaming::common::{EventProcessor, SubscriptionHandle};
use crate::streaming::event_parser::common::filter::EventTypeFilter;
use crate::streaming::event_parser::{Protocol, UnifiedEvent};
use crate::streaming::shred::TransactionWithSlot;
use log::error;
use solana_entry::entry::Entry;
use super::ShredStreamGrpc;
impl ShredStreamGrpc {
pub async fn shredstream_subscribe<F>(
&self,
protocols: Vec<Protocol>,
bot_wallet: Option<Pubkey>,
event_type_filter: Option<EventTypeFilter>,
callback: F,
) -> AnyResult<()>
where
F: Fn(Box<dyn UnifiedEvent>) + Send + Sync + 'static,
{
self.stop().await;
let mut metrics_handle = None;
if self.config.enable_metrics {
metrics_handle = self.metrics_manager.start_auto_monitoring().await;
}
let mut event_processor =
EventProcessor::new(self.metrics_manager.clone(), self.config.clone());
event_processor.set_protocols_and_event_type_filter(
protocols,
event_type_filter,
self.config.backpressure.clone(),
Some(Arc::new(callback)),
);
let mut client = (*self.shredstream_client).clone();
let request = tonic::Request::new(SubscribeEntriesRequest {});
let mut stream = client.subscribe_entries(request).await?.into_inner();
let event_processor_clone = event_processor.clone();
let stream_task = tokio::spawn(async move {
while let Some(message) = stream.next().await {
match message {
Ok(msg) => {
if let Ok(entries) = bincode::deserialize::<Vec<Entry>>(&msg.entries) {
for entry in entries {
for transaction in entry.transactions {
let transaction_with_slot = TransactionWithSlot::new(
transaction.clone(),
msg.slot,
chrono::Utc::now().timestamp_micros(),
);
if let Err(e) = event_processor_clone
.process_shred_transaction_with_metrics(
transaction_with_slot,
bot_wallet,
)
.await
{
error!("Error handling message: {e:?}");
}
}
}
}
continue;
}
Err(error) => {
error!("Stream error: {error:?}");
break;
}
}
}
});
let subscription_handle = SubscriptionHandle::new(stream_task, None, metrics_handle);
let mut handle_guard = self.subscription_handle.lock().await;
*handle_guard = Some(subscription_handle);
Ok(())
}
}