Refactor blocking loop

Signed-off-by: Danil-Grigorev <[email protected]>

Author: Danil-Grigorev

examples/broadcast_reflector.rs

@@ -1,18 +1,20 @@
-use futures::{future, pin_mut, stream, StreamExt};
+use futures::{future, lock::Mutex, pin_mut, stream, StreamExt};
 use k8s_openapi::api::{
     apps::v1::Deployment,
     core::v1::{ConfigMap, Secret},
 };
 use kube::{
     api::ApiResource,
     runtime::{
-        broadcaster, controller::Action, reflector::multi_dispatcher::MultiDispatcher, watcher, Controller,
+        broadcaster,
+        controller::Action,
+        reflector::multi_dispatcher::{BroadcastStream, MultiDispatcher},
+        watcher, Controller,
     },
     Api, Client, ResourceExt,
 };
 use std::{fmt::Debug, pin::pin, sync::Arc, time::Duration};
 use thiserror::Error;
-use tokio::sync::Mutex;
 use tracing::*;
 
 #[derive(Debug, Error)]
@@ -43,7 +45,7 @@ async fn main() -> anyhow::Result<()> {
 
     // multireflector stream
     let combo_stream = Arc::new(Mutex::new(stream::select_all(vec![])));
-    let watcher = broadcaster(writer.clone(), combo_stream.clone());
+    let watcher = broadcaster(writer.clone(), BroadcastStream::new(combo_stream.clone()));
 
     combo_stream.lock().await.push(
         watcher::watcher(

kube-runtime/src/reflector/mod.rs

@@ -18,19 +18,14 @@ use futures::{Stream, StreamExt};
 #[cfg(feature = "unstable-runtime-subscribe")]
 use kube_client::api::DynamicObject;
 #[cfg(feature = "unstable-runtime-subscribe")]
+use multi_dispatcher::BroadcastStream;
+#[cfg(feature = "unstable-runtime-subscribe")]
 use multi_dispatcher::MultiDispatcher;
 use std::hash::Hash;
 #[cfg(feature = "unstable-runtime-subscribe")] use std::pin::Pin;
 #[cfg(feature = "unstable-runtime-subscribe")] pub use store::store_shared;
 pub use store::{store, Store};
 
-#[cfg(feature = "unstable-runtime-subscribe")]
-/// Type for a shared stream of dynamic objects, which can be provided to [`broadcaster`]
-/// wrapped as [`Arc<Mutex<DynamicStream>>`], this type can be stored in context and provided
-/// to controllers, which can arbitrary modify existing event streams at runtime.
-pub type DynamicStream =
-    SelectAll<Pin<Box<dyn Stream<Item = Result<watcher::Event<DynamicObject>, watcher::Error>> + Send>>>;
-
 /// Cache objects from a [`watcher()`] stream into a local [`Store`]
 ///
 /// Observes the raw `Stream` of [`watcher::Event`] objects, and modifies the cache.
@@ -152,12 +147,15 @@ where
 
 // broadcaster uses a common stream of DynamicObject events to distribute to any subscribed typed watcher.
 #[cfg(feature = "unstable-runtime-subscribe")]
-pub fn broadcaster<W>(mut writer: MultiDispatcher, stream: Arc<Mutex<W>>) -> impl Stream<Item = W::Item>
+pub fn broadcaster<W>(
+    mut writer: MultiDispatcher,
+    mut stream: BroadcastStream<W>,
+) -> impl Stream<Item = W::Item>
 where
     W: Stream<Item = watcher::Result<watcher::Event<DynamicObject>>> + Unpin,
 {
     stream! {
-        while let Some(event) = stream.lock().await.next().await {
+        while let Some(event) = stream.next().await {
             match event {
                 Ok(ev) => {
                     writer.broadcast_event(&ev).await;

kube-runtime/src/reflector/multi_dispatcher.rs

@@ -1,5 +1,11 @@
-use std::hash::Hash;
+use std::{
+    hash::Hash,
+    pin::Pin,
+    sync::Arc,
+    task::{Context, Poll},
+};
 
+use futures::{lock::Mutex, FutureExt, Stream, StreamExt as _};
 use kube_client::{api::DynamicObject, Resource};
 use serde::de::DeserializeOwned;
 
@@ -43,6 +49,47 @@ impl MultiDispatcher {
 
     /// Broadcast an event to any downstream listeners subscribed on the store
     pub(crate) async fn broadcast_event(&mut self, event: &watcher::Event<DynamicObject>) {
-        self.dispatcher.broadcast(event.clone()).await
+        match event {
+            // Broadcast stores are pre-initialized
+            watcher::Event::InitDone => {}
+            ev => self.dispatcher.broadcast(ev.clone()).await,
+        }
+    }
+}
+
+/// See [`Scheduler::hold`]
+pub struct BroadcastStream<W> {
+    pub stream: Arc<Mutex<W>>,
+}
+
+impl<W> Clone for BroadcastStream<W> {
+    fn clone(&self) -> Self {
+        Self {
+            stream: self.stream.clone(),
+        }
+    }
+}
+
+impl<W> BroadcastStream<W>
+where
+    W: Stream<Item = watcher::Result<watcher::Event<DynamicObject>>> + Unpin,
+{
+    pub fn new(stream: Arc<Mutex<W>>) -> Self {
+        Self { stream }
+    }
+}
+
+impl<W> Stream for BroadcastStream<W>
+where
+    W: Stream<Item = watcher::Result<watcher::Event<DynamicObject>>> + Unpin,
+{
+    type Item = W::Item;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        if let Some(mut stream) = self.stream.try_lock() {
+            return stream.poll_next_unpin(cx);
+        }
+
+        Poll::Pending
     }
 }