Add shared stream interfaces (#1449)

Currently, there is no way for a controller to use a shared stream. If two or
more distinct controllers require a watch on the same type of resource,
separate watches have to be created. To better utilise resources, it would be
ideal for only one watch to exist and share multiple streams from the same
watch.

This change introduces shared streams through a couple of new primitives and
changes to public interfaces. Stores can now be created through a new
constructor that takes a buffer size. When a buffer size is specified, stores
will have an underlying channel that supports creating multiple readers from a
single watch passed into a reflector.

Furthermore, we introduce a new interface `for_shared_stream` that allows
threading shared streams through controllers. This feature is currently gated
behind a feature flag in the runtime crate.

N.B. a current `StreamSubscribe` interface exists for the same reasons
described above. The interface is not compatible with controllers and does not
exert backpressure and will be removed in a future iteration.

---------

Signed-off-by: Matei David <[email protected]>
Co-authored-by: Eirik A <[email protected]>

Cargo.toml

@@ -34,6 +34,8 @@ missing_docs = "deny"
 ahash = "0.8"
 anyhow = "1.0.71"
 assert-json-diff = "2.0.2"
+async-broadcast = "0.7.0"
+async-stream = "0.3.5"
 async-trait = "0.1.64"
 backoff = "0.4.0"
 base64 = "0.22.0"

examples/Cargo.toml

@@ -61,6 +61,10 @@ crossterm = "0.27.0"
 name = "configmapgen_controller"
 path = "configmapgen_controller.rs"
 
+[[example]]
+name = "shared_stream_controllers"
+path = "shared_stream_controllers.rs"
+
 [[example]]
 name = "crd_api"
 path = "crd_api.rs"

examples/shared_stream_controllers.rs

@@ -0,0 +1,187 @@
+use std::{sync::Arc, time::Duration};
+
+use futures::StreamExt;
+use k8s_openapi::api::core::v1::{Pod, PodCondition};
+use kube::{
+    api::{Patch, PatchParams},
+    runtime::{
+        controller::Action,
+        reflector::{self},
+        watcher, Config, Controller, WatchStreamExt,
+    },
+    Api, Client, ResourceExt,
+};
+use tracing::{debug, error, info, warn};
+
+use thiserror::Error;
+
+// Helper module that namespaces two constants describing a Kubernetes status condition
+pub mod condition {
+    pub static UNDOCUMENTED_TYPE: &str = "UndocumentedPort";
+    pub static STATUS_TRUE: &str = "True";
+}
+
+const SUBSCRIBE_BUFFER_SIZE: usize = 256;
+
+#[derive(Debug, Error)]
+enum Error {
+    #[error("Failed to patch pod: {0}")]
+    WriteFailed(#[source] kube::Error),
+
+    #[error("Missing po field: {0}")]
+    MissingField(&'static str),
+}
+
+#[derive(Clone)]
+struct Data {
+    client: Client,
+}
+
+/// A simple reconciliation function that will copy a pod's labels into the annotations.
+async fn reconcile_metadata(pod: Arc<Pod>, ctx: Arc<Data>) -> Result<Action, Error> {
+    let namespace = &pod.namespace().unwrap_or_default();
+    if namespace == "kube-system" {
+        return Ok(Action::await_change());
+    }
+
+    let mut pod = (*pod).clone();
+    pod.metadata.managed_fields = None;
+    // combine labels and annotations into a new map
+    let labels = pod.labels().clone().into_iter();
+    pod.annotations_mut().extend(labels);
+
+    let pod_api = Api::<Pod>::namespaced(
+        ctx.client.clone(),
+        pod.metadata
+            .namespace
+            .as_ref()
+            .ok_or_else(|| Error::MissingField(".metadata.name"))?,
+    );
+
+    pod_api
+        .patch(
+            &pod.name_any(),
+            &PatchParams::apply("controller-1"),
+            &Patch::Apply(&pod),
+        )
+        .await
+        .map_err(Error::WriteFailed)?;
+
+    Ok(Action::requeue(Duration::from_secs(300)))
+}
+
+/// Another reconiliation function that will add an 'UndocumentedPort' condition to pods that do
+/// do not have any ports declared across all containers.
+async fn reconcile_status(pod: Arc<Pod>, ctx: Arc<Data>) -> Result<Action, Error> {
+    for container in pod.spec.clone().unwrap_or_default().containers.iter() {
+        if container.ports.clone().unwrap_or_default().len() != 0 {
+            debug!(name = %pod.name_any(), "Skipped updating pod with documented ports");
+            return Ok(Action::await_change());
+        }
+    }
+
+    let pod_api = Api::<Pod>::namespaced(
+        ctx.client.clone(),
+        pod.metadata
+            .namespace
+            .as_ref()
+            .ok_or_else(|| Error::MissingField(".metadata.name"))?,
+    );
+
+    let undocumented_condition = PodCondition {
+        type_: condition::UNDOCUMENTED_TYPE.into(),
+        status: condition::STATUS_TRUE.into(),
+        ..Default::default()
+    };
+    let value = serde_json::json!({
+        "status": {
+            "name": pod.name_any(),
+            "kind": "Pod",
+            "conditions": vec![undocumented_condition]
+        }
+    });
+    pod_api
+        .patch_status(
+            &pod.name_any(),
+            &PatchParams::apply("controller-2"),
+            &Patch::Strategic(value),
+        )
+        .await
+        .map_err(Error::WriteFailed)?;
+
+    Ok(Action::requeue(Duration::from_secs(300)))
+}
+
+fn error_policy(obj: Arc<Pod>, error: &Error, _ctx: Arc<Data>) -> Action {
+    error!(%error, name = %obj.name_any(), "Failed reconciliation");
+    Action::requeue(Duration::from_secs(10))
+}
+
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    tracing_subscriber::fmt::init();
+
+    let client = Client::try_default().await?;
+    let pods = Api::<Pod>::namespaced(client.clone(), "default");
+    let config = Config::default().concurrency(2);
+    let ctx = Arc::new(Data { client });
+
+    // Create a shared store with a predefined buffer that will be shared between subscribers.
+    let (reader, writer) = reflector::store_shared(SUBSCRIBE_BUFFER_SIZE);
+    // Before threading an object watch through the store, create a subscriber.
+    // Any number of subscribers can be created from one writer.
+    let subscriber = writer
+        .subscribe()
+        .expect("subscribers can only be created from shared stores");
+
+    // Reflect a stream of pod watch events into the store and apply a backoff. For subscribers to
+    // be able to consume updates, the reflector must be shared.
+    let pod_watch = watcher(pods.clone(), Default::default())
+        .default_backoff()
+        .reflect_shared(writer)
+        .for_each(|res| async move {
+            match res {
+                Ok(event) => debug!("Received event on root stream {event:?}"),
+                Err(error) => error!(%error, "Unexpected error when watching resource"),
+            }
+        });
+
+    // Create the first controller using the reconcile_metadata function. Controllers accept
+    // subscribers through a dedicated interface.
+    let metadata_controller = Controller::for_shared_stream(subscriber.clone(), reader)
+        .with_config(config.clone())
+        .shutdown_on_signal()
+        .run(reconcile_metadata, error_policy, ctx.clone())
+        .for_each(|res| async move {
+            match res {
+                Ok(v) => info!("Reconciled metadata {v:?}"),
+                Err(error) => warn!(%error, "Failed to reconcile metadata"),
+            }
+        });
+
+    // Subscribers can be used to get a read handle on the store, if the initial handle has been
+    // moved or dropped.
+    let reader = subscriber.reader();
+    // Create the second controller using the reconcile_status function.
+    let status_controller = Controller::for_shared_stream(subscriber, reader)
+        .with_config(config)
+        .shutdown_on_signal()
+        .run(reconcile_status, error_policy, ctx)
+        .for_each(|res| async move {
+            match res {
+                Ok(v) => info!("Reconciled status {v:?}"),
+                Err(error) => warn!(%error, "Failed to reconcile status"),
+            }
+        });
+
+    // Drive streams to readiness. The initial watch (that is reflected) needs to be driven to
+    // consume events from the API Server and forward them to subscribers.
+    //
+    // Both controllers will operate on shared objects.
+    tokio::select! {
+        _ = futures::future::join(metadata_controller, status_controller) => {},
+        _ = pod_watch => {}
+    }
+
+    Ok(())
+}

kube-runtime/Cargo.toml

@@ -48,6 +48,8 @@ backoff.workspace = true
 async-trait.workspace = true
 hashbrown.workspace = true
 k8s-openapi.workspace = true
+async-broadcast.workspace = true
+async-stream.workspace = true
 
 [dev-dependencies]
 kube = { path = "../kube", features = ["derive", "client", "runtime"], version = "<1.0.0, >=0.60.0" }

kube-runtime/src/controller/mod.rs

@@ -123,6 +123,28 @@ where
     })
 }
 
