Merge pull request #6 from samply/feature/wait

wait count, logging
samply · Mar 14, 2024 · c7fb70d · c7fb70d
2 parents ceebff3 + 58db4b9
commit c7fb70d
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 9 deletions.
diff --git a/src/config.rs b/src/config.rs
@@ -46,10 +46,6 @@ struct CliArgs {
     #[clap(long, env, value_parser)]
     sites: String,
 
-    /// Wait for results count
-    #[clap(long, env, value_parser, default_value = "32")]
-    wait_count: usize,
-
     /// Credentials to use on the Beam Proxy
     #[clap(long, env, value_parser = parse_cors)]
     pub cors_origin: AllowOrigin,
@@ -73,7 +69,6 @@ pub(crate) struct Config {
     pub beam_app_id_long: AppId,
     pub api_key: String,
     pub sites: Vec<String>,
-    pub wait_count: usize,
     pub cors_origin: AllowOrigin,
     pub project: String,
     pub bind_addr: SocketAddr,
@@ -90,7 +85,6 @@ impl Config {
             beam_app_id_long: AppId::new_unchecked(cli_args.beam_app_id_long),
             api_key: cli_args.api_key,
             sites: cli_args.sites.split(';').map(|s| s.to_string()).collect(),
-            wait_count: cli_args.wait_count,
             cors_origin: cli_args.cors_origin,
             project: cli_args.project,
             bind_addr: cli_args.bind_addr,

diff --git a/src/main.rs b/src/main.rs
@@ -106,6 +106,8 @@ pub async fn main() {
         exit(2);
     }
 
+    info!("Beam ready");
+
     spawn_site_querying(shared_state.clone());
 
     let cors = CorsLayer::new()
@@ -185,14 +187,19 @@ async fn post_query(shared_state: SharedState, sites: Vec<String>) -> Result<(),
         info!("No sites to query");
         return Ok(());
     }
+    let wait_count = sites.len();
+    let site_display = sites.join(", ");
     let task = create_beam_task(sites);
+    info!("Querying sites {:?}", site_display);
     BEAM_CLIENT
         .post_task(&task)
         .await
         .map_err(|e| PrismError::BeamError(format!("Unable to post a query: {}", e)))?;
 
+    info!("Posted task {}", task.id);
+
     tokio::spawn(async move {
-        if let Err(e) = get_results(shared_state, task.id).await {
+        if let Err(e) = get_results(shared_state, task.id, wait_count).await {
             warn!("Failed to get results for {}: {e}", task.id);
         }
     });
@@ -224,13 +231,13 @@ async fn query_sites(
     Ok(())
 }
 
-async fn get_results(shared_state: SharedState, task_id: MsgId) -> Result<(), PrismError> {
+async fn get_results(shared_state: SharedState, task_id: MsgId, wait_count: usize) -> Result<(), PrismError> {
     let resp = BEAM_CLIENT
         .raw_beam_request(
             Method::GET,
             &format!(
                 "v1/tasks/{}/results?wait_count={}",
-                task_id, CONFIG.wait_count
+                task_id, wait_count
             ),
         )
         .header(
@@ -282,6 +289,7 @@ async fn get_results(shared_state: SharedState, task_id: MsgId) -> Result<(), Pr
             from.as_ref().split('.').nth(1).unwrap().to_string(), // extracting site name from app long name
             (criteria, std::time::SystemTime::now()),
         );
+        info!("Cached results from site {} for task {}", from.as_ref().split('.').nth(1).unwrap().to_string(), task_id);
     }
     Ok(())
 }