feat: compute node unregisters from meta for graceful shutdown

proto/stream_service.proto

@@ -108,10 +108,12 @@ message StreamingControlStreamRequest {
 
 message StreamingControlStreamResponse {
   message InitResponse {}
+  message ShutdownResponse {}
 
   oneof response {
     InitResponse init = 1;
     BarrierCompleteResponse complete_barrier = 2;
+    ShutdownResponse shutdown = 3;
   }
 }
 

src/common/common_service/src/observer_manager.rs

@@ -157,7 +157,7 @@ where
                 match self.rx.message().await {
                     Ok(resp) => {
                         if resp.is_none() {
-                            tracing::error!("Stream of notification terminated.");
+                            tracing::warn!("Stream of notification terminated.");
                             self.re_subscribe().await;
                             continue;
                         }

src/compute/src/server.rs

@@ -405,7 +405,7 @@ pub async fn compute_node_serve(
         meta_cache,
         block_cache,
     );
-    let config_srv = ConfigServiceImpl::new(batch_mgr, stream_mgr);
+    let config_srv = ConfigServiceImpl::new(batch_mgr, stream_mgr.clone());
     let health_srv = HealthServiceImpl::new();
 
     let telemetry_manager = TelemetryManager::new(
@@ -469,8 +469,12 @@ pub async fn compute_node_serve(
     // Wait for the shutdown signal.
     shutdown.cancelled().await;
 
-    // TODO(shutdown): gracefully unregister from the meta service (need to cautious since it may
-    // trigger auto-scaling)
+    // Unregister from the meta service, then...
+    // - batch queries will not be scheduled to this compute node,
+    // - streaming actors will not be scheduled to this compute node after next recovery.
+    meta_client.try_unregister().await;
+    // Shutdown the streaming manager.
+    let _ = stream_mgr.shutdown().await;
 
     // NOTE(shutdown): We can't simply join the tonic server here because it only returns when all
     // existing connections are closed, while we have long-running streaming calls that never

src/meta/src/barrier/info.rs

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::collections::{HashMap, HashSet};
+use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
 
 use risingwave_common::catalog::TableId;
 use risingwave_pb::common::PbWorkerNode;
@@ -137,11 +137,26 @@ impl InflightActorInfo {
             .into_iter()
             .map(|node| (node.id, node))
             .collect::<HashMap<_, _>>();
-        for (actor_id, location) in &self.actor_location_map {
-            if !new_node_map.contains_key(location) {
-                warn!(actor_id, location, node = ?self.node_map.get(location), "node with running actors is deleted");
+
+        let mut deleted_actors = BTreeMap::new();
+        for (&actor_id, &location) in &self.actor_location_map {
+            if !new_node_map.contains_key(&location) {
+                deleted_actors
+                    .entry(location)
+                    .or_insert_with(BTreeSet::new)
+                    .insert(actor_id);
             }
         }
+        for (node_id, actors) in deleted_actors {
+            let node = self.node_map.get(&node_id);
+            warn!(
+                node_id,
+                ?node,
+                ?actors,
+                "node with running actors is deleted"
+            );
+        }
+
         self.node_map = new_node_map;
     }
 

src/meta/src/barrier/mod.rs

@@ -985,24 +985,18 @@ impl GlobalBarrierManagerContext {
     }
 
     fn report_complete_event(&self, duration_sec: f64, command_ctx: &CommandContext) {
-        {
-            {
-                {
-                    // Record barrier latency in event log.
-                    use risingwave_pb::meta::event_log;
-                    let event = event_log::EventBarrierComplete {
-                        prev_epoch: command_ctx.prev_epoch.value().0,
-                        cur_epoch: command_ctx.curr_epoch.value().0,
-                        duration_sec,
-                        command: command_ctx.command.to_string(),
-                        barrier_kind: command_ctx.kind.as_str_name().to_string(),
-                    };
-                    self.env
-                        .event_log_manager_ref()
-                        .add_event_logs(vec![event_log::Event::BarrierComplete(event)]);
-                }
-            }
-        }
+        // Record barrier latency in event log.
+        use risingwave_pb::meta::event_log;
+        let event = event_log::EventBarrierComplete {
+            prev_epoch: command_ctx.prev_epoch.value().0,
+            cur_epoch: command_ctx.curr_epoch.value().0,
+            duration_sec,
+            command: command_ctx.command.to_string(),
+            barrier_kind: command_ctx.kind.as_str_name().to_string(),
+        };
+        self.env
+            .event_log_manager_ref()
+            .add_event_logs(vec![event_log::Event::BarrierComplete(event)]);
     }
 }
 

src/meta/src/barrier/recovery.rs

@@ -234,9 +234,12 @@ impl GlobalBarrierManager {
                 .committed_epoch
                 .into(),
         );
+
         // Mark blocked and abort buffered schedules, they might be dirty already.
         self.scheduled_barriers
             .abort_and_mark_blocked("cluster is under recovering");
+        // Clear all control streams to release resources (connections to compute nodes) first.
+        self.control_stream_manager.clear();
 
         tracing::info!("recovery start!");
         let retry_strategy = Self::get_retry_strategy();
@@ -288,6 +291,7 @@ impl GlobalBarrierManager {
                     // Resolve actor info for recovery. If there's no actor to recover, most of the
                     // following steps will be no-op, while the compute nodes will still be reset.
                     // FIXME: Transactions should be used.
+                    // TODO(error-handling): attach context to the errors and log them together, instead of inspecting everywhere.
                     let mut info = if !self.env.opts.disable_automatic_parallelism_control
                         && background_streaming_jobs.is_empty()
                     {

src/meta/src/barrier/rpc.rs

@@ -167,25 +167,39 @@ impl ControlStreamManager {
         Ok(())
     }
 
+    /// Clear all nodes and response streams in the manager.
+    pub(super) fn clear(&mut self) {
+        *self = Self::new(self.context.clone());
+    }
+
     async fn next_response(
         &mut self,
     ) -> Option<(WorkerId, MetaResult<StreamingControlStreamResponse>)> {
         let (worker_id, response_stream, result) = self.response_streams.next().await?;
-        if result.is_ok() {
-            self.response_streams
-                .push(into_future(worker_id, response_stream));
+
+        match result.as_ref().map(|r| r.response.as_ref().unwrap()) {
+            Ok(streaming_control_stream_response::Response::Shutdown(_)) | Err(_) => {
+                // Do not add it back to the `response_streams` so that it will not be polled again.
+            }
+            _ => {
+                self.response_streams
+                    .push(into_future(worker_id, response_stream));
+            }
         }
+
         Some((worker_id, result))
     }
 
     pub(super) async fn next_complete_barrier_response(
         &mut self,
     ) -> MetaResult<(WorkerId, u64, BarrierCompleteResponse)> {
+        use streaming_control_stream_response::Response;
+
         loop {
             let (worker_id, result) = pending_on_none(self.next_response()).await;
             match result {
-                Ok(resp) => match resp.response {
-                    Some(streaming_control_stream_response::Response::CompleteBarrier(resp)) => {
+                Ok(resp) => match resp.response.unwrap() {
+                    Response::CompleteBarrier(resp) => {
                         let node = self
                             .nodes
                             .get_mut(&worker_id)
@@ -196,26 +210,37 @@ impl ControlStreamManager {
                             .expect("should exist when get collect resp");
                         break Ok((worker_id, command.prev_epoch.value().0, resp));
                     }
-                    resp => {
-                        break Err(anyhow!("get unexpected resp: {:?}", resp).into());
+                    Response::Shutdown(_) => {
+                        let _ = self
+                            .nodes
+                            .remove(&worker_id)
+                            .expect("should exist when get shutdown resp");
+                        // TODO: if there's no actor running on the node, we can ignore and not trigger recovery.
+                        break Err(anyhow!("worker node {worker_id} is shutting down").into());
+                    }
+                    Response::Init(_) => {
+                        // This arm should be unreachable.
+                        break Err(anyhow!("get unexpected init response").into());
                     }
                 },
                 Err(err) => {
-                    let mut node = self
+                    let node = self
                         .nodes
                         .remove(&worker_id)
                         .expect("should exist when get collect resp");
                     // Note: No need to use `?` as the backtrace is from meta and not useful.
                     warn!(node = ?node.worker, err = %err.as_report(), "get error from response stream");
-                    if let Some(command) = node.inflight_barriers.pop_front() {
+
+                    if let Some(command) = node.inflight_barriers.into_iter().next() {
+                        // FIXME: this future can be cancelled during collection, so the error collection
+                        // might not work as expected.
                         let errors = self.collect_errors(node.worker.id, err).await;
                         let err = merge_node_rpc_errors("get error from control stream", errors);
                         self.context.report_collect_failure(&command, &err);
                         break Err(err);
                     } else {
                         // for node with no inflight barrier, simply ignore the error
-                        info!(node = ?node.worker, "no inflight barrier no node. Ignore error");
-                        continue;
+                        info!(node = ?node.worker, error = %err.as_report(), "no inflight barrier in the node, ignore error");
                     }
                 }
             }
@@ -239,6 +264,7 @@ impl ControlStreamManager {
             })
             .await;
         }
+        tracing::debug!(?errors, "collected stream errors");
         errors
     }
 }

src/meta/src/stream/scale.rs

@@ -2618,15 +2618,20 @@ impl GlobalStreamManager {
                 notification = local_notification_rx.recv() => {
                     let notification = notification.expect("local notification channel closed in loop of stream manager");
 
+                    // Only maintain the cache for streaming compute nodes.
+                    let worker_is_streaming_compute = |worker: &WorkerNode| {
+                        worker.get_type() == Ok(WorkerType::ComputeNode)
+                            && worker.property.as_ref().unwrap().is_streaming
+                    };
+
                     match notification {
                         LocalNotification::WorkerNodeActivated(worker) => {
-                            match (worker.get_type(), worker.property.as_ref()) {
-                                (Ok(WorkerType::ComputeNode), Some(prop)) if prop.is_streaming => {
-                                    tracing::info!("worker {} activated notification received", worker.id);
-                                }
-                                _ => continue
+                            if !worker_is_streaming_compute(&worker) {
+                                continue;
                             }
 
+                            tracing::info!(worker = worker.id, "worker activated notification received");
+
                             let prev_worker = worker_cache.insert(worker.id, worker.clone());
 
                             match prev_worker {
@@ -2645,11 +2650,14 @@ impl GlobalStreamManager {
                         // Since our logic for handling passive scale-in is within the barrier manager,
                         // there’s not much we can do here. All we can do is proactively remove the entries from our cache.
                         LocalNotification::WorkerNodeDeleted(worker) => {
+                            if !worker_is_streaming_compute(&worker) {
+                                continue;
+                            }
+
                             match worker_cache.remove(&worker.id) {
                                 Some(prev_worker) => {
                                     tracing::info!(worker = prev_worker.id, "worker removed from stream manager cache");
                                 }
-
                                 None => {
                                     tracing::warn!(worker = worker.id, "worker not found in stream manager cache, but it was removed");
                                 }

src/rpc_client/src/meta_client.rs

@@ -14,6 +14,8 @@
 
 use std::collections::HashMap;
 use std::fmt::{Debug, Display};
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering::Relaxed;
 use std::sync::Arc;
 use std::thread;
 use std::time::{Duration, SystemTime};
@@ -115,6 +117,7 @@ pub struct MetaClient {
     inner: GrpcMetaClient,
     meta_config: MetaConfig,
     cluster_id: String,
+    shutting_down: Arc<AtomicBool>,
 }
 
 impl MetaClient {
@@ -276,6 +279,7 @@ impl MetaClient {
             inner: grpc_meta_client,
             meta_config: meta_config.to_owned(),
             cluster_id: add_worker_resp.cluster_id,
+            shutting_down: Arc::new(false.into()),
         };
 
         static REPORT_PANIC: std::sync::Once = std::sync::Once::new();
@@ -322,8 +326,12 @@ impl MetaClient {
         let resp = self.inner.heartbeat(request).await?;
         if let Some(status) = resp.status {
             if status.code() == risingwave_pb::common::status::Code::UnknownWorker {
-                tracing::error!("worker expired: {}", status.message);
-                std::process::exit(1);
+                // Ignore the error if we're already shutting down.
+                // Otherwise, exit the process.
+                if !self.shutting_down.load(Relaxed) {
+                    tracing::error!(message = status.message, "worker expired");
+                    std::process::exit(1);
+                }
             }
         }
         Ok(())
@@ -745,6 +753,7 @@ impl MetaClient {
             host: Some(self.host_addr.to_protobuf()),
         };
         self.inner.delete_worker_node(request).await?;
+        self.shutting_down.store(true, Relaxed);
         Ok(())
     }
 

src/storage/src/hummock/event_handler/mod.rs

@@ -217,7 +217,7 @@ impl Drop for LocalInstanceGuard {
                     instance_id: self.instance_id,
                 })
                 .unwrap_or_else(|err| {
-                    tracing::error!(
+                    tracing::debug!(
                         error = %err.as_report(),
                         table_id = %self.table_id,
                         instance_id = self.instance_id,

src/storage/src/hummock/store/hummock_storage.rs

@@ -37,7 +37,6 @@ use risingwave_rpc_client::HummockMetaClient;
 use thiserror_ext::AsReport;
 use tokio::sync::mpsc::{unbounded_channel, UnboundedSender};
 use tokio::sync::oneshot;
-use tracing::error;
 
 use super::local_hummock_storage::LocalHummockStorage;
 use super::version::{read_filter_for_version, CommittedVersion, HummockVersionReader};
@@ -75,7 +74,7 @@ impl Drop for HummockStorageShutdownGuard {
         let _ = self
             .shutdown_sender
             .send(HummockEvent::Shutdown)
-            .inspect_err(|e| error!(event = ?e.0, "unable to send shutdown"));
+            .inspect_err(|e| tracing::debug!(event = ?e.0, "unable to send shutdown"));
     }
 }
 

src/stream/src/error.rs

@@ -86,12 +86,14 @@ pub enum ErrorKind {
         actor_id: ActorId,
         reason: &'static str,
     },
+
     #[error("Secret error: {0}")]
     Secret(
         #[from]
         #[backtrace]
         SecretError,
     ),
+
     #[error(transparent)]
     Uncategorized(
         #[from]