feat(reflection): Expose ReflectionService (#2066)

hyperium · Nov 26, 2024 · 94587ce · 94587ce
1 parent bdccf58
commit 94587ce
Show file tree

Hide file tree

Showing 3 changed files with 73 additions and 10 deletions.
diff --git a/tonic-reflection/Cargo.toml b/tonic-reflection/Cargo.toml
@@ -26,6 +26,7 @@ server = ["prost-types", "dep:tokio", "dep:tokio-stream"]
 default = ["server"]
 
 [dependencies]
+pin-project = "1"
 prost = "0.13"
 prost-types = {version = "0.13", optional = true}
 tokio = { version = "1.0", features = ["sync", "rt"], optional = true }

diff --git a/tonic-reflection/src/server/v1.rs b/tonic-reflection/src/server/v1.rs
@@ -1,7 +1,8 @@
-use std::sync::Arc;
+use std::{fmt, sync::Arc};
 
+use pin_project::pin_project;
 use tokio::sync::mpsc;
-use tokio_stream::{wrappers::ReceiverStream, StreamExt};
+use tokio_stream::{wrappers::ReceiverStream, Stream, StreamExt};
 use tonic::{Request, Response, Status, Streaming};
 
 use super::ReflectionServiceState;
@@ -13,14 +14,15 @@ use crate::pb::v1::{
     ServerReflectionResponse, ServiceResponse,
 };
 
+/// An implementation for `ServerReflection`.
 #[derive(Debug)]
-pub(super) struct ReflectionService {
+pub struct ReflectionService {
     state: Arc<ReflectionServiceState>,
 }
 
 #[tonic::async_trait]
 impl ServerReflection for ReflectionService {
-    type ServerReflectionInfoStream = ReceiverStream<Result<ServerReflectionResponse, Status>>;
+    type ServerReflectionInfoStream = ServerReflectionInfoStream;
 
     async fn server_reflection_info(
         &self,
@@ -91,7 +93,9 @@ impl ServerReflection for ReflectionService {
             }
         });
 
-        Ok(Response::new(ReceiverStream::new(resp_rx)))
+        Ok(Response::new(ServerReflectionInfoStream(
+            ReceiverStream::new(resp_rx),
+        )))
     }
 }
 
@@ -102,3 +106,30 @@ impl From<ReflectionServiceState> for ReflectionService {
         }
     }
 }
+
+/// A response stream.
+#[pin_project]
+pub struct ServerReflectionInfoStream(
+    #[pin] ReceiverStream<Result<ServerReflectionResponse, Status>>,
+);
+
+impl Stream for ServerReflectionInfoStream {
+    type Item = Result<ServerReflectionResponse, Status>;
+
+    fn poll_next(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Option<Self::Item>> {
+        self.project().0.poll_next(cx)
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.0.size_hint()
+    }
+}
+
+impl fmt::Debug for ServerReflectionInfoStream {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_tuple("ServerReflectionInfoStream").finish()
+    }
+}
diff --git a/tonic-reflection/src/server/v1alpha.rs b/tonic-reflection/src/server/v1alpha.rs
@@ -1,7 +1,8 @@
-use std::sync::Arc;
+use std::{fmt, sync::Arc};
 
+use pin_project::pin_project;
 use tokio::sync::mpsc;
-use tokio_stream::{wrappers::ReceiverStream, StreamExt};
+use tokio_stream::{wrappers::ReceiverStream, Stream, StreamExt};
 use tonic::{Request, Response, Status, Streaming};
 
 use super::ReflectionServiceState;
@@ -13,14 +14,15 @@ use crate::pb::v1alpha::{
     ServerReflectionResponse, ServiceResponse,
 };
 
+/// An implementation for `ServerReflection`.
 #[derive(Debug)]
-pub(super) struct ReflectionService {
+pub struct ReflectionService {
     state: Arc<ReflectionServiceState>,
 }
 
 #[tonic::async_trait]
 impl ServerReflection for ReflectionService {
-    type ServerReflectionInfoStream = ReceiverStream<Result<ServerReflectionResponse, Status>>;
+    type ServerReflectionInfoStream = ServerReflectionInfoStream;
 
     async fn server_reflection_info(
         &self,
@@ -91,7 +93,9 @@ impl ServerReflection for ReflectionService {
             }
         });
 
-        Ok(Response::new(ReceiverStream::new(resp_rx)))
+        Ok(Response::new(ServerReflectionInfoStream(
+            ReceiverStream::new(resp_rx),
+        )))
     }
 }
 
@@ -102,3 +106,30 @@ impl From<ReflectionServiceState> for ReflectionService {
         }
     }
 }
+
+/// A response stream.
+#[pin_project]
+pub struct ServerReflectionInfoStream(
+    #[pin] ReceiverStream<Result<ServerReflectionResponse, Status>>,
+);
+
+impl Stream for ServerReflectionInfoStream {
+    type Item = Result<ServerReflectionResponse, Status>;
+
+    fn poll_next(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Option<Self::Item>> {
+        self.project().0.poll_next(cx)
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.0.size_hint()
+    }
+}
+
+impl fmt::Debug for ServerReflectionInfoStream {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_tuple("ServerReflectionInfoStream").finish()
+    }
+}