raftstore & raftstore-v2:control grpc server according to slowness. (#…

…15088)

close #15086

Make raftstore & raftstore-v2 able to control the grpc service, according to the slowness.

Signed-off-by: lucasliang <nkcs_lykx@hotmail.com>

Co-authored-by: tonyxuqqi <tonyxuqi@outlook.com>

Cargo.toml

@@ -128,6 +128,7 @@ serde = { version = "1.0", features = ["derive"] }
 serde_derive = "1.0"
 serde_ignored = "0.1"
 serde_json = { version = "1.0", features = ["preserve_order"] }
+service = { workspace = true }
 slog = { workspace = true }
 slog-global = { workspace = true }
 smallvec = "1.4"
@@ -254,6 +255,7 @@ members = [
   "components/resource_metering",
   "components/security",
   "components/server",
+  "components/service",
   "components/snap_recovery",
   "components/sst_importer",
   "components/test_backup",
@@ -331,6 +333,7 @@ resource_control = { path = "components/resource_control" }
 resource_metering = { path = "components/resource_metering" }
 security = { path = "components/security" }
 server = { path = "components/server" }
+service = { path = "components/service" }
 snap_recovery = { path = "components/snap_recovery" }
 sst_importer = { path = "components/sst_importer", default-features = false }
 test_backup = { path = "components/test_backup" }

cmd/tikv-server/src/main.rs

@@ -217,8 +217,11 @@ fn main() {
         process::exit(1)
     }
 
+    let (service_event_tx, service_event_rx) = tikv_util::mpsc::unbounded(); // pipe for controling service
     match config.storage.engine {
-        EngineType::RaftKv => server::server::run_tikv(config),
-        EngineType::RaftKv2 => server::server2::run_tikv(config),
+        EngineType::RaftKv => server::server::run_tikv(config, service_event_tx, service_event_rx),
+        EngineType::RaftKv2 => {
+            server::server2::run_tikv(config, service_event_tx, service_event_rx)
+        }
     }
 }

components/raftstore-v2/Cargo.toml

@@ -55,6 +55,7 @@ raftstore = { workspace = true }
 rand = "0.8.3"
 resource_control = { workspace = true }
 resource_metering = { workspace = true }
+service = { workspace = true }
 slog = "2.3"
 smallvec = "1.4"
 sst_importer = { workspace = true }

components/raftstore-v2/src/batch/store.rs

@@ -40,6 +40,7 @@ use raftstore::{
     },
 };
 use resource_metering::CollectorRegHandle;
