Added wait while shutdown to move repos

Summary: As titled

Reviewed By: RajivTS

Differential Revision: D85244548

fbshipit-source-id: d98772982275370199a2c14e14e32132e9176ee3

Author: Prashant Pal

eden/mononoke/cmdlib/sharding/src/facebook.rs

@@ -33,6 +33,7 @@ use stats::prelude::*;
 use tokio::runtime::Handle;
 use tokio::sync::RwLock;
 use tokio::sync::oneshot::Receiver;
+use tokio::sync::oneshot::Sender;
 use tokio::task::JoinHandle;
 use tokio::time;
 
@@ -898,8 +899,8 @@ impl ShardedProcessHandler {
         }
     }
 
-    pub async fn repo_map_empty(&self) -> bool {
-        self.repo_map.read().await.is_empty()
+    pub async fn repo_map_len(&self) -> usize {
+        self.repo_map.read().await.len()
     }
 }
 
@@ -1191,7 +1192,7 @@ impl ShardedProcessExecutor {
         logger: &Logger,
         terminate_signal_receiver: Receiver<bool>,
     ) -> Result<()> {
-        self.block_and_execute_with_quiesce_timeout(logger, terminate_signal_receiver, None)
+        self.block_and_execute_with_quiesce_timeout(logger, terminate_signal_receiver, None, None)
             .await
     }
 
@@ -1200,6 +1201,7 @@ impl ShardedProcessExecutor {
         logger: &Logger,
         terminate_signal_receiver: Receiver<bool>,
         quiesce_timeout: Option<Duration>,
+        quiesce_completion_sender: Option<Sender<bool>>,
     ) -> Result<()> {
         info!(logger, "Initiating sharded execution for service");
         let shards = self.client.get_my_shards()?;
@@ -1226,20 +1228,24 @@ impl ShardedProcessExecutor {
                         STATS::manual_shard_eviction_by_timeout.add_value(1);
                         break;
                     }
-                    repo_empty = self.handler.repo_map_empty() => {
-                        if repo_empty {
+                    repo_count = self.handler.repo_map_len() => {
+                        if repo_count == 0 {
                             info!(logger, "All repos moved, evicting...");
                             STATS::manual_shard_eviction_by_repomap.add_value(1);
                             break;
                         } else {
-                            info!(logger, "repos present, not evicting...");
+                            info!(logger, "repos present count={}, not evicting...", repo_count);
                             // Sleep a bit before next check to avoid busy loop
                             tokio::time::sleep(std::time::Duration::from_secs(1)).await;
                         }
                     }
                 }
             }
         }
+
+        if let Some(quiesce_completion_sender) = quiesce_completion_sender {
+            let _ = quiesce_completion_sender.send(true);
+        }
         Ok(())
     }
 }

eden/mononoke/cmdlib/sharding/src/oss.rs

@@ -14,6 +14,7 @@ use fbinit::FacebookInit;
 use slog::Logger;
 use tokio::runtime::Handle;
 use tokio::sync::oneshot::Receiver;
+use tokio::sync::oneshot::Sender;
 
 use crate::RepoShardedProcess;
 
@@ -46,6 +47,7 @@ impl ShardedProcessExecutor {
         _logger: &Logger,
         _terminate_signal_receiver: Receiver<bool>,
         _quiesce_timeout: Option<Duration>,
+        _quiesce_completion_sender: Option<Sender<bool>>,
     ) -> Result<()> {
         unimplemented!("ShardedProcessExecutor is supported only for fbcode build")
     }

eden/mononoke/git_server/src/main.rs

@@ -67,6 +67,7 @@ use mononoke_app::monitoring::MonitoringAppExtension;
 use ods_counters::OdsCounterManager;
 use rate_limiting::RateLimitEnvironment;
 use slog::info;
+use slog::warn;
 use tokio::net::TcpListener;
 
 use crate::middleware::Ods3Middleware;
@@ -220,6 +221,7 @@ fn main(fb: FacebookInit) -> Result<(), Error> {
     };
     let will_exit = Arc::new(AtomicBool::new(false));
     let (sender, receiver) = tokio::sync::oneshot::channel::<bool>();
+    let (quiesce_sender, quiesce_receiver) = tokio::sync::oneshot::channel::<bool>();
     let runtime = app.runtime().clone();
     // Service name is used for shallow or deep sharding. If sharding itself is disabled, provide
     // service name as None while opening repos.
@@ -252,6 +254,13 @@ fn main(fb: FacebookInit) -> Result<(), Error> {
         })
     };
 
