cache: Only spawn a single task per cache entry (#770)

olix0r · web-flow · commit cf9900de443a · 2020-12-12T14:59:38.000-08:00
In b0cf2fb, the proxy moved eviction to be controlled by a background task; but this scheme spawned a task every time a service was dropped. This change modifies the cache to only spawn a single task per cached service and use a notification channel to signal drops. This reduces memory pressure, especially for HTTP caches.
diff --git a/linkerd/app/core/src/svc.rs b/linkerd/app/core/src/svc.rs
@@ -8,8 +8,10 @@ use linkerd2_concurrency_limit as concurrency_limit;
 pub use linkerd2_stack::{self as stack, layer, NewService};
 pub use linkerd2_stack_tracing::{InstrumentMake, InstrumentMakeLayer};
 pub use linkerd2_timeout as timeout;
-use std::task::{Context, Poll};
-use std::time::Duration;
+use std::{
+    task::{Context, Poll},
+    time::Duration,
+};
 use tower::layer::util::{Identity, Stack as Pair};
 pub use tower::layer::Layer;
 pub use tower::make::MakeService;
@@ -229,7 +231,7 @@ impl<S> Stack<S> {
 
     pub fn push_cache<T>(self, idle: Duration) -> Stack<cache::Cache<T, S>>
     where
-        T: Eq + std::hash::Hash + Send + Sync + 'static,
+        T: Clone + Eq + std::fmt::Debug + std::hash::Hash + Send + Sync + 'static,
         S: NewService<T> + 'static,
         S::Service: Send + Sync + 'static,
     {
diff --git a/linkerd/cache/Cargo.toml b/linkerd/cache/Cargo.toml
@@ -10,7 +10,7 @@ futures = "0.3"
 linkerd2-error = { path = "../error" }
 linkerd2-stack = { path = "../stack" }
 parking_lot = "0.11"
-tokio = { version = "0.3", default-features = false, features = ["rt", "time"] }
+tokio = { version = "0.3", default-features = false, features = ["rt", "sync", "time"] }
 tower = { version = "0.4", default-features = false, features = ["util"] }
 tracing = "0.1.22"
 tracing-futures = "0.2"
diff --git a/linkerd/cache/src/lib.rs b/linkerd/cache/src/lib.rs
@@ -1,5 +1,6 @@
 #![deny(warnings, rust_2018_idioms)]
 
+use futures::prelude::*;
 use linkerd2_stack::{layer, NewService};
 use parking_lot::RwLock;
 use std::{
@@ -8,9 +9,8 @@ use std::{
     sync::{Arc, Weak},
     task::{Context, Poll},
 };
-use tokio::time;
-use tracing::{debug, debug_span, trace};
-use tracing_futures::Instrument;
+use tokio::{sync::Notify, time};
+use tracing::{debug, instrument, trace};
 
 #[derive(Clone)]
 pub struct Cache<T, N>
@@ -19,72 +19,125 @@ where
     N: NewService<T>,
 {
     inner: N,
-    services: Services<T, N::Service>,
+    services: Arc<Services<T, N::Service>>,
     idle: time::Duration,
 }
 
 #[derive(Clone, Debug)]
-pub struct Cached<S, T>
+pub struct Cached<S>
 where
     S: Send + Sync + 'static,
-    T: Eq + Hash + Send + Sync + 'static,
 {
     inner: S,
-    handle: Option<Handle<S, T>>,
-    idle: time::Duration,
-}
-
-#[derive(Debug)]
-struct Handle<S, T> {
-    handle: Arc<T>,
-    cache: Weak<RwLock<HashMap<T, (S, Weak<T>)>>>,
+    // Notifies entry's eviction task that a drop has occurred.
+    handle: Arc<Notify>,
 }
 
-type Services<T, S> = Arc<RwLock<HashMap<T, (S, Weak<T>)>>>;
+type Services<T, S> = RwLock<HashMap<T, (S, Weak<Notify>)>>;
 
 // === impl Cache ===
 
 impl<T, N> Cache<T, N>
 where
-    T: Eq + Hash + Send + Sync + 'static,
+    T: Clone + std::fmt::Debug + Eq + Hash + Send + Sync + 'static,
     N: NewService<T> + 'static,
     N::Service: Send + Sync + 'static,
 {
     pub fn layer(idle: time::Duration) -> impl layer::Layer<N, Service = Self> + Clone {
         layer::mk(move |inner| Self::new(idle, inner))
     }
 
-    pub fn new(idle: time::Duration, inner: N) -> Self {
-        let services = Services::default();
-
+    fn new(idle: time::Duration, inner: N) -> Self {
+        let services = Arc::new(Services::default());
         Self {
             inner,
             services,
             idle,
         }
     }
+
+    fn spawn_idle(
+        target: T,
+        idle: time::Duration,
+        cache: &Arc<Services<T, N::Service>>,
+    ) -> Arc<Notify> {
+        // Spawn a background task that holds the handle. Every time the handle
+        // is notified, it resets the idle timeout. Every time teh idle timeout
+        // expires, the handle is checked and the service is dropped if there
+        // are no active handles.
+        let handle = Arc::new(Notify::new());
+        tokio::spawn(Self::evict(
+            target,
+            idle,
+            handle.clone(),
+            Arc::downgrade(&cache),
+        ));
+        handle
+    }
+
+    #[instrument(level = "debug", skip(idle, reset, cache))]
+    async fn evict(
+        target: T,
+        idle: time::Duration,
+        mut reset: Arc<Notify>,
+        cache: Weak<Services<T, N::Service>>,
+    ) {
+        // Wait for the handle to be notified before starting to track idleness.
+        reset.notified().await;
+        debug!("Awaiting idleness");
+
+        // Wait for either the reset to be notified or the idle timeout to
+        // elapse.
+        loop {
+            futures::select_biased! {
+                // If the reset was notified, restart the timer.
+                _ = reset.notified().fuse() => {
+                    trace!("Reset");
+                }
+                _ = time::sleep(idle).fuse() => match cache.upgrade() {
+                    Some(cache) => match Arc::try_unwrap(reset) {
+                        // If this is the last reference to the handle after the
+                        // idle timeout, remove the cache entry.
+                        Ok(_) => {
+                            let removed = cache.write().remove(&target).is_some();
+                            debug_assert!(removed, "Cache item must exist: {:?}", target);
+                            debug!("Cache entry dropped");
+                            return;
+                        }
+                        // Otherwise, another handle has been acquired, so
+                        // restore our reset reference for the next iteration.
+                        Err(r) => {
+                            trace!("The handle is still active");
+                            reset = r;
+                        }
+                    },
+                    None => {
+                        trace!("Cache already dropped");
+                        return;
+                    }
+                },
+            }
+        }
+    }
 }
 
 impl<T, N> NewService<T> for Cache<T, N>
 where
-    T: Clone + Eq + Hash + Send + Sync + 'static,
-    N: NewService<T>,
+    T: Clone + std::fmt::Debug + Eq + Hash + Send + Sync + 'static,
+    N: NewService<T> + 'static,
     N::Service: Clone + Send + Sync + 'static,
 {
-    type Service = Cached<N::Service, T>;
-
-    fn new_service(&mut self, target: T) -> Cached<N::Service, T> {
-        let cache = Arc::downgrade(&self.services);
+    type Service = Cached<N::Service>;
 
+    fn new_service(&mut self, target: T) -> Cached<N::Service> {
         // We expect the item to be available in most cases, so initially obtain
         // only a read lock.
         if let Some((svc, weak)) = self.services.read().get(&target) {
             if let Some(handle) = weak.upgrade() {
                 trace!("Using cached service");
                 return Cached {
                     inner: svc.clone(),
-                    idle: self.idle,
-                    handle: Some(Handle { cache, handle }),
+                    handle,
                 };
             }
         }
@@ -97,58 +150,37 @@ where
                 let (svc, weak) = entry.get();
                 match weak.upgrade() {
                     Some(handle) => {
-                        trace!("Using cached service");
+                        trace!(?target, "Using cached service");
                         Cached {
                             inner: svc.clone(),
-                            idle: self.idle,
-                            handle: Some(Handle { cache, handle }),
+                            handle,
                         }
                     }
                     None => {
-                        debug!("Replacing defunct service");
-                        let inner = self.inner.new_service(target.clone());
-                        let handle = Arc::new(target);
+                        debug!(?target, "Replacing defunct service");
+                        let handle = Self::spawn_idle(target.clone(), self.idle, &self.services);
+                        let inner = self.inner.new_service(target);
                         entry.insert((inner.clone(), Arc::downgrade(&handle)));
-                        Cached {
-                            inner,
-                            idle: self.idle,
-                            handle: Some(Handle { cache, handle }),
-                        }
+                        Cached { inner, handle }
                     }
                 }
             }
             Entry::Vacant(entry) => {
-                debug!("Caching new service");
-                let inner = self.inner.new_service(target.clone());
-                let handle = Arc::new(target);
+                debug!(?target, "Caching new service");
+                let handle = Self::spawn_idle(target.clone(), self.idle, &self.services);
+                let inner = self.inner.new_service(target);
                 entry.insert((inner.clone(), Arc::downgrade(&handle)));
-                Cached {
-                    inner,
-                    idle: self.idle,
-                    handle: Some(Handle { cache, handle }),
-                }
+                Cached { inner, handle }
             }
         }
     }
 }
 
-// === impl Handle ===
-
-impl<S, T> Clone for Handle<S, T> {
-    fn clone(&self) -> Self {
-        Self {
-            cache: self.cache.clone(),
-            handle: self.handle.clone(),
-        }
-    }
-}
-
 // === impl Cached ===
 
-impl<Req, S, T> tower::Service<Req> for Cached<S, T>
+impl<Req, S> tower::Service<Req> for Cached<S>
 where
     S: tower::Service<Req> + Send + Sync + 'static,
-    T: Eq + Hash + Send + Sync + 'static,
 {
     type Response = S::Response;
     type Error = S::Error;
@@ -165,66 +197,35 @@ where
     }
 }
 
-impl<S, T> Drop for Cached<S, T>
+impl<S> Drop for Cached<S>
 where
     S: Send + Sync + 'static,
-    T: Eq + Hash + Send + Sync + 'static,
 {
     fn drop(&mut self) {
-        // If the eviction task is still running, wait for an idle timeout,
-        // retaining the handle for this service. If there are no new instances
-        // holding the handle after that timeout, drop the last strong reference
-        // and signal the eviction task.
-        trace!("Dropping cached service");
-        if let Some(Handle { cache, handle }) = self.handle.take() {
-            let idle = self.idle;
-            tokio::spawn(
-                async move {
-                    trace!(?idle, "Waiting for timeout to elapse");
-                    time::sleep(idle).await;
-                    if let Some(cache) = cache.upgrade() {
-                        // Only evict the service if we can claim the target
-                        // from the handle, ensuring that there are no other
-                        // strong references to the handle.
-                        if let Ok(target) = Arc::try_unwrap(handle) {
-                            trace!("Evicting target");
-                            let mut services = cache.write();
-                            if services.remove(&target).is_some() {
-                                trace!("Evicted");
-                            } else {
-                                trace!("Service already evicted");
-                            }
-                        } else {
-                            trace!("Handle acquired by another instance");
-                        }
-                    } else {
-                        trace!("Cache dropped");
-                    }
-                }
-                .instrument(debug_span!("evict")),
-            );
-        }
+        self.handle.notify_one();
     }
 }
 
 #[cfg(test)]
 #[tokio::test]
-async fn cached_idle_retain() {
+async fn test_idle_retain() {
     let _ = tracing_subscriber::fmt::try_init();
     time::pause();
 
-    let cache = Services::default();
-    let handle = Arc::new(());
+    let idle = time::Duration::from_secs(10);
+    let cache = Arc::new(Services::default());
+
+    let handle = Cache::<(), fn(()) -> ()>::spawn_idle((), idle, &cache);
     cache.write().insert((), ((), Arc::downgrade(&handle)));
-    let c0 = Cached {
-        inner: (),
-        idle: time::Duration::from_secs(10),
-        handle: Some(Handle {
-            handle,
-            cache: Arc::downgrade(&cache),
-        }),
-    };
-    let handle = Arc::downgrade(&c0.handle.as_ref().unwrap().handle);
+    let c0 = Cached { inner: (), handle };
+
+    let handle = Arc::downgrade(&c0.handle);
+
+    // Let an idle timeout elapse and ensured the held service has not been
+    // evicted.
+    time::sleep(idle * 2).await;
+    assert!(handle.upgrade().is_some());
+    assert!(cache.read().contains_key(&()));
 
     // Drop the original cached instance and elapse only half of the idle
     // timeout.
@@ -237,11 +238,8 @@ async fn cached_idle_retain() {
     // new cached instance.
     let c1 = Cached {
         inner: (),
-        idle: time::Duration::from_secs(10),
-        handle: Some(Handle {
-            handle: handle.upgrade().unwrap(),
-            cache: Arc::downgrade(&cache),
-        }),
+        // Retain the handle from the first instance.
+        handle: handle.upgrade().unwrap(),
     };
 
     // Drop the new cache instance. Wait the remainder of the first idle timeout
