tmp

Signed-off-by: Matt Klein <[email protected]>

Author: mattklein123

Cargo.lock

@@ -48,7 +48,7 @@ arbitrary       = { version = "1.4.1", features = ["derive"] }
 arc-swap        = "1.7.1"
 assert_matches  = "1.5.0"
 async-trait     = "0.1.85"
-axum            = { version = "0.8.1", features = ["http2", "macros"] }
+axum            = { version = "0.8.1", features = ["http2", "macros", "ws"] }
 axum-server     = { version = "0.7.1", features = ["tls-rustls-no-provider"] }
 backoff         = "0.4.0"
 base64          = "0.22.1"

Cargo.toml

@@ -11,13 +11,7 @@ use crate::error::{Error, Result};
 use crate::service::ServiceMethod;
 use crate::status::{Code, Status};
 use crate::{
-  finalize_decompression,
-  StreamingApi,
-  StreamingApiReceiver,
-  CONNECT_PROTOCOL_VERSION,
-  CONTENT_ENCODING_SNAPPY,
-  GRPC_STATUS,
-  TRANSFER_ENCODING_TRAILERS,
+  finalize_decompression, BodyFramer, SingleFrame, StreamingApi, StreamingApiReceiver, CONNECT_PROTOCOL_VERSION, CONTENT_ENCODING_SNAPPY, GRPC_STATUS, TRANSFER_ENCODING_TRAILERS
 };
 use assert_matches::debug_assert_matches;
 use axum::body::Body;
@@ -32,8 +26,10 @@ use bd_grpc_codec::{
   GRPC_ENCODING_HEADER,
 };
 use bd_time::TimeDurationExt;
+use futures::StreamExt;
 use http::header::{CONTENT_ENCODING, CONTENT_TYPE, TRANSFER_ENCODING};
 use http::{HeaderMap, Uri};
+use http_body::Frame;
 use http_body_util::{BodyExt, StreamBody};
 use hyper::body::Incoming;
 use hyper_util::client::legacy::connect::{Connect, HttpConnector};
@@ -296,9 +292,13 @@ impl<C: Connect + Clone + Send + Sync + 'static> Client<C> {
     validate: bool,
     compression: Option<bd_grpc_codec::Compression>,
     optimize_for: OptimizeFor,
-  ) -> Result<StreamingApi<OutgoingType, IncomingType>> {
+  ) -> Result<StreamingApi<OutgoingType, IncomingType, BodyFramer>> {
     let (tx, rx) = mpsc::channel(1);
-    let body = StreamBody::new(ReceiverStream::new(rx));
+    let body = StreamBody::new(ReceiverStream::new(rx).map(|f| match f {
+      Ok(SingleFrame::Data(data)) => Ok(Frame::data(data)),
+      Ok(SingleFrame::Trailers(trailers)) => Ok(Frame::trailers(trailers)),
+      Err(e) => Err(e),
+    }));
 
     match compression {
       None => {},
@@ -347,7 +347,7 @@ impl<C: Connect + Clone + Send + Sync + 'static> Client<C> {
     validate: bool,
     optimize_for: OptimizeFor,
     connect_protocol: Option<ConnectProtocolType>,
-  ) -> Result<StreamingApiReceiver<IncomingType>> {
+  ) -> Result<StreamingApiReceiver<IncomingType, BodyFramer>> {
     let mut encoder = Encoder::new(None);
     let response = self
       .common_request(
@@ -360,7 +360,9 @@ impl<C: Connect + Clone + Send + Sync + 'static> Client<C> {
     let (parts, body) = response.into_parts();
     Ok(StreamingApiReceiver::new(
       parts.headers,
-      Body::new(body),
+      BodyFramer {
+        body: Body::new(body),
+      },
       validate,
       optimize_for,
       None,

bd-grpc/src/client.rs

@@ -10,22 +10,7 @@ use crate::compression::{Compression, ConnectSafeCompressionLayer};
 use crate::connect_protocol::ConnectProtocolType;
 use crate::generated::proto::test::{EchoRequest, EchoResponse};
 use crate::{
-  make_server_streaming_router,
-  make_unary_router,
-  new_grpc_response,
-  Code,
-  Error,
-  Handler,
-  Result,
-  ServerStreamingHandler,
-  ServiceMethod,
-  Status,
-  StreamStats,
-  StreamingApi,
-  StreamingApiSender,
-  CONNECT_PROTOCOL_VERSION,
-  CONTENT_TYPE,
-  CONTENT_TYPE_PROTO,
+  make_server_streaming_router, make_unary_router, new_grpc_response, BodyFramer, Code, Error, Handler, Result, ServerStreamingHandler, ServiceMethod, SingleFrame, Status, StreamStats, StreamingApi, StreamingApiSender, CONNECT_PROTOCOL_VERSION, CONTENT_TYPE, CONTENT_TYPE_PROTO
 };
 use assert_matches::assert_matches;
 use async_trait::async_trait;
@@ -39,8 +24,9 @@ use bd_server_stats::stats::CounterWrapper;
 use bd_server_stats::test::util::stats::{self, Helper};
 use bd_time::TimeDurationExt;
 use bytes::Bytes;
-use futures::poll;
+use futures::{poll, StreamExt};
 use http::{Extensions, HeaderMap};
+use http_body::Frame;
 use http_body_util::StreamBody;
 use parking_lot::Mutex;
 use prometheus::labels;
@@ -389,11 +375,17 @@ async fn read_stop() {
       let (parts, body) = request.into_parts();
       let (response_sender, response_body) = mpsc::channel(1);
       let response = new_grpc_response(
-        Body::new(StreamBody::new(ReceiverStream::new(response_body))),
+        Body::new(StreamBody::new(ReceiverStream::new(response_body).map(
+          |f| match f {
+            Ok(SingleFrame::Data(data)) => Ok(Frame::data(data)),
+            Ok(SingleFrame::Trailers(trailers)) => Ok(Frame::trailers(trailers)),
+            Err(e) => Err(e),
+          },
+        ))),
         None,
         None,
       );
-      let api = StreamingApi::<EchoResponse, EchoRequest>::new(
+      let api = StreamingApi::<EchoResponse, EchoRequest, BodyFramer>::new(
         response_sender,
         parts.headers,
         body,