-
Notifications
You must be signed in to change notification settings - Fork 824
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Refactor to share code between do_put and do_exchange calls #5728
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -15,7 +15,7 @@ | |
// specific language governing permissions and limitations | ||
// under the License. | ||
|
||
use std::task::Poll; | ||
use std::{pin::Pin, task::Poll}; | ||
|
||
use crate::{ | ||
decode::FlightRecordBatchStream, | ||
|
@@ -28,6 +28,7 @@ use crate::{ | |
use arrow_schema::Schema; | ||
use bytes::Bytes; | ||
use futures::{ | ||
channel::oneshot::{Receiver, Sender}, | ||
future::ready, | ||
ready, | ||
stream::{self, BoxStream}, | ||
|
@@ -364,33 +365,18 @@ impl FlightClient { | |
&mut self, | ||
request: S, | ||
) -> Result<BoxStream<'static, Result<PutResult>>> { | ||
let (sender, mut receiver) = futures::channel::oneshot::channel(); | ||
let (sender, receiver) = futures::channel::oneshot::channel(); | ||
|
||
// Intercepts client errors and sends them to the oneshot channel above | ||
let mut request = Box::pin(request); // Pin to heap | ||
let mut sender = Some(sender); // Wrap into Option so can be taken | ||
let request_stream = futures::stream::poll_fn(move |cx| { | ||
Poll::Ready(match ready!(request.poll_next_unpin(cx)) { | ||
Some(Ok(data)) => Some(data), | ||
Some(Err(e)) => { | ||
let _ = sender.take().unwrap().send(e); | ||
None | ||
} | ||
None => None, | ||
}) | ||
}); | ||
let request = Box::pin(request); // Pin to heap | ||
let request_stream = FallibleRequestStream::new(sender, request); | ||
|
||
let request = self.make_request(request_stream); | ||
let mut response_stream = self.inner.do_put(request).await?.into_inner(); | ||
let response_stream = self.inner.do_put(request).await?.into_inner(); | ||
|
||
// Forwards errors from the error oneshot with priority over responses from server | ||
let error_stream = futures::stream::poll_fn(move |cx| { | ||
if let Poll::Ready(Ok(err)) = receiver.poll_unpin(cx) { | ||
return Poll::Ready(Some(Err(err))); | ||
} | ||
let next = ready!(response_stream.poll_next_unpin(cx)); | ||
Poll::Ready(next.map(|x| x.map_err(FlightError::Tonic))) | ||
}); | ||
let response_stream = Box::pin(response_stream); | ||
let error_stream = FallibleTonicResponseStream::new(receiver, response_stream); | ||
|
||
// combine the response from the server and any error from the client | ||
Ok(error_stream.boxed()) | ||
|
@@ -433,33 +419,17 @@ impl FlightClient { | |
&mut self, | ||
request: S, | ||
) -> Result<FlightRecordBatchStream> { | ||
let (sender, mut receiver) = futures::channel::oneshot::channel(); | ||
let (sender, receiver) = futures::channel::oneshot::channel(); | ||
|
||
let request = Box::pin(request); | ||
// Intercepts client errors and sends them to the oneshot channel above | ||
let mut request = Box::pin(request); // Pin to heap | ||
let mut sender = Some(sender); // Wrap into Option so can be taken | ||
let request_stream = futures::stream::poll_fn(move |cx| { | ||
Poll::Ready(match ready!(request.poll_next_unpin(cx)) { | ||
Some(Ok(data)) => Some(data), | ||
Some(Err(e)) => { | ||
let _ = sender.take().unwrap().send(e); | ||
None | ||
} | ||
None => None, | ||
}) | ||
}); | ||
let request_stream = FallibleRequestStream::new(sender, request); | ||
|
||
let request = self.make_request(request_stream); | ||
let mut response_stream = self.inner.do_exchange(request).await?.into_inner(); | ||
let response_stream = self.inner.do_exchange(request).await?.into_inner(); | ||
|
||
// Forwards errors from the error oneshot with priority over responses from server | ||
let error_stream = futures::stream::poll_fn(move |cx| { | ||
if let Poll::Ready(Ok(err)) = receiver.poll_unpin(cx) { | ||
return Poll::Ready(Some(Err(err))); | ||
} | ||
let next = ready!(response_stream.poll_next_unpin(cx)); | ||
Poll::Ready(next.map(|x| x.map_err(FlightError::Tonic))) | ||
}); | ||
let response_stream = Box::pin(response_stream); | ||
let error_stream = FallibleTonicResponseStream::new(receiver, response_stream); | ||
|
||
// combine the response from the server and any error from the client | ||
Ok(FlightRecordBatchStream::new_from_flight_data(error_stream)) | ||
|
@@ -704,3 +674,99 @@ impl FlightClient { | |
request | ||
} | ||
} | ||
|
||
/// Wrapper around fallible stream such that when | ||
/// it encounters an error it uses the oneshot sender to | ||
/// notify the error and stop any further streaming. See `do_put` or | ||
/// `do_exchange` for it's uses. | ||
struct FallibleRequestStream<T, E> { | ||
/// sender to notify error | ||
sender: Option<Sender<E>>, | ||
/// fallible stream | ||
fallible_stream: Pin<Box<dyn Stream<Item = std::result::Result<T, E>> + Send + 'static>>, | ||
} | ||
|
||
impl<T, E> FallibleRequestStream<T, E> { | ||
fn new( | ||
sender: Sender<E>, | ||
fallible_stream: Pin<Box<dyn Stream<Item = std::result::Result<T, E>> + Send + 'static>>, | ||
) -> Self { | ||
Self { | ||
sender: Some(sender), | ||
fallible_stream, | ||
} | ||
} | ||
} | ||
|
||
impl<T, E> Stream for FallibleRequestStream<T, E> { | ||
type Item = T; | ||
|
||
fn poll_next( | ||
self: std::pin::Pin<&mut Self>, | ||
cx: &mut std::task::Context<'_>, | ||
) -> std::task::Poll<Option<Self::Item>> { | ||
let pinned = self.get_mut(); | ||
let mut request_streams = pinned.fallible_stream.as_mut(); | ||
match ready!(request_streams.poll_next_unpin(cx)) { | ||
Some(Ok(data)) => Poll::Ready(Some(data)), | ||
Some(Err(e)) => { | ||
// unwrap() here is safe, ownership of sender will | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Another potential improvement might be to do something like if let Some(sender) = pinned.sender.take() {
sender.send();
} And that way avoid a panic However, I see this unwrap is simply moved here from elsewhere in the PR so I think we can improve it in the future There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Ah good point, it would've taken less time to swap the code than writing the justification comments. |
||
// be moved only once as this stream will not be polled | ||
// again | ||
let _ = pinned.sender.take().unwrap().send(e); | ||
Poll::Ready(None) | ||
} | ||
None => Poll::Ready(None), | ||
} | ||
} | ||
} | ||
|
||
/// Wrapper for a tonic response stream that can produce a tonic | ||
/// error. This is tied to a oneshot receiver which can be notified | ||
/// of other errors. When it receives an error through receiver | ||
/// end, it prioritises that error to be sent back. See `do_put` or | ||
/// `do_exchange` for it's uses | ||
struct FallibleTonicResponseStream<T> { | ||
/// Receiver for FlightError | ||
receiver: Receiver<FlightError>, | ||
/// Tonic response stream | ||
response_stream: | ||
Pin<Box<dyn Stream<Item = std::result::Result<T, tonic::Status>> + Send + 'static>>, | ||
} | ||
|
||
impl<T> FallibleTonicResponseStream<T> { | ||
fn new( | ||
receiver: Receiver<FlightError>, | ||
response_stream: Pin< | ||
Box<dyn Stream<Item = std::result::Result<T, tonic::Status>> + Send + 'static>, | ||
>, | ||
) -> Self { | ||
Self { | ||
receiver, | ||
response_stream, | ||
} | ||
} | ||
} | ||
|
||
impl<T> Stream for FallibleTonicResponseStream<T> { | ||
type Item = Result<T>; | ||
|
||
fn poll_next( | ||
self: Pin<&mut Self>, | ||
cx: &mut std::task::Context<'_>, | ||
) -> Poll<Option<Self::Item>> { | ||
let pinned = self.get_mut(); | ||
let receiver = &mut pinned.receiver; | ||
// Prioritise sending the error that's been notified over | ||
// polling the response_stream | ||
if let Poll::Ready(Ok(err)) = receiver.poll_unpin(cx) { | ||
return Poll::Ready(Some(Err(err))); | ||
}; | ||
|
||
match ready!(pinned.response_stream.poll_next_unpin(cx)) { | ||
Some(Ok(res)) => Poll::Ready(Some(Ok(res))), | ||
Some(Err(status)) => Poll::Ready(Some(Err(FlightError::Tonic(status)))), | ||
None => Poll::Ready(None), | ||
} | ||
} | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
BTW I especially appreciate the understandable comments in this PR 🙏