From bc67dcf878a0dcd36ccf4dd03fc0959657cbc7d1 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sat, 7 Sep 2024 00:33:05 -0400 Subject: [PATCH] don't re-establish connection if we already have one --- .../count/count_all_model_usages_on_date.rs | 3 +++ crates/service/job/analytics_job/src/main.rs | 4 ++-- .../job/analytics_job/src/tasks/mod.rs | 2 +- .../update_model_usage_counts_table/mod.rs | 1 + .../update_model_usage_counts_table.rs} | 23 ++++++++++++++----- .../update_weekly_model_usage_counts/mod.rs | 1 - 6 files changed, 24 insertions(+), 10 deletions(-) create mode 100644 crates/service/job/analytics_job/src/tasks/update_model_usage_counts_table/mod.rs rename crates/service/job/analytics_job/src/tasks/{update_weekly_model_usage_counts/update_weekly_model_usage_counts.rs => update_model_usage_counts_table/update_model_usage_counts_table.rs} (77%) delete mode 100644 crates/service/job/analytics_job/src/tasks/update_weekly_model_usage_counts/mod.rs diff --git a/crates/lib/mysql_queries/src/queries/model_weights/count/count_all_model_usages_on_date.rs b/crates/lib/mysql_queries/src/queries/model_weights/count/count_all_model_usages_on_date.rs index 88e6551508..96a29af825 100644 --- a/crates/lib/mysql_queries/src/queries/model_weights/count/count_all_model_usages_on_date.rs +++ b/crates/lib/mysql_queries/src/queries/model_weights/count/count_all_model_usages_on_date.rs @@ -43,6 +43,9 @@ pub async fn count_all_model_usages_on_date<'e, 'c, E>( // See https://stackoverflow.com/a/14769096 // Basically it requires an interval of timestamps: // (f.created_at >= CURDATE() AND f.created_at < CURDATE() + INTERVAL 1 DAY) + // + // NB(3): We also join against model_weight_usage_counts to get the previously + // cached count. We should rewrite that value if there's a mismatch. let query = sqlx::query_as!( RawResult, r#" diff --git a/crates/service/job/analytics_job/src/main.rs b/crates/service/job/analytics_job/src/main.rs index eb33d36b4f..7e78aed56e 100644 --- a/crates/service/job/analytics_job/src/main.rs +++ b/crates/service/job/analytics_job/src/main.rs @@ -26,7 +26,7 @@ use errors::AnyhowResult; use crate::job_state::{JobState, SleepConfigs}; use crate::tasks::calculate_old_model_analytics::calculate_old_model_analytics_loop::calculate_old_model_analytics_loop; -use crate::tasks::update_weekly_model_usage_counts::update_weekly_model_usage_counts::update_weekly_model_usage_counts; +use crate::tasks::update_model_usage_counts_table::update_model_usage_counts_table::update_model_usage_counts_table; pub mod job_state; pub mod tasks; @@ -68,7 +68,7 @@ async fn main() -> AnyhowResult<()> { let job_state_2 = job_state.clone(); let handle_2 = tokio::task::spawn(async move { - let _r = update_weekly_model_usage_counts(job_state_2).await; + let _r = update_model_usage_counts_table(job_state_2).await; }); futures::future::join_all([ diff --git a/crates/service/job/analytics_job/src/tasks/mod.rs b/crates/service/job/analytics_job/src/tasks/mod.rs index d609e1f6c8..fababc4ff1 100644 --- a/crates/service/job/analytics_job/src/tasks/mod.rs +++ b/crates/service/job/analytics_job/src/tasks/mod.rs @@ -1,2 +1,2 @@ pub mod calculate_old_model_analytics; -pub mod update_weekly_model_usage_counts; \ No newline at end of file +pub mod update_model_usage_counts_table; \ No newline at end of file diff --git a/crates/service/job/analytics_job/src/tasks/update_model_usage_counts_table/mod.rs b/crates/service/job/analytics_job/src/tasks/update_model_usage_counts_table/mod.rs new file mode 100644 index 0000000000..46897b72f6 --- /dev/null +++ b/crates/service/job/analytics_job/src/tasks/update_model_usage_counts_table/mod.rs @@ -0,0 +1 @@ +pub mod update_model_usage_counts_table; \ No newline at end of file diff --git a/crates/service/job/analytics_job/src/tasks/update_weekly_model_usage_counts/update_weekly_model_usage_counts.rs b/crates/service/job/analytics_job/src/tasks/update_model_usage_counts_table/update_model_usage_counts_table.rs similarity index 77% rename from crates/service/job/analytics_job/src/tasks/update_weekly_model_usage_counts/update_weekly_model_usage_counts.rs rename to crates/service/job/analytics_job/src/tasks/update_model_usage_counts_table/update_model_usage_counts_table.rs index a3b3f0dba7..c248c79ba8 100644 --- a/crates/service/job/analytics_job/src/tasks/update_weekly_model_usage_counts/update_weekly_model_usage_counts.rs +++ b/crates/service/job/analytics_job/src/tasks/update_model_usage_counts_table/update_model_usage_counts_table.rs @@ -2,6 +2,8 @@ use std::time::Duration; use chrono::{Days, NaiveDate, Utc}; use log::{error, info}; +use sqlx::MySql; +use sqlx::pool::PoolConnection; use datetimes::generate_dates_inclusive::generate_dates_inclusive; use errors::AnyhowResult; @@ -10,11 +12,11 @@ use mysql_queries::queries::model_weights::count::count_all_model_usages_on_date use crate::job_state::JobState; -pub async fn update_weekly_model_usage_counts(job_state: JobState) -> AnyhowResult<()> { +pub async fn update_model_usage_counts_table(job_state: JobState) -> AnyhowResult<()> { loop { - info!("update weekly model usage counts"); + info!("update model usage counts table"); - match calculate_old_model_analytics(&job_state).await { + match connect_and_calculate_loop(&job_state).await { Ok(_) => { tokio::time::sleep(Duration::from_millis(job_state.sleep_config.between_job_batch_wait_millis)).await; } @@ -26,15 +28,24 @@ pub async fn update_weekly_model_usage_counts(job_state: JobState) -> AnyhowResu } } -pub async fn calculate_old_model_analytics(job_state: &JobState) -> AnyhowResult<()> { +async fn connect_and_calculate_loop(job_state: &JobState) -> AnyhowResult<()> { let mut connection = job_state.mysql_pool.acquire().await?; + loop { + calculate_with_connection(&mut connection, job_state).await?; + } +} + +async fn calculate_with_connection( + connection: &mut PoolConnection, + job_state: &JobState +) -> AnyhowResult<()> { let dates = get_dates(); for date in dates { info!("Finding usages for date: {:?}", date); - let usages = count_all_model_usages_on_date(&mut *connection, date).await?; + let usages = count_all_model_usages_on_date(&mut **connection, date).await?; info!("Records found: {}", usages.counts.len()); @@ -58,7 +69,7 @@ pub async fn calculate_old_model_analytics(job_state: &JobState) -> AnyhowResult date, usage_count: usage.latest_usage_count, insert_on_zero: false, - mysql_executor: &mut *connection, + mysql_executor: &mut **connection, phantom: Default::default(), }).await?; } diff --git a/crates/service/job/analytics_job/src/tasks/update_weekly_model_usage_counts/mod.rs b/crates/service/job/analytics_job/src/tasks/update_weekly_model_usage_counts/mod.rs deleted file mode 100644 index 0654e3d125..0000000000 --- a/crates/service/job/analytics_job/src/tasks/update_weekly_model_usage_counts/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod update_weekly_model_usage_counts; \ No newline at end of file