+use service::service_manager::GrpcServiceManager;
 use slog::{warn, Logger};
 use sst_importer::SstImporter;
 use tikv_util::{
@@ -678,6 +679,7 @@ impl<EK: KvEngine, ER: RaftEngine> StoreSystem<EK, ER> {
         pd_worker: LazyWorker<pd::Task>,
         sst_importer: Arc<SstImporter>,
         key_manager: Option<Arc<DataKeyManager>>,
+        grpc_service_mgr: GrpcServiceManager,
     ) -> Result<()>
     where
         T: Transport + 'static,
@@ -793,6 +795,7 @@ impl<EK: KvEngine, ER: RaftEngine> StoreSystem<EK, ER> {
             auto_split_controller,
             store_meta.lock().unwrap().region_read_progress.clone(),
             collector_reg_handle,
+            grpc_service_mgr,
             self.logger.clone(),
             self.shutdown.clone(),
             cfg.clone(),

components/raftstore-v2/src/worker/pd/mod.rs

@@ -19,6 +19,7 @@ use raftstore::store::{
     WriteStats, NUM_COLLECT_STORE_INFOS_PER_HEARTBEAT,
 };
 use resource_metering::{Collector, CollectorRegHandle, RawRecords};
+use service::service_manager::GrpcServiceManager;
 use slog::{error, warn, Logger};
 use tikv_util::{
     config::VersionTrack,
@@ -217,6 +218,9 @@ where
     // For slowness detection
     slowness_stats: slowness::SlownessStatistics,
 
+    // For grpc server.
+    grpc_service_manager: GrpcServiceManager,
+
     logger: Logger,
     shutdown: Arc<AtomicBool>,
     cfg: Arc<VersionTrack<Config>>,
@@ -242,6 +246,7 @@ where
         auto_split_controller: AutoSplitController,
         region_read_progress: RegionReadProgressRegistry,
         collector_reg_handle: CollectorRegHandle,
+        grpc_service_manager: GrpcServiceManager,
         logger: Logger,
         shutdown: Arc<AtomicBool>,
         cfg: Arc<VersionTrack<Config>>,
@@ -279,6 +284,7 @@ where
             concurrency_manager,
             causal_ts_provider,
             slowness_stats,
+            grpc_service_manager,
             logger,
             shutdown,
             cfg,

components/raftstore-v2/src/worker/pd/store.rs

@@ -231,6 +231,8 @@ where
         stats.set_cpu_usages(self.store_stat.store_cpu_usages.clone().into());
         stats.set_read_io_rates(self.store_stat.store_read_io_rates.clone().into());
         stats.set_write_io_rates(self.store_stat.store_write_io_rates.clone().into());
+        // Update grpc server status
+        stats.set_is_grpc_paused(self.grpc_service_manager.is_paused());
 
         let mut interval = pdpb::TimeInterval::default();
         interval.set_start_timestamp(self.store_stat.last_report_ts.into_inner());
@@ -268,6 +270,7 @@ where
 
         let resp = self.pd_client.store_heartbeat(stats, None, None);
         let logger = self.logger.clone();
+        let mut grpc_service_manager = self.grpc_service_manager.clone();
         let f = async move {
             match resp.await {
                 Ok(mut resp) => {
@@ -282,6 +285,27 @@ where
                             "Ignored AwakenRegions in raftstore-v2 as no hibernated regions in raftstore-v2"
                         );
                     }
+                    // Control grpc server.
+                    else if let Some(op) = resp.control_grpc.take() {
+                        info!(logger, "forcely control grpc server";
+                                "is_grpc_server_paused" => grpc_service_manager.is_paused(),
+                                "event" => ?op,
+                        );
+                        match op.get_ctrl_event() {
+                            pdpb::ControlGrpcEvent::Pause => {
+                                if let Err(e) = grpc_service_manager.pause() {
+                                    warn!(logger, "failed to send service event to PAUSE grpc server";
+                                        "err" => ?e);
+                                }
+                            }
+                            pdpb::ControlGrpcEvent::Resume => {
+                                if let Err(e) = grpc_service_manager.resume() {
+                                    warn!(logger, "failed to send service event to RESUME grpc server";
+                                        "err" => ?e);
+                                }
+                            }
+                        }
+                    }
                 }
                 Err(e) => {
                     error!(logger, "store heartbeat failed"; "err" => ?e);

components/raftstore-v2/tests/integrations/cluster.rs

@@ -45,6 +45,7 @@ use raftstore_v2::{
     Bootstrap, SimpleWriteEncoder, StateStorage, StoreSystem,
 };
 use resource_metering::CollectorRegHandle;
+use service::service_manager::GrpcServiceManager;
 use slog::{debug, o, Logger};
 use sst_importer::SstImporter;
 use tempfile::TempDir;
@@ -353,6 +354,7 @@ impl RunningState {
                 pd_worker,
                 importer,
                 key_manager,
+                GrpcServiceManager::dummy(),
             )
             .unwrap();
 

components/raftstore/Cargo.toml

@@ -77,6 +77,7 @@ resource_metering = { workspace = true }
 serde = "1.0"
 serde_derive = "1.0"
 serde_with = "1.4"
+service = { workspace = true }
 slog = { workspace = true }
 slog-global = { workspace = true }
 smallvec = "1.4"

components/raftstore/src/store/fsm/peer.rs

@@ -2824,7 +2824,9 @@ where
                 }
             }
             // It's v2 only message and ignore does no harm.
-            ExtraMessageType::MsgGcPeerResponse | ExtraMessageType::MsgFlushMemtable => (),
+            ExtraMessageType::MsgGcPeerResponse
+            | ExtraMessageType::MsgFlushMemtable
+            | ExtraMessageType::MsgRefreshBuckets => (),
         }
     }
 

components/raftstore/src/store/fsm/store.rs

@@ -45,6 +45,7 @@ use protobuf::Message;
 use raft::StateRole;
 use resource_control::{channel::unbounded, ResourceGroupManager};
 use resource_metering::CollectorRegHandle;
+use service::service_manager::GrpcServiceManager;
 use sst_importer::SstImporter;
 use tikv_alloc::trace::TraceEvent;
 use tikv_util::{
@@ -1576,6 +1577,7 @@ impl<EK: KvEngine, ER: RaftEngine> RaftBatchSystem<EK, ER> {
         collector_reg_handle: CollectorRegHandle,
         health_service: Option<HealthService>,
         causal_ts_provider: Option<Arc<CausalTsProviderImpl>>, // used for rawkv apiv2
+        grpc_service_mgr: GrpcServiceManager,
     ) -> Result<()> {
         assert!(self.workers.is_none());
         // TODO: we can get cluster meta regularly too later.
@@ -1715,6 +1717,7 @@ impl<EK: KvEngine, ER: RaftEngine> RaftBatchSystem<EK, ER> {
             health_service,
             causal_ts_provider,
             snap_generator_pool,
+            grpc_service_mgr,
         )?;
         Ok(())
     }
@@ -1733,6 +1736,7 @@ impl<EK: KvEngine, ER: RaftEngine> RaftBatchSystem<EK, ER> {
         health_service: Option<HealthService>,
         causal_ts_provider: Option<Arc<CausalTsProviderImpl>>, // used for rawkv apiv2
         snap_generator_pool: FuturePool,
+        grpc_service_mgr: GrpcServiceManager,
     ) -> Result<()> {
         let cfg = builder.cfg.value().clone();
         let store = builder.store.clone();
@@ -1824,6 +1828,7 @@ impl<EK: KvEngine, ER: RaftEngine> RaftBatchSystem<EK, ER> {
             health_service,
             coprocessor_host,
             causal_ts_provider,
+            grpc_service_mgr,
         );
         assert!(workers.pd_worker.start_with_timer(pd_runner));