Skip to main content

yellowstone_fumarole_client/connectors/
dataplane.rs

1use {
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}