service-profiles: Wrap receiver types (#1038)

olix0r · web-flow · commit f4a2a6b6df9f · 2021-06-09T11:59:39.000-07:00
We should be able to cache profile resolutions independently of the
stack; but the fact that the the profile modue's `Receiver` type is just
an alias for `watch::Receiver&lt;Profile&gt;` means that we can't attach
the additional metadata needed for cache entry retention.

In anticipation of caching profiles by lookup address, this change
creates an explicit `Receiver` type that hides the underlying
`watch::Receiver`, exposing only methods used by other crates/modules.

This change also replaces the boxed watch stream with a
`tokio_stream::wrappers::WatchStream`--also wrapped by a
profile-specific type--to avoid needless allocation and support
retaining additional resources.
diff --git a/Cargo.lock b/Cargo.lock
@@ -1303,7 +1303,6 @@ dependencies = [
 name = "linkerd-service-profiles"
 version = "0.1.0"
 dependencies = [
- "async-stream",
  "bytes",
  "futures",
  "http",
@@ -1323,6 +1322,7 @@ dependencies = [
  "regex",
  "thiserror",
  "tokio",
+ "tokio-stream",
  "tonic",
  "tower",
  "tracing",
diff --git a/linkerd/app/gateway/src/gateway.rs b/linkerd/app/gateway/src/gateway.rs
@@ -63,7 +63,7 @@ where
             None => return Gateway::BadDomain(http.target.name().clone()),
         };
 
-        let logical_addr = match profile.borrow().addr.clone() {
+        let logical_addr = match profile.logical_addr() {
             Some(addr) => addr,
             None => return Gateway::BadDomain(http.target.name().clone()),
         };
diff --git a/linkerd/app/gateway/src/lib.rs b/linkerd/app/gateway/src/lib.rs
@@ -119,7 +119,7 @@ where
                 let profile = p.ok_or_else(|| {
                     DiscoveryRejected::new("no profile discovered for gateway target")
                 })?;
-                let logical_addr = profile.borrow().addr.clone().ok_or_else(|| {
+                let logical_addr = profile.logical_addr().ok_or_else(|| {
                     DiscoveryRejected::new(
                         "profile for gateway target does not have a logical address",
                     )
diff --git a/linkerd/app/inbound/fuzz/Cargo.lock b/linkerd/app/inbound/fuzz/Cargo.lock
@@ -1132,7 +1132,6 @@ dependencies = [
 name = "linkerd-service-profiles"
 version = "0.1.0"
 dependencies = [
- "async-stream",
  "bytes",
  "futures",
  "http",
@@ -1150,6 +1149,7 @@ dependencies = [
  "regex",
  "thiserror",
  "tokio",
+ "tokio-stream",
  "tonic",
  "tower",
  "tracing",
diff --git a/linkerd/app/outbound/src/ingress.rs b/linkerd/app/outbound/src/ingress.rs
@@ -97,8 +97,7 @@ impl Outbound<svc::BoxNewHttp<http::Endpoint>> {
             .push_request_filter(
                 |(profile, http): (Option<profiles::Receiver>, Http<NameAddr>)| {
                     if let Some(profile) = profile {
-                        let addr = profile.borrow().addr.clone();
-                        if let Some(logical_addr) = addr {
+                        if let Some(logical_addr) = profile.logical_addr() {
                             return Ok(http::Logical {
                                 profile,
                                 logical_addr,
diff --git a/linkerd/app/outbound/src/logical.rs b/linkerd/app/outbound/src/logical.rs
@@ -58,7 +58,7 @@ impl<P> svc::Param<LogicalAddr> for Logical<P> {
 // Used for skipping HTTP detection
 impl svc::Param<Option<http::detect::Skip>> for Logical<()> {
     fn param(&self) -> Option<http::detect::Skip> {
-        if self.profile.borrow().opaque_protocol {
+        if self.profile.is_opaque_protocol() {
             Some(http::detect::Skip)
         } else {
             None
diff --git a/linkerd/app/outbound/src/switch_logical.rs b/linkerd/app/outbound/src/switch_logical.rs
@@ -36,28 +36,21 @@ impl<S> Outbound<S> {
             .push_switch(
                 move |(profile, target): (Option<profiles::Receiver>, T)| -> Result<_, Never> {
                     if let Some(rx) = profile {
-                        let profiles::Profile {
-                            ref addr,
-                            ref endpoint,
-                            opaque_protocol,
-                            ..
-                        } = *rx.borrow();
-
                         // If the profile provides an endpoint, then the target is single endpoint and
                         // not a logical/load-balanced service.
-                        if let Some((addr, metadata)) = endpoint.clone() {
+                        if let Some((addr, metadata)) = rx.endpoint() {
                             return Ok(svc::Either::A(Endpoint::from_metadata(
                                 addr,
                                 metadata,
                                 no_tls_reason,
-                                opaque_protocol,
+                                rx.is_opaque_protocol(),
                             )));
                         }
 
                         // Otherwise, if the profile provides a (named) logical address, then we build a
                         // logical stack so we apply routes, traffic splits, and load balancing.
-                        if let Some(logical_addr) = addr.clone() {
-                            return Ok(svc::Either::B(Logical::new(logical_addr, rx.clone())));
+                        if let Some(logical_addr) = rx.logical_addr() {
+                            return Ok(svc::Either::B(Logical::new(logical_addr, rx)));
                         }
                     }
 
@@ -151,7 +144,7 @@ mod tests {
         });
 
         let orig_dst = OrigDstAddr(SocketAddr::new([192, 0, 2, 20].into(), 2020));
-        let svc = stack.new_service((Some(profile), orig_dst));
+        let svc = stack.new_service((Some(profile.into()), orig_dst));
         let (server_io, _client_io) = io::duplex(1);
         svc.oneshot(server_io).await.expect("service must succeed");
     }
@@ -182,7 +175,7 @@ mod tests {
         });
 
         let orig_dst = OrigDstAddr(SocketAddr::new([192, 0, 2, 20].into(), 2020));
-        let svc = stack.new_service((Some(profile), orig_dst));
+        let svc = stack.new_service((Some(profile.into()), orig_dst));
         let (server_io, _client_io) = io::duplex(1);
         svc.oneshot(server_io).await.expect("service must succeed");
     }
diff --git a/linkerd/app/outbound/src/tcp/logical.rs b/linkerd/app/outbound/src/tcp/logical.rs
@@ -140,12 +140,12 @@ mod tests {
 
         // We create a logical target to be resolved to endpoints.
         let logical_addr = LogicalAddr("xyz.example.com:4444".parse().unwrap());
-        let (_tx, profile) = tokio::sync::watch::channel(Profile {
+        let (_tx, rx) = tokio::sync::watch::channel(Profile {
             addr: Some(logical_addr.clone()),
             ..Default::default()
         });
         let logical = Logical {
-            profile,
+            profile: rx.into(),
             logical_addr: logical_addr.clone(),
             protocol: (),
         };
@@ -203,12 +203,12 @@ mod tests {
 
         // We create a logical target to be resolved to endpoints.
         let logical_addr = LogicalAddr("xyz.example.com:4444".parse().unwrap());
-        let (_tx, profile) = tokio::sync::watch::channel(Profile {
+        let (_tx, rx) = tokio::sync::watch::channel(Profile {
             addr: Some(logical_addr.clone()),
             ..Default::default()
         });
         let logical = Logical {
-            profile,
+            profile: rx.into(),
             logical_addr: logical_addr.clone(),
             protocol: (),
         };
diff --git a/linkerd/app/test/src/profile.rs b/linkerd/app/test/src/profile.rs
@@ -16,7 +16,7 @@ pub fn only_default() -> Receiver {
 pub fn only(profile: Profile) -> Receiver {
     let (tx, rx) = channel(profile);
     tokio::spawn(async move { tx.closed().await });
-    rx
+    rx.into()
 }
 
 pub fn resolver() -> crate::resolver::Profiles {
diff --git a/linkerd/app/test/src/resolver.rs b/linkerd/app/test/src/resolver.rs
@@ -158,7 +158,7 @@ impl Profiles {
             .endpoints
             .lock()
             .unwrap()
-            .insert(addr.into(), Some(rx));
+            .insert(addr.into(), Some(rx.into()));
         tx
     }
 
@@ -169,7 +169,7 @@ impl Profiles {
             .endpoints
             .lock()
             .unwrap()
-            .insert(addr.into(), Some(rx));
+            .insert(addr.into(), Some(rx.into()));
         self
     }
 
diff --git a/linkerd/service-profiles/Cargo.toml b/linkerd/service-profiles/Cargo.toml
@@ -25,7 +25,7 @@ linkerd-tonic-watch = { path = "../tonic-watch" }
 rand = { version = "0.8", features = ["small_rng"] }
 regex = "1.0.0"
 tokio = { version = "1", features = ["macros", "rt", "sync", "time"] }
-async-stream = "0.3"
+tokio-stream = { version = "0.1", features = ["sync"] }
 tonic = { version = "0.4", default-features = false }
 tower = { version = "0.4.7", features = [ "ready-cache", "retry", "util"] }
 thiserror = "1"
diff --git a/linkerd/service-profiles/src/client.rs b/linkerd/service-profiles/src/client.rs
@@ -73,7 +73,10 @@ where
         let w = self.watch.clone();
         Box::pin(async move {
             match w.spawn_watch(addr).await {
-                Ok(rsp) => Ok(Some(rsp.into_inner())),
+                Ok(rsp) => {
+                    let rx = rsp.into_inner();
+                    Ok(Some(rx.into()))
+                }
                 Err(status) => {
                     debug!(%status, "Ignoring profile");
                     Ok::<_, Never>(None)
diff --git a/linkerd/service-profiles/src/http/mod.rs b/linkerd/service-profiles/src/http/mod.rs
@@ -1,4 +1,3 @@
-use crate::Receiver;
 use indexmap::IndexMap;
 use regex::Regex;
 use std::{
diff --git a/linkerd/service-profiles/src/http/route_request.rs b/linkerd/service-profiles/src/http/route_request.rs
@@ -1,12 +1,11 @@
-use super::{Receiver, RequestMatch, Route};
-use crate::Profile;
-use futures::{future::ErrInto, prelude::*, ready, Stream};
+use super::{RequestMatch, Route};
+use crate::{Profile, Receiver, ReceiverStream};
+use futures::{future::ErrInto, prelude::*, ready};
 use linkerd_error::Error;
 use linkerd_stack::{layer, NewService, Param, Proxy};
 use std::{
     collections::HashMap,
     marker::PhantomData,
-    pin::Pin,
     task::{Context, Poll},
 };
 use tracing::{debug, trace};
@@ -34,7 +33,7 @@ pub struct NewRouteRequest<M, N, R> {
 
 pub struct RouteRequest<T, S, N, R> {
     target: T,
-    rx: Pin<Box<dyn Stream<Item = Profile> + Send + Sync>>,
+    rx: ReceiverStream,
     inner: S,
     new_route: N,
     http_routes: Vec<(RequestMatch, Route)>,
@@ -62,13 +61,13 @@ where
     type Service = RouteRequest<T, M::Service, N, N::Service>;
 
     fn new_service(&mut self, target: T) -> Self::Service {
-        let rx = crate::stream_profile(target.param());
+        let rx = target.param();
         let inner = self.inner.new_service(target.clone());
         let default = self
             .new_route
             .new_service((self.default.clone(), target.clone()));
         RouteRequest {
-            rx,
+            rx: rx.into(),
             target,
             inner,
             default,
@@ -94,7 +93,7 @@ where
 
     fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
         let mut update = None;
-        while let Poll::Ready(Some(up)) = self.rx.as_mut().poll_next(cx) {
+        while let Poll::Ready(Some(up)) = self.rx.poll_next_unpin(cx) {
             tracing::trace!(update = ?up, "updated profile");
             update = Some(up);
         }
diff --git a/linkerd/service-profiles/src/lib.rs b/linkerd/service-profiles/src/lib.rs
@@ -2,7 +2,7 @@
 #![forbid(unsafe_code)]
 #![allow(clippy::inconsistent_struct_constructor)]
 
-use futures::stream::Stream;
+use futures::Stream;
 use linkerd_addr::{Addr, NameAddr};
 use linkerd_error::Error;
 use linkerd_proxy_api_resolve::Metadata;
@@ -15,6 +15,7 @@ use std::{
     task::{Context, Poll},
 };
 use thiserror::Error;
+use tokio::sync::watch;
 use tower::util::{Oneshot, ServiceExt};
 
 mod client;
@@ -26,7 +27,15 @@ pub mod split;
 
 pub use self::client::Client;
 
-pub type Receiver = tokio::sync::watch::Receiver<Profile>;
+#[derive(Clone, Debug)]
+pub struct Receiver {
+    inner: tokio::sync::watch::Receiver<Profile>,
+}
+
+#[derive(Debug)]
+struct ReceiverStream {
+    inner: tokio_stream::wrappers::WatchStream<Profile>,
+}
 
 #[derive(Clone, Debug, Default)]
 pub struct Profile {
@@ -113,20 +122,48 @@ where
     }
 }
 
-fn stream_profile(mut rx: Receiver) -> Pin<Box<dyn Stream<Item = Profile> + Send + Sync>> {
-    Box::pin(async_stream::stream! {
-        loop {
-            let val = rx.borrow().clone();
-            yield val;
-            // This is a loop with a return condition rather than a while loop,
-            // because we want to yield the *first* value immediately, rather
-            // than waiting for the profile to change again.
-            if let Err(_) = rx.changed().await {
-                tracing::trace!("profile sender dropped");
-                return;
-            }
-        }
-    })
+// === impl Receiver ===
+
+impl From<watch::Receiver<Profile>> for Receiver {
+    fn from(inner: watch::Receiver<Profile>) -> Self {
+        Self { inner }
+    }
+}
+
+impl Receiver {
+    pub fn logical_addr(&self) -> Option<LogicalAddr> {
+        self.inner.borrow().addr.clone()
+    }
+
+    pub fn is_opaque_protocol(&self) -> bool {
+        self.inner.borrow().opaque_protocol
+    }
+
+    pub fn endpoint(&self) -> Option<(SocketAddr, Metadata)> {
+        self.inner.borrow().endpoint.clone()
+    }
+
+    fn targets(&self) -> Vec<Target> {
+        self.inner.borrow().targets.clone()
+    }
+}
+
+// === impl ReceiverStream ===
+
+impl From<Receiver> for ReceiverStream {
+    fn from(Receiver { inner }: Receiver) -> Self {
+        let inner = tokio_stream::wrappers::WatchStream::new(inner);
+        ReceiverStream { inner }
+    }
+}
+
+impl Stream for ReceiverStream {
+    type Item = Profile;
+
+    #[inline]
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        Pin::new(&mut self.inner).poll_next(cx)
+    }
 }
 
 // === impl LookupAddr ===
diff --git a/linkerd/service-profiles/src/split.rs b/linkerd/service-profiles/src/split.rs
@@ -1,4 +1,4 @@
-use crate::{LogicalAddr, Profile, Receiver, Target};
+use crate::{LogicalAddr, Profile, Receiver, ReceiverStream, Target};
 use futures::{prelude::*, ready};
 use indexmap::IndexSet;
 use linkerd_addr::NameAddr;
@@ -32,7 +32,7 @@ pub struct NewSplit<N, S, Req> {
 
 pub struct Split<T, N, S, Req> {
     rng: SmallRng,
-    rx: Pin<Box<dyn Stream<Item = Profile> + Send + Sync>>,
+    rx: ReceiverStream,
     target: T,
     new_service: N,
     distribution: WeightedIndex<u32>,
@@ -62,7 +62,7 @@ where
 
     fn new_service(&mut self, target: T) -> Self::Service {
         let rx: Receiver = target.param();
-        let mut targets = rx.borrow().targets.clone();
+        let mut targets = rx.targets();
         if targets.is_empty() {
             let LogicalAddr(addr) = target.param();
             targets.push(Target { addr, weight: 1 })
@@ -83,7 +83,7 @@ where
         }
 
         Split {
-            rx: crate::stream_profile(rx),
+            rx: rx.into(),
             target,
             new_service,
             services,
@@ -112,8 +112,8 @@ where
 
     fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
         let mut update = None;
-        while let Poll::Ready(Some(up)) = self.rx.as_mut().poll_next(cx) {
-            update = Some(up.clone());
+        while let Poll::Ready(Some(up)) = self.rx.poll_next_unpin(cx) {
+            update = Some(up);
         }
 
         // Every time the profile updates, rebuild the distribution, reusing

Original file line number	Diff line number	Diff line change
`@@ -119,7 +119,7 @@ where`
`119`	`119`	`let profile = p.ok_or_else(\|\| {`
`120`	`120`	`DiscoveryRejected::new("no profile discovered for gateway target")`
`121`	`121`	`})?;`
`122`		`- let logical_addr = profile.borrow().addr.clone().ok_or_else(\|\| {`
	`122`	`+ let logical_addr = profile.logical_addr().ok_or_else(\|\| {`
`123`	`123`	`DiscoveryRejected::new(`
`124`	`124`	`"profile for gateway target does not have a logical address",`
`125`	`125`	`)`
Original file line number	Diff line number	Diff line change
`@@ -16,7 +16,7 @@ pub fn only_default() -> Receiver {`
`16`	`16`	`pub fn only(profile: Profile) -> Receiver {`
`17`	`17`	`let (tx, rx) = channel(profile);`
`18`	`18`	`tokio::spawn(async move { tx.closed().await });`
`19`		`- rx`
	`19`	`+ rx.into()`
`20`	`20`	`}`
`21`	`21`
`22`	`22`	`pub fn resolver() -> crate::resolver::Profiles {`