Merge pull request #8322 from systeminit/victor/eng-3338-migrate-func-runs-and-func-run-logs-back-to-the-si-database

feat: add func runs and func run logs migration to sdf

Author: vbustamante

Cargo.lock

@@ -164,7 +164,7 @@ pub(crate) struct Args {
     #[arg(long, env = "SI_BACKFILL_CACHE_TYPES")]
     pub(crate) backfill_cache_types: Option<String>,
 
-    /// Key batch size for PostgreSQL queries during backfill
+    /// Batch size for PostgreSQL queries during backfill (layer cache and func runs)
     #[arg(long, default_value = "1000", env = "SI_BACKFILL_KEY_BATCH_SIZE")]
     pub(crate) backfill_key_batch_size: usize,
 
@@ -180,6 +180,14 @@ pub(crate) struct Args {
     #[arg(long, default_value = "5", env = "SI_BACKFILL_MAX_CONCURRENT_UPLOADS")]
     pub(crate) backfill_max_concurrent_uploads: usize,
 
+    /// Cutoff ID for func runs backfill (optional, start from a specific func run ID)
+    #[arg(long, env = "SI_BACKFILL_FUNC_RUNS_CUTOFF_ID")]
+    pub(crate) backfill_func_runs_cutoff_id: Option<String>,
+
+    /// Cutoff ID for func run logs backfill (optional, start from a specific func run log ID)
+    #[arg(long, env = "SI_BACKFILL_FUNC_RUN_LOGS_CUTOFF_ID")]
+    pub(crate) backfill_func_run_logs_cutoff_id: Option<String>,
+
     /// Veritech encryption key file location [default: /run/sdf/veritech_encryption.key]
     #[arg(long)]
     pub(crate) veritech_encryption_key_path: Option<PathBuf>,
@@ -502,6 +510,17 @@ fn build_config_map(args: Args, config_map: &mut ConfigMap) -> &ConfigMap {
         i64::try_from(args.backfill_max_concurrent_uploads).unwrap_or(5),
     );
 
+    if let Some(backfill_func_runs_cutoff_id) = args.backfill_func_runs_cutoff_id {
+        config_map.set("backfill_func_runs_cutoff_id", backfill_func_runs_cutoff_id);
+    }
+
+    if let Some(backfill_func_run_logs_cutoff_id) = args.backfill_func_run_logs_cutoff_id {
+        config_map.set(
+            "backfill_func_run_logs_cutoff_id",
+            backfill_func_run_logs_cutoff_id,
+        );
+    }
+
     config_map.set("nats.connection_name", NAME);
     config_map.set("pg.application_name", NAME);
     config_map.set("layer_db_config.pg_pool_config.application_name", NAME);

bin/sdf/src/args.rs

@@ -9,6 +9,7 @@ use innit_client::InnitClient;
 use sdf_server::{
     BackfillConfig,
     Config,
+    FuncRunsBackfiller,
     LayerCacheBackfiller,
     Migrator,
     Server,
@@ -157,6 +158,18 @@ async fn async_main() -> Result<()> {
                 telemetry_shutdown,
             )
             .await
+        } else if config.migration_mode().is_backfill_func_runs() {
+            backfill_func_runs(
+                config,
+                main_tracker,
+                main_token,
+                helping_tasks_tracker,
+                helping_tasks_token,
+                telemetry_tracker,
+                telemetry_token,
+                telemetry_shutdown,
+            )
+            .await
         } else {
             run_server(
                 config,
@@ -396,3 +409,47 @@ async fn generate_symmetric_key(
         .await
         .map_err(Into::into)
 }
+
+#[inline]
+#[allow(clippy::too_many_arguments)]
+async fn backfill_func_runs(
+    config: Config,
+    main_tracker: TaskTracker,
+    main_token: CancellationToken,
+    helping_tasks_tracker: TaskTracker,
+    helping_tasks_token: CancellationToken,
+    telemetry_tracker: TaskTracker,
+    telemetry_token: CancellationToken,
+    telemetry_shutdown: TelemetryShutdownGuard,
+) -> Result<()> {
+    // Extract parameters before config is moved
+    let func_run_cutoff_id = config
+        .backfill_func_runs_cutoff_id()
+        .and_then(|id_str| id_str.parse().ok());
+    let func_run_log_cutoff_id = config
+        .backfill_func_run_logs_cutoff_id()
+        .and_then(|id_str| id_str.parse().ok());
+    let batch_size = config.backfill_key_batch_size() as i64;
+
+    let handle = main_tracker.spawn(
+        FuncRunsBackfiller::upload_all_func_runs_and_logs_concurrently(
+            config,
+            helping_tasks_tracker.clone(),
+            helping_tasks_token.clone(),
+            main_token.clone(),
+            func_run_cutoff_id,
+            func_run_log_cutoff_id,
+            batch_size,
+        ),
+    );
+
+    shutdown::graceful_with_handle(handle)
+        .group(main_tracker, main_token)
+        .group(helping_tasks_tracker, helping_tasks_token)
+        .group(telemetry_tracker, telemetry_token)
+        .telemetry_guard(telemetry_shutdown.into_future())
+        .timeout(GRACEFUL_SHUTDOWN_TIMEOUT)
+        .wait()
+        .await
+        .map_err(Into::into)
+}