alamb commented on code in PR #5728:
URL: https://github.com/apache/arrow-rs/pull/5728#discussion_r1591462072
##########
arrow-flight/src/client.rs:
##########
@@ -704,3 +679,75 @@ impl FlightClient {
request
}
}
+
+struct FallibleRequestStream {
Review Comment:
looks good to me 👍
Perhaps we could add some comments about what this struct is doing and what
it is for.
##########
arrow-flight/src/client.rs:
##########
@@ -704,3 +679,75 @@ impl FlightClient {
request
}
}
+
+struct FallibleRequestStream {
+ sender: Option<Sender<FlightError>>,
+ request_streams: Pin<Box<dyn Stream<Item = Result<FlightData>> + Send +
'static>>,
+}
+
+impl Stream for FallibleRequestStream {
+ type Item = FlightData;
+
+ 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.request_streams.as_mut();
+ match request_streams.poll_next_unpin(cx) {
+ Poll::Ready(s) => {
+ match s {
+ Some(Ok(data)) => Poll::Ready(Some(data)),
+ Some(Err(e)) => {
+ let _ = pinned.sender.take().unwrap().send(e);
+ Poll::Ready(None)
+ },
+ None => Poll::Ready(None),
+ }
+
+ },
+ Poll::Pending => Poll::Pending,
+ }
+ }
+}
+
+impl FallibleRequestStream {
+ fn new(sender: Sender<FlightError>, request_streams: Pin<Box<dyn
Stream<Item = Result<FlightData>> + Send + 'static>>) -> Self {
+ Self {
+ sender: Some(sender),
+ request_streams
+ }
+ }
+}
+
+struct FallibleResponseStream<T> {
+ receiver: Receiver<FlightError>,
+ response_streams: Pin<Box<dyn Stream<Item = std::result::Result<T,
tonic::Status>> + Send + 'static>>,
+}
+
+impl <T> Stream for FallibleResponseStream<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;
+ if let Poll::Ready(Ok(err)) = receiver.poll_unpin(cx) {
+ return Poll::Ready(Some(Err(err)));
+ };
+
+ match pinned.response_streams.poll_next_unpin(cx) {
Review Comment:
Likewise here, calling `ready!()` can probably simplify this a bit
##########
arrow-flight/src/client.rs:
##########
@@ -704,3 +679,75 @@ impl FlightClient {
request
}
}
+
+struct FallibleRequestStream {
+ sender: Option<Sender<FlightError>>,
+ request_streams: Pin<Box<dyn Stream<Item = Result<FlightData>> + Send +
'static>>,
+}
+
+impl Stream for FallibleRequestStream {
+ type Item = FlightData;
+
+ 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.request_streams.as_mut();
+ match request_streams.poll_next_unpin(cx) {
Review Comment:
I think you can avoid a level of nesting using the `ready!` macro here
so like
```rust
match ready!(request_streams.poll_next_unpin(cx) {
Some(Ok(data)) => ...
Some(Err(e)) => ...
None => ...
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
