Skip to content

Commit

Permalink
Add graceful shutdown during DB loading
Browse files Browse the repository at this point in the history
  • Loading branch information
danielle-tfh committed Jan 21, 2025
1 parent f4dbe47 commit 545a003
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 2 deletions.
2 changes: 1 addition & 1 deletion .github/workflows/temp-branch-build-and-push.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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 }}'
Expand Down
8 changes: 7 additions & 1 deletion iris-mpc/src/bin/server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 || {
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -1199,6 +1201,10 @@ async fn server_main(config: Config) -> eyre::Result<()> {
elapsed,
record_counter as f64 / elapsed.as_secs_f64()
);
if download_shutdown_handler.is_shutting_down() {
tracing::warn!("Shutdown requested by shutdown_handler.");
return Err(eyre::eyre!("Shutdown requested"));
}
}

// if the serial id hasn't been loaded before, count is as unique record
Expand Down Expand Up @@ -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,
Expand Down

0 comments on commit 545a003

Please sign in to comment.