yellowstone_fumarole_client/connectors/
dataplane.rs1use {
2 crate::{
3 FumaroleGrpcConnector,
4 core::{
5 ports::FumaroleDataplaneConnector,
6 runtime::{DataplaneErrorKind, DataplaneStreamError},
7 },
8 proto::{DataCommand, DataResponse},
9 },
10 futures::{Future, Sink, Stream, StreamExt},
11 std::{error::Error as _, pin::Pin, task::Poll},
12 tokio::sync::mpsc,
13 tokio_stream::wrappers::ReceiverStream,
14 tonic::{Code, Streaming, transport},
15};
16
17impl From<tonic::Status> for DataplaneStreamError {
18 fn from(status: tonic::Status) -> Self {
19 let message = status.message().to_ascii_lowercase();
20 if let Some(source) = status.source() {
21 if source.downcast_ref::<transport::Error>().is_some() {
22 return Self::new(
23 DataplaneErrorKind::RecoverableTransport,
24 status.to_string(),
25 Some(Box::new(status)),
26 );
27 }
28 }
29
30 let kind = match status.code() {
31 Code::Unavailable
32 | Code::Internal
33 | Code::Aborted
34 | Code::ResourceExhausted
35 | Code::DataLoss
36 | Code::Unknown
37 | Code::Cancelled
38 | Code::DeadlineExceeded => DataplaneErrorKind::RecoverableTransport,
39 Code::NotFound => DataplaneErrorKind::SlotNotFound,
40 Code::InvalidArgument if message.contains("filter") => {
41 DataplaneErrorKind::InvalidSubscribeFilter
42 }
43 _ => DataplaneErrorKind::NonRecoverable,
44 };
45
46 Self::new(kind, status.to_string(), Some(Box::new(status)))
47 }
48}
49
50type DataplaneSinkSendError = mpsc::error::SendError<DataCommand>;
51
52fn create_dataplane_sink(
53 tx: mpsc::Sender<DataCommand>,
54) -> impl Sink<DataCommand, Error = DataplaneSinkSendError> + Send {
55 futures::sink::unfold(tx, |tx, cmd| async move {
56 tx.send(cmd).await?;
57 Ok::<_, DataplaneSinkSendError>(tx)
58 })
59}
60
61impl FumaroleDataplaneConnector for FumaroleGrpcConnector {
62 type DataplaneSubscribeError = tonic::Status;
63 type DataplaneSinkError = DataplaneSinkSendError;
64 type DataplaneSink = Pin<Box<dyn Sink<DataCommand, Error = Self::DataplaneSinkError> + Send>>;
65 type DataplaneStream = TonicDataplaneStreamAdapter;
66 type DataplaneSubscribeFut = Pin<
67 Box<
68 dyn Future<
69 Output = Result<
70 (Self::DataplaneSink, Self::DataplaneStream),
71 Self::DataplaneSubscribeError,
72 >,
73 > + Send,
74 >,
75 >;
76
77 fn subscribe_data(&self) -> Self::DataplaneSubscribeFut {
78 let mut client = self.connect_lazy();
79 Box::pin(async move {
80 let (tx, rx) = mpsc::channel(100);
81 let response = client.subscribe_data(ReceiverStream::new(rx)).await?;
82 let sink: Self::DataplaneSink = Box::pin(create_dataplane_sink(tx));
83 let stream: Self::DataplaneStream = TonicDataplaneStreamAdapter {
84 inner: response.into_inner(),
85 };
86 Ok((sink, stream))
87 })
88 }
89}
90
91pub struct TonicDataplaneStreamAdapter {
92 inner: Streaming<DataResponse>,
93}
94
95impl Stream for TonicDataplaneStreamAdapter {
96 type Item = Result<DataResponse, DataplaneStreamError>;
97
98 fn poll_next(
99 mut self: Pin<&mut Self>,
100 cx: &mut std::task::Context<'_>,
101 ) -> Poll<Option<Self::Item>> {
102 match self.inner.poll_next_unpin(cx) {
103 Poll::Ready(Some(Ok(response))) => Poll::Ready(Some(Ok(response))),
104 Poll::Ready(Some(Err(status))) => {
105 Poll::Ready(Some(Err(DataplaneStreamError::from(status))))
106 }
107 Poll::Ready(None) => Poll::Ready(None),
108 Poll::Pending => Poll::Pending,
109 }
110 }
111}