Repository: tikv/tikv
File: components/raftstore/src/store/worker/pd.rs (Rust)

Code before:
                            for mut demote in plan.take_demotes().into_iter() {
                                if let Err(e) = handle.send_demote_peers(
                                    demote.get_region_id(),
                                    demote.take_failed_voters().into_vec(),
                                    syncer.clone(),
                                ) {
                                    error!("fail to send update peer list message for recovery"; "err" => ?e);
                                }
                            }
                        }
                    }
                    // Forcely awaken all hibernated regions if there existed slow stores in this
                    // cluster.
                    if let Some(awaken_regions) = resp.awaken_regions.take() {
                        info!("forcely awaken hibernated regions in this store");
                        let _ = router.send_store_msg(StoreMsg::AwakenRegions {
                            abnormal_stores: awaken_regions.get_abnormal_stores().to_vec(),
                        });
                    }
                    // Control grpc server.
                    if let Some(op) = resp.control_grpc.take() {
                        if let Err(e) =
                            scheduler.schedule(Task::ControlGrpcServer(op.get_ctrl_event()))
                        {
                            warn!("fail to schedule control grpc task"; "err" => ?e);
                        }
                    }
                    // NodeState for this store.
                    {
                        let state = (|| {
                            #[cfg(feature = "failpoints")]
                            fail_point!("manually_set_store_offline", |_| {
                                metapb::NodeState::Removing
                            });
                            resp.get_state()
                        })();
                        match state {
                            metapb::NodeState::Removing | metapb::NodeState::Removed => {
                                let is_offlined = snap_mgr.is_offlined();
                                if !is_offlined {
                                    snap_mgr.set_offline(true);
                                    info!("store is offlined by pd");
                                }
                            }
                            // As for NodeState::Preparing | Serving, if `is_offlined == true`,
                            // it means the store has been re-added into the cluster and
                            // the offline operation is terminated. Therefore, the state
                            // `is_offlined` should be reset with `false`.
                            _ => {
                                let is_offlined = snap_mgr.is_offlined();
                                if is_offlined {

Review comment (suggestion):
Just a note: in some cases there are more than 50K hibernate regions (e.g. #17101), a chunk size of `1024` would mean we'll be sending 50 `StoreMsg::AwakenRegions` messages. Probably fine, but worth being aware of.

Code after:
                                if let Err(e) = handle.send_destroy_peer(delete, syncer.clone()) {
                                    error!("fail to send destroy peer message for recovery"; "err" => ?e);
                                }
                            }
                            for mut demote in plan.take_demotes().into_iter() {
                                if let Err(e) = handle.send_demote_peers(
                                    demote.get_region_id(),
                                    demote.take_failed_voters().into_vec(),
                                    syncer.clone(),
                                ) {
                                    error!("fail to send update peer list message for recovery"; "err" => ?e);
                                }
                            }
                        }
                    }
                    // Force awaken all hibernated regions if there are slow stores in this
                    // cluster. To mitigate the impact of stalls when awakening too many regions,
                    // we break up all regions into small batches.
                    if let Some(awaken_regions) = resp.awaken_regions.take() {
                        info!("forcely awaken hibernated regions in this store");
                        let abnormal_stores = awaken_regions.get_abnormal_stores().to_vec();
                        // The chunk_size of 1024 is an empirical batch size that balances
                        // between generating too many `StoreMsg::AwakenRegions` messages and
                        // keeping each batch small enough to avoid stalling Raftstore for
                        // extended periods.
                        let chunk_size = 1024;
                        let mut region_ids = Vec::with_capacity(chunk_size);

                        for (id, _) in region_peers.read().unwrap().iter() {
                            region_ids.push(*id);

                            // Send batch when it reaches the chunk size
                            if region_ids.len() >= chunk_size {
                                let _ = router.send_store_msg(StoreMsg::AwakenRegions {
                                    abnormal_stores: abnormal_stores.clone(),
                                    region_ids: std::mem::take(&mut region_ids),
                                });
                                region_ids.reserve(chunk_size);
                            }
                        }

                        // Send remaining regions if any
                        if !region_ids.is_empty() {
                            let _ = router.send_store_msg(StoreMsg::AwakenRegions {
                                abnormal_stores: abnormal_stores.clone(),
                                region_ids,
                            });
                        }
                    }
                    // Control grpc server.
                    if let Some(op) = resp.control_grpc.take() {