[batch] Fix attempt resources after update trigger to use deduped res…

…ource id (#13083) Must go in after #13082.
hail-is · May 25, 2023 · b0c8dea · b0c8dea
1 parent 8294fd4
commit b0c8dea
Show file tree

Hide file tree

Showing 3 changed files with 134 additions and 4 deletions.
diff --git a/batch/sql/estimated-current.sql b/batch/sql/estimated-current.sql
@@ -540,7 +540,7 @@ BEGIN
 
     INSERT INTO aggregated_billing_project_user_resources_v3 (billing_project, user, resource_id, token, `usage`)
     SELECT batches.billing_project, batches.`user`,
-      attempt_resources.resource_id,
+      attempt_resources.deduped_resource_id,
       rand_token,
       msec_diff_rollup * quantity
     FROM attempt_resources
@@ -564,7 +564,7 @@ BEGIN
 
     INSERT INTO aggregated_batch_resources_v3 (batch_id, resource_id, token, `usage`)
     SELECT attempt_resources.batch_id,
-      attempt_resources.resource_id,
+      attempt_resources.deduped_resource_id,
       rand_token,
       msec_diff_rollup * quantity
     FROM attempt_resources
@@ -585,7 +585,7 @@ BEGIN
 
     INSERT INTO aggregated_job_resources_v3 (batch_id, job_id, resource_id, `usage`)
     SELECT attempt_resources.batch_id, attempt_resources.job_id,
-      attempt_resources.resource_id,
+      attempt_resources.deduped_resource_id,
       msec_diff_rollup * quantity
     FROM attempt_resources
     JOIN aggregated_job_resources_v2 ON
@@ -611,7 +611,7 @@ BEGIN
     SELECT cur_billing_date,
       batches.billing_project,
       batches.`user`,
-      attempt_resources.resource_id,
+      attempt_resources.deduped_resource_id,
       rand_token,
       msec_diff_rollup * quantity
     FROM attempt_resources

diff --git a/batch/sql/fix-billing-triggers-deduped-resource-id.sql b/batch/sql/fix-billing-triggers-deduped-resource-id.sql
@@ -0,0 +1,127 @@
+DELIMITER $$
+
+DROP TRIGGER IF EXISTS attempts_after_update $$
+CREATE TRIGGER attempts_after_update AFTER UPDATE ON attempts
+FOR EACH ROW
+BEGIN
+  DECLARE job_cores_mcpu INT;
+  DECLARE cur_billing_project VARCHAR(100);
+  DECLARE msec_diff_rollup BIGINT;
+  DECLARE cur_n_tokens INT;
+  DECLARE rand_token INT;
+  DECLARE cur_billing_date DATE;
+
+  SELECT n_tokens INTO cur_n_tokens FROM globals LOCK IN SHARE MODE;
+  SET rand_token = FLOOR(RAND() * cur_n_tokens);
+
+  SELECT cores_mcpu INTO job_cores_mcpu FROM jobs
+  WHERE batch_id = NEW.batch_id AND job_id = NEW.job_id;
+
+  SELECT billing_project INTO cur_billing_project FROM batches WHERE id = NEW.batch_id;
+
+  SET msec_diff_rollup = (GREATEST(COALESCE(NEW.rollup_time - NEW.start_time, 0), 0) -
+                          GREATEST(COALESCE(OLD.rollup_time - OLD.start_time, 0), 0));
+
+  SET cur_billing_date = CAST(UTC_DATE() AS DATE);
+
+  IF msec_diff_rollup != 0 THEN
+    INSERT INTO aggregated_billing_project_user_resources_v2 (billing_project, user, resource_id, token, `usage`)
+    SELECT billing_project, `user`,
+      resource_id,
+      rand_token,
+      msec_diff_rollup * quantity
+    FROM attempt_resources
+    JOIN batches ON batches.id = attempt_resources.batch_id
+    WHERE batch_id = NEW.batch_id AND job_id = NEW.job_id AND attempt_id = NEW.attempt_id
+    ON DUPLICATE KEY UPDATE `usage` = `usage` + msec_diff_rollup * quantity;
+
+    INSERT INTO aggregated_billing_project_user_resources_v3 (billing_project, user, resource_id, token, `usage`)
+    SELECT batches.billing_project, batches.`user`,
+      attempt_resources.deduped_resource_id,
+      rand_token,
+      msec_diff_rollup * quantity
+    FROM attempt_resources
+    JOIN batches ON batches.id = attempt_resources.batch_id
+    INNER JOIN aggregated_billing_project_user_resources_v2 ON
+      aggregated_billing_project_user_resources_v2.billing_project = batches.billing_project AND
+      aggregated_billing_project_user_resources_v2.user = batches.user AND
+      aggregated_billing_project_user_resources_v2.resource_id = attempt_resources.resource_id AND
+      aggregated_billing_project_user_resources_v2.token = rand_token
+    WHERE attempt_resources.batch_id = NEW.batch_id AND attempt_resources.job_id = NEW.job_id AND attempt_id = NEW.attempt_id AND migrated = 1
+    ON DUPLICATE KEY UPDATE `usage` = aggregated_billing_project_user_resources_v3.`usage` + msec_diff_rollup * quantity;
+
+    INSERT INTO aggregated_batch_resources_v2 (batch_id, resource_id, token, `usage`)
+    SELECT batch_id,
+      resource_id,
+      rand_token,
+      msec_diff_rollup * quantity
+    FROM attempt_resources
+    WHERE batch_id = NEW.batch_id AND job_id = NEW.job_id AND attempt_id = NEW.attempt_id
+    ON DUPLICATE KEY UPDATE `usage` = `usage` + msec_diff_rollup * quantity;
+
+    INSERT INTO aggregated_batch_resources_v3 (batch_id, resource_id, token, `usage`)
+    SELECT attempt_resources.batch_id,
+      attempt_resources.deduped_resource_id,
+      rand_token,
+      msec_diff_rollup * quantity
+    FROM attempt_resources
+    JOIN aggregated_batch_resources_v2 ON
+      aggregated_batch_resources_v2.batch_id = attempt_resources.batch_id AND
+      aggregated_batch_resources_v2.resource_id = attempt_resources.resource_id AND
+      aggregated_batch_resources_v2.token = rand_token
+    WHERE attempt_resources.batch_id = NEW.batch_id AND attempt_resources.job_id = NEW.job_id AND attempt_id = NEW.attempt_id AND migrated = 1
+    ON DUPLICATE KEY UPDATE `usage` = aggregated_batch_resources_v3.`usage` + msec_diff_rollup * quantity;
+
+    INSERT INTO aggregated_job_resources_v2 (batch_id, job_id, resource_id, `usage`)
+    SELECT batch_id, job_id,
+      resource_id,
+      msec_diff_rollup * quantity
+    FROM attempt_resources
+    WHERE batch_id = NEW.batch_id AND job_id = NEW.job_id AND attempt_id = NEW.attempt_id
+    ON DUPLICATE KEY UPDATE `usage` = `usage` + msec_diff_rollup * quantity;
+
+    INSERT INTO aggregated_job_resources_v3 (batch_id, job_id, resource_id, `usage`)
+    SELECT attempt_resources.batch_id, attempt_resources.job_id,
+      attempt_resources.deduped_resource_id,
+      msec_diff_rollup * quantity
+    FROM attempt_resources
+    JOIN aggregated_job_resources_v2 ON
+      aggregated_job_resources_v2.batch_id = attempt_resources.batch_id AND
+      aggregated_job_resources_v2.job_id = attempt_resources.job_id AND
+      aggregated_job_resources_v2.resource_id = attempt_resources.resource_id
+    WHERE attempt_resources.batch_id = NEW.batch_id AND attempt_resources.job_id = NEW.job_id AND attempt_id = NEW.attempt_id AND migrated = 1
+    ON DUPLICATE KEY UPDATE `usage` = aggregated_job_resources_v3.`usage` + msec_diff_rollup * quantity;
+
+    INSERT INTO aggregated_billing_project_user_resources_by_date_v2 (billing_date, billing_project, user, resource_id, token, `usage`)
+    SELECT cur_billing_date,
+      billing_project,
+      `user`,
+      resource_id,
+      rand_token,
+      msec_diff_rollup * quantity
+    FROM attempt_resources
+    JOIN batches ON batches.id = attempt_resources.batch_id
+    WHERE batch_id = NEW.batch_id AND job_id = NEW.job_id AND attempt_id = NEW.attempt_id
+    ON DUPLICATE KEY UPDATE `usage` = `usage` + msec_diff_rollup * quantity;
+
+    INSERT INTO aggregated_billing_project_user_resources_by_date_v3 (billing_date, billing_project, user, resource_id, token, `usage`)
+    SELECT cur_billing_date,
+      batches.billing_project,
+      batches.`user`,
+      attempt_resources.deduped_resource_id,
+      rand_token,
+      msec_diff_rollup * quantity
+    FROM attempt_resources
+    JOIN batches ON batches.id = attempt_resources.batch_id
+    JOIN aggregated_billing_project_user_resources_by_date_v2 ON
+      aggregated_billing_project_user_resources_by_date_v2.billing_date = cur_billing_date AND
+      aggregated_billing_project_user_resources_by_date_v2.billing_project = batches.billing_project AND
+      aggregated_billing_project_user_resources_by_date_v2.user = batches.user AND
+      aggregated_billing_project_user_resources_by_date_v2.resource_id = attempt_resources.resource_id AND
+      aggregated_billing_project_user_resources_by_date_v2.token = rand_token
+    WHERE attempt_resources.batch_id = NEW.batch_id AND attempt_resources.job_id = NEW.job_id AND attempt_id = NEW.attempt_id AND migrated = 1
+    ON DUPLICATE KEY UPDATE `usage` = aggregated_billing_project_user_resources_by_date_v3.`usage` + msec_diff_rollup * quantity;
+  END IF;
+END $$
+
+DELIMITER ;
diff --git a/build.yaml b/build.yaml
@@ -2198,6 +2198,9 @@ steps:
       - name: set-test-and-dev-pools-to-16-core-max-3
         script: /io/sql/set-test-and-dev-pools-to-16-core-max-3.py
         online: true
+      - name: fix-billing-triggers-deduped-resource-id
+        script: /io/sql/fix-billing-triggers-deduped-resource-id.sql
+        online: true
     inputs:
       - from: /repo/batch/sql
         to: /io/sql