Move resolve api to async stream

Signed-off-by: Zahari Dichev <zaharidichev@gmail.com>

Author: zaharidichev

Cargo.lock

@@ -1213,6 +1213,7 @@ dependencies = [
 name = "linkerd2-proxy-api-resolve"
 version = "0.1.0"
 dependencies = [
+ "async-stream",
  "futures 0.3.5",
  "http 0.2.1",
  "http-body",

linkerd/proxy/api-resolve/Cargo.toml

@@ -9,6 +9,7 @@ Implements the Resolve trait using the proxy's gRPC API
 """
 
 [dependencies]
+async-stream = "0.2.1"
 futures = "0.3"
 linkerd2-identity = { path = "../../identity" }
 linkerd2-proxy-api = { git = "https://github.com/linkerd/linkerd2-proxy-api", tag = "v0.1.13" }

linkerd/proxy/api-resolve/src/lib.rs

@@ -1,4 +1,5 @@
 #![deny(warnings, rust_2018_idioms)]
+#![recursion_limit = "512"]
 
 use linkerd2_identity as identity;
 use linkerd2_proxy_api as api;

linkerd/proxy/api-resolve/src/resolve.rs

@@ -3,11 +3,12 @@ use crate::core::resolve::{self, Update};
 use crate::metadata::Metadata;
 use crate::pb;
 use api::destination_client::DestinationClient;
-use futures::{ready, Stream};
+use async_stream::try_stream;
+use futures::future;
+use futures::stream::StreamExt;
+use futures::Stream;
 use http_body::Body as HttpBody;
-use pin_project::pin_project;
 use std::error::Error;
-use std::future::Future;
 use std::pin::Pin;
 use std::task::{Context, Poll};
 use tonic::{
@@ -25,13 +26,7 @@ pub struct Resolve<S> {
     context_token: String,
 }
 
-#[pin_project]
-pub struct Resolution {
-    #[pin]
-    inner: grpc::Streaming<api::Update>,
-}
-
-// === impl Resolver ===
+// === impl Resolve ===
 
 impl<S> Resolve<S>
 where
@@ -65,6 +60,9 @@ where
     }
 }
 
+type UpdatesStream =
+    Pin<Box<dyn Stream<Item = Result<Update<Metadata>, grpc::Status>> + Send + 'static>>;
+
 impl<T, S> Service<T> for Resolve<S>
 where
     T: ToString,
@@ -75,10 +73,9 @@ where
     <S::ResponseBody as HttpBody>::Error: Into<Box<dyn Error + Send + Sync + 'static>> + Send,
     S::Future: Send,
 {
-    type Response = Resolution;
+    type Response = UpdatesStream;
     type Error = grpc::Status;
-    type Future =
-        Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;
+    type Future = futures::future::Ready<Result<Self::Response, Self::Error>>;
 
     fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
         // The future returned by the Tonic generated `DestinationClient`'s `get` method will drive the service to readiness before calling it, so we can always return `Ready` here.
@@ -88,69 +85,72 @@ where
     fn call(&mut self, target: T) -> Self::Future {
         let path = target.to_string();
         debug!(dst = %path, context = %self.context_token, "Resolving");
-        let mut svc = self.service.clone();
         let req = api::GetDestination {
             path,
             scheme: self.scheme.clone(),
             context_token: self.context_token.clone(),
         };
-        Box::pin(async move {
-            let rsp = svc.get(grpc::Request::new(req)).await?;
-            trace!(metadata = ?rsp.metadata());
-            Ok(Resolution {
-                inner: rsp.into_inner(),
-            })
-        })
+
+        future::ok(Box::pin(resolution(self.service.clone(), req)))
     }
 }
 
-impl Stream for Resolution {
-    type Item = Result<resolve::Update<Metadata>, grpc::Status>;
+fn resolution<S>(
+    mut client: DestinationClient<S>,
+    req: api::GetDestination,
+) -> impl Stream<Item = Result<resolve::Update<Metadata>, grpc::Status>>
+where
+    S: GrpcService<BoxBody> + Clone + Send + 'static,
+    S::Error: Into<Box<dyn Error + Send + Sync>> + Send,
+    S::ResponseBody: Send,
+    <S::ResponseBody as Body>::Data: Send,
+    <S::ResponseBody as HttpBody>::Error: Into<Box<dyn Error + Send + Sync + 'static>> + Send,
+    S::Future: Send,
+{
+    try_stream! {
+        let rsp = client.get(grpc::Request::new(req)).await?;
+        trace!(metadata = ?rsp.metadata());
+        let mut stream = rsp.into_inner();
 
-    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        let mut this = self.project();
-        loop {
-            match ready!(this.inner.as_mut().poll_next(cx)) {
-                Some(update) => match update?.update {
-                    Some(api::update::Update::Add(api::WeightedAddrSet {
-                        addrs,
-                        metric_labels,
-                    })) => {
-                        let addr_metas = addrs
-                            .into_iter()
-                            .filter_map(|addr| pb::to_addr_meta(addr, &metric_labels))
-                            .collect::<Vec<_>>();
-                        if !addr_metas.is_empty() {
-                            debug!(endpoints = %addr_metas.len(), "Add");
-                            return Poll::Ready(Some(Ok(Update::Add(addr_metas))));
-                        }
+        while let Some(update) = stream.next().await {
+            match update?.update {
+                Some(api::update::Update::Add(api::WeightedAddrSet {
+                    addrs,
+                    metric_labels,
+                })) => {
+                    let addr_metas = addrs
+                        .into_iter()
+                        .filter_map(|addr| pb::to_addr_meta(addr, &metric_labels))
+                        .collect::<Vec<_>>();
+                    if !addr_metas.is_empty() {
+                        debug!(endpoints = %addr_metas.len(), "Add");
+                        yield Update::Add(addr_metas);
                     }
+                }
 
-                    Some(api::update::Update::Remove(api::AddrSet { addrs })) => {
-                        let sock_addrs = addrs
-                            .into_iter()
-                            .filter_map(pb::to_sock_addr)
-                            .collect::<Vec<_>>();
-                        if !sock_addrs.is_empty() {
-                            debug!(endpoints = %sock_addrs.len(), "Remove");
-                            return Poll::Ready(Some(Ok(Update::Remove(sock_addrs))));
-                        }
+                Some(api::update::Update::Remove(api::AddrSet { addrs })) => {
+                    let sock_addrs = addrs
+                        .into_iter()
+                        .filter_map(pb::to_sock_addr)
+                        .collect::<Vec<_>>();
+                    if !sock_addrs.is_empty() {
+                        debug!(endpoints = %sock_addrs.len(), "Remove");
+                        yield Update::Remove(sock_addrs);
                     }
+                }
 
-                    Some(api::update::Update::NoEndpoints(api::NoEndpoints { exists })) => {
-                        info!("No endpoints");
-                        let update = if exists {
-                            Update::Empty
-                        } else {
-                            Update::DoesNotExist
-                        };
-                        return Poll::Ready(Some(Ok(update.into())));
-                    }
+                Some(api::update::Update::NoEndpoints(api::NoEndpoints { exists })) => {
+                    info!("No endpoints");
+                    let update = if exists {
+                        Update::Empty
+                    } else {
+                        Update::DoesNotExist
+                    };
+                    yield update.into();
+                }
 
-                    None => {} // continue
-                },
-                None => return Poll::Ready(None),
-            };
+                None => {} // continue
+            }
         }
     }
 }