+    let quiesce_timeout_secs = justknobs::get_as::<u64>(
+        "scm/mononoke:shardmanager_shutdown_timeout_secs",
+        Some("git_server"),
+    )
+    .unwrap();
+    let quiesce_timeout = std::time::Duration::from_secs(quiesce_timeout_secs);
+
     let requests_counter = Arc::new(AtomicI64::new(0));
     let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
     let server = {
@@ -352,17 +361,12 @@ fn main(fb: FacebookInit) -> Result<(), Error> {
                     let logger = app.logger().clone();
                     {
                         async move {
-                            let timeout_secs = justknobs::get_as::<u64>(
-                                "scm/mononoke:shardmanager_shutdown_timeout_secs",
-                                Some("git_server"),
-                            )
-                            .unwrap();
-                            let quiesce_timeout = std::time::Duration::from_secs(timeout_secs);
                             executor
                                 .block_and_execute_with_quiesce_timeout(
                                     &logger,
                                     receiver,
                                     Some(quiesce_timeout),
+                                    Some(quiesce_sender),
                                 )
                                 .await
                         }
@@ -398,6 +402,7 @@ fn main(fb: FacebookInit) -> Result<(), Error> {
             Ok(())
         }
     };
+    let cloned_logger = app.logger().clone();
     app.run_until_terminated(
         server,
         move || {
@@ -406,14 +411,21 @@ fn main(fb: FacebookInit) -> Result<(), Error> {
         },
         args.shutdown_timeout_args.shutdown_grace_period,
         async move {
+            match quiesce_receiver.await {
+                Ok(_) => info!(
+                    &cloned_logger.clone(),
+                    "received signal from quiesce sender"
+                ),
+                Err(_) => warn!(&cloned_logger, "quiesce sender dropped"),
+            };
             let _ = shutdown_tx.send(());
             // Currently we kill off in-flight requests as soon as we've closed the listener.
             // If this is a problem in prod, this would be the point at which to wait
             // for all connections to shut down.
             // To do this properly, we'd need to track the `Connection` futures that Gotham
             // gets from Hyper, tell them to gracefully shutdown, then wait for them to complete
         },
-        args.shutdown_timeout_args.shutdown_timeout,
+        args.shutdown_timeout_args.shutdown_timeout + quiesce_timeout,
         // TODO
         Some(requests_counter),
     )?;

eden/mononoke/mononoke_macros/just_knobs_defaults/just_knobs.json

@@ -105,6 +105,7 @@
     "scm/mononoke:retry_query_from_replica_with_consistency_check_jitter": 50,
     "scm/mononoke:retry_query_from_replica_with_consistency_check_interval_ms": 200,
     "scm/mononoke:retry_query_from_replica_with_consistency_check_hlc_drift_tolerance_ns": 0,
-    "scm/mononoke:cgdm_reloading_interval_secs": 5
+    "scm/mononoke:cgdm_reloading_interval_secs": 5,
+    "scm/mononoke:shardmanager_shutdown_timeout_secs": 10
   }
 }

eden/mononoke/scs/scs_server/src/main.rs

@@ -51,6 +51,7 @@ use panichandler::Fate;
 use scs_methods::source_control_impl::SourceControlServiceImpl;
 use sharding_ext::RepoShard;
 use slog::info;
+use slog::warn;
 use source_control_services::make_SourceControlService_server;
 use sql_construct::SqlConstruct;
 use sql_storage::Destination;
@@ -285,6 +286,7 @@ fn main(fb: FacebookInit) -> Result<(), Error> {
 
     let will_exit = Arc::new(AtomicBool::new(false));
     let (sm_shutdown_sender, sm_shutdown_receiver) = tokio::sync::oneshot::channel::<bool>();
+    let (quiesce_sender, quiesce_receiver) = tokio::sync::oneshot::channel::<bool>();
 
     if let Some(max_memory) = args.max_memory {
         memory::set_max_memory(max_memory);
@@ -398,6 +400,13 @@ fn main(fb: FacebookInit) -> Result<(), Error> {
         writer.write_all(b"\n")?;
     }
 
+    let timeout_secs = justknobs::get_as::<u64>(
+        "scm/mononoke:shardmanager_shutdown_timeout_secs",
+        Some("scs_server"),
+    )
+    .unwrap();
+    let quiesce_timeout = std::time::Duration::from_secs(timeout_secs);
+
     if let Some(executor) = args.sharded_executor_args.build_executor(
         fb,
         runtime.clone(),
@@ -411,18 +420,13 @@ fn main(fb: FacebookInit) -> Result<(), Error> {
         runtime.spawn({
             let logger = logger.clone();
             {
-                let timeout_secs = justknobs::get_as::<u64>(
-                    "scm/mononoke:shardmanager_shutdown_timeout_secs",
-                    Some("scs_server"),
-                )
-                .unwrap();
-                let quiesce_timeout = std::time::Duration::from_secs(timeout_secs);
                 async move {
                     executor
                         .block_and_execute_with_quiesce_timeout(
                             &logger,
                             sm_shutdown_receiver,
                             Some(quiesce_timeout),
+                            Some(quiesce_sender),
                         )
                         .await
                 }
@@ -434,6 +438,7 @@ fn main(fb: FacebookInit) -> Result<(), Error> {
     // stats aggregation.
     app.start_stats_aggregation()?;
 
+    let cloned_logger = logger.clone();
     app.wait_until_terminated(
         move || {
             let _ = sm_shutdown_sender.send(true);
@@ -442,13 +447,20 @@ fn main(fb: FacebookInit) -> Result<(), Error> {
         args.shutdown_timeout_args.shutdown_grace_period,
         async {
             // Note that async blocks are lazy, so this isn't called until first poll
+            match quiesce_receiver.await {
+                Ok(_) => info!(
+                    &cloned_logger.clone(),
+                    "received signal from quiesce sender"
+                ),
+                Err(_) => warn!(&cloned_logger, "quiesce sender dropped"),
+            };
             let _ = task::spawn_blocking(move || {
                 // Calling `stop` blocks until the service has completed all requests.
                 service_framework.stop();
             })
             .await;
         },
-        args.shutdown_timeout_args.shutdown_timeout,
+        args.shutdown_timeout_args.shutdown_timeout + quiesce_timeout,
         None,
     )?;