Add graceful shutdown during DB loading

.github/workflows/temp-branch-build-and-push.yaml

@@ -3,7 +3,7 @@ name: Branch - Build and push docker image
 on:
   push:
     branches:
-      - "add-error-handling-skipped-requests"
+      - "add-graceful-shutdown-during-s3-sync"
 
 concurrency:
   group: '${{ github.workflow }} @ ${{ github.event.pull_request.head.label || github.head_ref || github.ref }}'

iris-mpc/src/bin/server.rs

@@ -989,6 +989,7 @@ async fn server_main(config: Config) -> eyre::Result<()> {
     let load_chunks_parallelism = config.load_chunks_parallelism;
     let db_chunks_bucket_name = config.db_chunks_bucket_name.clone();
     let db_chunks_folder_name = config.db_chunks_folder_name.clone();
+    let download_shutdown_handler = Arc::clone(&shutdown_handler);
 
     let (tx, rx) = oneshot::channel();
     background_tasks.spawn_blocking(move || {
@@ -1068,6 +1069,7 @@ async fn server_main(config: Config) -> eyre::Result<()> {
                         "Initialize iris db: Loading from DB (parallelism: {})",
                         parallelism
                     );
+                    let download_shutdown_handler = Arc::clone(&download_shutdown_handler);
                     let s3_store = S3Store::new(s3_client_clone, db_chunks_bucket_name);
                     tokio::runtime::Handle::current().block_on(async {
                         let mut stream = match config.enable_s3_importer {
@@ -1148,6 +1150,10 @@ async fn server_main(config: Config) -> eyre::Result<()> {
                         let mut n_loaded_from_db = 0;
                         let mut n_loaded_from_s3 = 0;
                         while let Some(result) = stream.try_next().await? {
+                            if download_shutdown_handler.is_shutting_down() {
+                                tracing::warn!("Shutdown requested by shutdown_handler.");
+                                return Err(eyre::eyre!("Shutdown requested"));
+                            }
                             time_waiting_for_stream += now_load_summary.elapsed();
                             now_load_summary = Instant::now();
                             let index = result.index();
@@ -1282,7 +1288,7 @@ async fn server_main(config: Config) -> eyre::Result<()> {
     let sns_client_bg = sns_client.clone();
     let config_bg = config.clone();
     let store_bg = store.clone();
-    let shutdown_handler_bg = shutdown_handler.clone();
+    let shutdown_handler_bg = Arc::clone(&shutdown_handler);
     let _result_sender_abort = background_tasks.spawn(async move {
         while let Some(ServerJobResult {
             merged_results,