+/// Enqueues the object itself for reconciliation when the object is behind a
+/// shared pointer
+#[cfg(feature = "unstable-runtime-subscribe")]
+fn trigger_self_shared<K, S>(
+    stream: S,
+    dyntype: K::DynamicType,
+) -> impl Stream<Item = Result<ReconcileRequest<K>, S::Error>>
+where
+    // Input stream has item as some Arc'd Resource (via
+    // Controller::for_shared_stream)
+    S: TryStream<Ok = Arc<K>>,
+    K: Resource,
+    K::DynamicType: Clone,
+{
+    trigger_with(stream, move |obj| {
+        Some(ReconcileRequest {
+            obj_ref: ObjectRef::from_obj_with(obj.as_ref(), dyntype.clone()),
+            reason: ReconcileReason::ObjectUpdated,
+        })
+    })
+}
+
 /// Enqueues any mapper returned `K` types for reconciliation
 fn trigger_others<S, K, I>(
     stream: S,
@@ -703,6 +725,117 @@ where
         }
     }
 
+    /// This is the same as [`Controller::for_stream`]. Instead of taking an
+    /// `Api` (e.g. [`Controller::new`]), a stream of resources is used. Shared
+    /// streams can be created out-of-band by subscribing on a store `Writer`.
+    /// Through this interface, multiple controllers can use the same root
+    /// (shared) input stream of resources to keep memory overheads smaller.
+    ///
+    /// **N.B**: This constructor requires an
+    /// [`unstable`](https://github.com/kube-rs/kube/blob/main/kube-runtime/Cargo.toml#L17-L21)
+    /// feature.
+    ///
+    /// Prefer [`Controller::new`] or [`Controller::for_stream`] if you do not
+    /// need to share the stream.
+    ///
+    /// ## Warning:
+    ///
+    /// You **must** ensure the root stream (i.e. stream created through a `reflector()`)
+    /// is driven to readiness independently of this controller to ensure the
+    /// watcher never deadlocks.
+    ///
+    /// # Example:
+    ///
+    /// ```no_run
+    /// # use futures::StreamExt;
+    /// # use k8s_openapi::api::apps::v1::Deployment;
+    /// # use kube::runtime::controller::{Action, Controller};
+    /// # use kube::runtime::{predicates, watcher, reflector, WatchStreamExt};
+    /// # use kube::{Api, Client, Error, ResourceExt};
+    /// # use std::sync::Arc;
+    /// # async fn reconcile(_: Arc<Deployment>, _: Arc<()>) -> Result<Action, Error> { Ok(Action::await_change()) }
+    /// # fn error_policy(_: Arc<Deployment>, _: &kube::Error, _: Arc<()>) -> Action { Action::await_change() }
+    /// # async fn doc(client: kube::Client) {
+    /// let api: Api<Deployment> = Api::default_namespaced(client);
+    /// let (reader, writer) = reflector::store_shared(128);
+    /// let subscriber = writer
+    ///     .subscribe()
+    ///     .expect("subscribers can only be created from shared stores");
+    /// let deploys = watcher(api, watcher::Config::default())
+    ///     .default_backoff()
+    ///     .reflect(writer)
+    ///     .applied_objects()
+    ///     .for_each(|ev| async move {
+    ///         match ev {
+    ///             Ok(obj) => tracing::info!("got obj {obj:?}"),
+    ///             Err(error) => tracing::error!(%error, "received error")
+    ///         }
+    ///     });
+    ///
+    /// let controller = Controller::for_shared_stream(subscriber, reader)
+    ///     .run(reconcile, error_policy, Arc::new(()))
+    ///     .for_each(|ev| async move {
+    ///         tracing::info!("reconciled {ev:?}")
+    ///     });
+    ///
+    /// // Drive streams using a select statement
+    /// tokio::select! {
+    ///   _ = deploys => {},
+    ///   _ = controller => {},
+    /// }
+    /// # }
+    #[cfg(feature = "unstable-runtime-subscribe")]
+    pub fn for_shared_stream(trigger: impl Stream<Item = Arc<K>> + Send + 'static, reader: Store<K>) -> Self
+    where
+        K::DynamicType: Default,
+    {
+        Self::for_shared_stream_with(trigger, reader, Default::default())
+    }
+
+    /// This is the same as [`Controller::for_stream`]. Instead of taking an
+    /// `Api` (e.g. [`Controller::new`]), a stream of resources is used. Shared
+    /// streams can be created out-of-band by subscribing on a store `Writer`.
+    /// Through this interface, multiple controllers can use the same root
+    /// (shared) input stream of resources to keep memory overheads smaller.
+    ///
+    /// **N.B**: This constructor requires an
+    /// [`unstable`](https://github.com/kube-rs/kube/blob/main/kube-runtime/Cargo.toml#L17-L21)
+    /// feature.
+    ///
+    /// Prefer [`Controller::new`] or [`Controller::for_stream`] if you do not
+    /// need to share the stream.
+    ///
+    /// This variant constructor is used for [`dynamic`] types found through
+    /// discovery. Prefer [`Controller::for_shared_stream`] for static types (i.e.
+    /// known at compile time).
+    ///
+    /// [`dynamic`]: kube_client::core::dynamic
+    #[cfg(feature = "unstable-runtime-subscribe")]
+    pub fn for_shared_stream_with(
+        trigger: impl Stream<Item = Arc<K>> + Send + 'static,
+        reader: Store<K>,
+        dyntype: K::DynamicType,
+    ) -> Self {
+        let mut trigger_selector = stream::SelectAll::new();
+        let self_watcher = trigger_self_shared(trigger.map(Ok), dyntype.clone()).boxed();
+        trigger_selector.push(self_watcher);
+        Self {
+            trigger_selector,
+            trigger_backoff: Box::<DefaultBackoff>::default(),
+            graceful_shutdown_selector: vec![
+                // Fallback future, ensuring that we never terminate if no additional futures are added to the selector
+                future::pending().boxed(),
+            ],
+            forceful_shutdown_selector: vec![
+                // Fallback future, ensuring that we never terminate if no additional futures are added to the selector
+                future::pending().boxed(),
+            ],
+            dyntype,
+            reader,
+            config: Default::default(),
+        }
+    }
+
     /// Specify the configuration for the controller's behavior.
     #[must_use]
     pub fn with_config(mut self, config: Config) -> Self {