[WIP] Implement instance inhibit for node failure while streaming response

Author: kthui

lib/bindings/python/rust/lib.rs

@@ -643,14 +643,15 @@ impl Client {
         let annotated = annotated.unwrap_or(false);
 
         let (tx, rx) = tokio::sync::mpsc::channel(32);
-        let client = self.router.clone();
+        let router = self.router.clone();
 
         pyo3_async_runtimes::tokio::future_into_py(py, async move {
-            let stream = client.round_robin(request.into()).await.map_err(to_pyerr)?;
+            let (stream, instance_id) = router.round_robin_ret_id(request.into()).await.map_err(to_pyerr)?;
             tokio::spawn(process_stream(stream, tx));
             Ok(AsyncResponseStream::new(
                 Arc::new(Mutex::new(rx)),
                 annotated,
+                Some((router.clone(), instance_id)),
             ))
         })
     }
@@ -667,14 +668,15 @@ impl Client {
         let annotated = annotated.unwrap_or(false);
 
         let (tx, rx) = tokio::sync::mpsc::channel(32);
-        let client = self.router.clone();
+        let router = self.router.clone();
 
         pyo3_async_runtimes::tokio::future_into_py(py, async move {
-            let stream = client.random(request.into()).await.map_err(to_pyerr)?;
+            let (stream, instance_id) = router.random_ret_id(request.into()).await.map_err(to_pyerr)?;
             tokio::spawn(process_stream(stream, tx));
             Ok(AsyncResponseStream::new(
                 Arc::new(Mutex::new(rx)),
                 annotated,
+                Some((router.clone(), instance_id)),
             ))
         })
     }
@@ -692,10 +694,10 @@ impl Client {
         let annotated = annotated.unwrap_or(false);
 
         let (tx, rx) = tokio::sync::mpsc::channel(32);
-        let client = self.router.clone();
+        let router = self.router.clone();
 
         pyo3_async_runtimes::tokio::future_into_py(py, async move {
-            let stream = client
+            let stream = router
                 .direct(request.into(), instance_id)
                 .await
                 .map_err(to_pyerr)?;
@@ -705,6 +707,7 @@ impl Client {
             Ok(AsyncResponseStream::new(
                 Arc::new(Mutex::new(rx)),
                 annotated,
+                Some((router.clone(), instance_id)),
             ))
         })
     }
@@ -731,6 +734,7 @@ impl Client {
             Ok(AsyncResponseStream::new(
                 Arc::new(Mutex::new(rx)),
                 annotated,
+                None,
             ))
         })
     }
@@ -778,22 +782,27 @@ async fn process_stream(
 
 #[pyclass]
 struct AsyncResponseStream {
+    // Pipe for receiving responses.
     rx: Arc<Mutex<tokio::sync::mpsc::Receiver<RsAnnotated<PyObject>>>>,
     // Return response in Annotated wrapper.
     annotated: bool,
-    // For iterator to track complete final revceival.
+    // For iterator to track complete final receival.
     is_complete_final: Arc<AtomicBool>,
+    // Track the (router, instance_id) for reporting complete final errors.
+    routing: Option<(rs::pipeline::PushRouter<serde_json::Value, serde_json::Value>, i64)>,
 }
 
 impl AsyncResponseStream {
     pub fn new(
         rx: Arc<Mutex<tokio::sync::mpsc::Receiver<RsAnnotated<PyObject>>>>,
         annotated: bool,
+        routing: Option<(rs::pipeline::PushRouter<serde_json::Value, serde_json::Value>, i64)>,
     ) -> Self {
         AsyncResponseStream {
             rx,
             annotated,
             is_complete_final: Arc::new(AtomicBool::new(false)),
+            routing,
         }
     }
 }
@@ -811,6 +820,7 @@ impl AsyncResponseStream {
         let rx = self.rx.clone();
         let annotated = self.annotated;
         let is_complete_final = self.is_complete_final.clone();
+        let routing = self.routing.clone();
 
         pyo3_async_runtimes::tokio::future_into_py(py, async move {
             loop {
@@ -849,6 +859,9 @@ impl AsyncResponseStream {
                     }
                     None => {
                         if !(*is_complete_final).load(Ordering::Relaxed) {
+                            if let Some((router, instance_id)) = &routing {
+                                router.client.report_instance_down(*instance_id).await;
+                            }
                             return Err(PyConnectionError::new_err(
                                 "Stream closed unexpectedly before complete final",
                             ));

lib/runtime/src/pipeline/network/egress/push_router.rs

@@ -109,6 +109,16 @@ where
 
     /// Issue a request to the next available instance in a round-robin fashion
     pub async fn round_robin(&self, request: SingleIn<T>) -> anyhow::Result<ManyOut<U>> {
+        match self.round_robin_ret_id(request).await {
+            Ok((stream, _)) => Ok(stream),
+            Err(err) => Err(err),
+        }
+    }
+
+    pub async fn round_robin_ret_id(
+        &self,
+        request: SingleIn<T>,
+    ) -> anyhow::Result<(ManyOut<U>, i64)> {
         let slf = self;
         let routing_algorithm = move || async move {
             let counter = slf.round_robin_counter.fetch_add(1, Ordering::Relaxed);
@@ -135,6 +145,13 @@ where
 
     /// Issue a request to a random endpoint
     pub async fn random(&self, request: SingleIn<T>) -> anyhow::Result<ManyOut<U>> {
+        match self.random_ret_id(request).await {
+            Ok((stream, _)) => Ok(stream),
+            Err(err) => Err(err),
+        }
+    }
+
+    pub async fn random_ret_id(&self, request: SingleIn<T>) -> anyhow::Result<(ManyOut<U>, i64)> {
         let slf = self;
         let routing_algorithm = move || async move {
             let instance_id = {
@@ -178,8 +195,13 @@ where
             }
             Ok(instance_id)
         };
-        self.generate_with_fault_tolerance(routing_algorithm, request)
+        match self
+            .generate_with_fault_tolerance(routing_algorithm, request)
             .await
+        {
+            Ok((stream, _)) => Ok(stream),
+            Err(err) => Err(err),
+        }
     }
 
     pub async fn r#static(&self, request: SingleIn<T>) -> anyhow::Result<ManyOut<U>> {
@@ -194,7 +216,7 @@ where
         &self,
         routing_algorithm: F,
         request: SingleIn<T>,
-    ) -> anyhow::Result<ManyOut<U>>
+    ) -> anyhow::Result<(ManyOut<U>, i64)>
     where
         F: FnOnce() -> R,
         R: Future<Output = anyhow::Result<i64>>,
@@ -212,7 +234,10 @@ where
                 }
             }
         }
-        stream
+        match stream {
+            Ok(stream) => Ok((stream, instance_id)),
+            Err(err) => Err(err),
+        }
     }
 }