Add shared stream interfaces

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,188 @@
+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, watcher, Config, Controller, WatchStreamExt},
+    Api, Client, ResourceExt,
+};
+use tokio::sync::mpsc;
+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> {
+    if pod.name_any() == "kube-system" {
+        return Ok(Action::await_change());
+    }
+
+    let labels = pod.labels();
+    if labels.is_empty() {
+        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 = "0.7.0"
+async-stream = "0.3.5"
 
 [dev-dependencies]
 kube = { path = "../kube", features = ["derive", "client", "runtime"], version = "<1.0.0, >=0.60.0" }

kube-runtime/src/controller/mod.rs

@@ -122,6 +122,23 @@ where
     })
 }
 
+pub fn trigger_self_shared<K, S>(
+    stream: S,
+    dyntype: K::DynamicType,
+) -> impl Stream<Item = Result<ReconcileRequest<K>, S::Error>>
+where
+    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,
@@ -701,6 +718,42 @@ where
         }
     }
 
+    // TODO: do an entrypoint for shared streams of owned objects
+    //
+    // Is it better to use a concrete type (i.e. a SubscribeHandle as a trigger)
+    // or to pass in the reader out-of-band?
+    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())
+    }
+
+    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 {