From b24c6a50b7d85ec42b356834709458d3ab17ecb5 Mon Sep 17 00:00:00 2001
From: Tristan Vuong <85768771+tristanvuong2021@users.noreply.github.com>
Date: Thu, 5 Dec 2024 08:38:58 -0800
Subject: [PATCH] perf: Recreate metric reuse index (#1949)

The current index isn't actually capable of being used and is leading to
sequential scans. The new index along with the corresponding query
changes allows the index to actually be used, cutting the execution time
to 1/10th of what it is currently.
---
 .../v2/postgres/readers/MetricReader.kt       | 32 ++++++-------
 .../v2/postgres/writers/CreateReport.kt       |  5 ++-
 .../reporting/postgres/changelog-v2.yaml      |  3 ++
 .../postgres/recreate-metric-reuse-index.sql  | 45 +++++++++++++++++++
 4 files changed, 68 insertions(+), 17 deletions(-)
 create mode 100644 src/main/resources/reporting/postgres/recreate-metric-reuse-index.sql

diff --git a/src/main/kotlin/org/wfanet/measurement/reporting/deploy/v2/postgres/readers/MetricReader.kt b/src/main/kotlin/org/wfanet/measurement/reporting/deploy/v2/postgres/readers/MetricReader.kt
index 13be9b177bd..bada1a5aec7 100644
--- a/src/main/kotlin/org/wfanet/measurement/reporting/deploy/v2/postgres/readers/MetricReader.kt
+++ b/src/main/kotlin/org/wfanet/measurement/reporting/deploy/v2/postgres/readers/MetricReader.kt
@@ -71,6 +71,7 @@ class MetricReader(private val readContext: ReadContext) {
     val metricSpec: MetricSpec,
     val metricDetails: Metric.Details,
     val createTime: Instant,
+    val state: Metric.State,
   )
 
   private data class MetricInfo(
@@ -344,37 +345,34 @@ class MetricReader(private val readContext: ReadContext) {
       """
         .trimIndent()
 
-    val sqlJoins: String =
-      """
-      JOIN Metrics USING(MeasurementConsumerId)
-      JOIN MetricCalculationSpecReportingMetrics USING(MeasurementConsumerId, MetricId)
-      """
-        .trimIndent()
-
     val sql =
       StringBuilder(
         """
-          $sqlSelect
-          FROM
-            MeasurementConsumers
-            $sqlJoins
-          WHERE Metrics.MeasurementConsumerId = $1
+          WITH ReusableMetricIds AS (
+            SELECT MetricId
+            FROM
+              Metrics JOIN MetricCalculationSpecReportingMetrics USING(MeasurementConsumerId, MetricId)
+            WHERE Metrics.MeasurementConsumerId = $1
             AND (
               Metrics.ReportingSetId,
               Metrics.TimeIntervalStart,
               Metrics.TimeIntervalEndExclusive,
               MetricCalculationSpecReportingMetrics.MetricCalculationSpecId
             ) IN (VALUES ${ValuesListBoundStatement.VALUES_LIST_PLACEHOLDER})
-            AND (Metrics.State = $2 OR Metrics.State = $3)
+          )
+          $sqlSelect
+          FROM
+            MeasurementConsumers
+            JOIN Metrics USING(MeasurementConsumerId)
+            JOIN MetricCalculationSpecReportingMetrics USING(MeasurementConsumerId, MetricId)
+          WHERE Metrics.MetricId IN (SELECT MetricId FROM ReusableMetricIDs)
         """
           .trimIndent()
       )
 
     val statement =
-      valuesListBoundStatement(valuesStartIndex = 3, paramCount = 4, sql.toString()) {
+      valuesListBoundStatement(valuesStartIndex = 1, paramCount = 4, sql.toString()) {
         bind("$1", measurementConsumerId)
-        bind("$2", Metric.State.SUCCEEDED)
-        bind("$3", Metric.State.RUNNING)
         reportingMetricKeys.forEach {
           addValuesBinding {
             bindValuesParam(0, it.reportingSetId)
@@ -411,6 +409,7 @@ class MetricReader(private val readContext: ReadContext) {
       val metricCalculationSpecId: InternalId = row["MetricCalculationSpecId"]
       val metricId: InternalId = row["MetricId"]
       val externalMetricId: String = row["ExternalMetricId"]
+      val state: Metric.State = row.getProtoEnum("State", Metric.State::forNumber)
       val metricTimeIntervalStart: Instant = row["MetricsTimeIntervalStart"]
       val metricTimeIntervalEnd: Instant = row["MetricsTimeIntervalEndExclusive"]
       val metricType: MetricSpec.TypeCase = MetricSpec.TypeCase.forNumber(row["MetricType"])
@@ -481,6 +480,7 @@ class MetricReader(private val readContext: ReadContext) {
           metricSpec = metricSpec,
           metricDetails = metricDetails,
           createTime = createTime,
+          state = state,
         )
       }
     }
diff --git a/src/main/kotlin/org/wfanet/measurement/reporting/deploy/v2/postgres/writers/CreateReport.kt b/src/main/kotlin/org/wfanet/measurement/reporting/deploy/v2/postgres/writers/CreateReport.kt
index 7d01f800681..ad2d033d47a 100644
--- a/src/main/kotlin/org/wfanet/measurement/reporting/deploy/v2/postgres/writers/CreateReport.kt
+++ b/src/main/kotlin/org/wfanet/measurement/reporting/deploy/v2/postgres/writers/CreateReport.kt
@@ -187,7 +187,10 @@ class CreateReport(private val request: CreateReportRequest) : PostgresWriter<Re
                 it.metricDetails,
               )
             val oldValue: MetricReader.ReportingMetric? = get(key)
-            if (oldValue == null || it.createTime.isAfter(oldValue.createTime)) {
+            if (
+              it.state != Metric.State.FAILED &&
+                (oldValue == null || it.createTime.isAfter(oldValue.createTime))
+            ) {
               put(key, it)
             }
           }
diff --git a/src/main/resources/reporting/postgres/changelog-v2.yaml b/src/main/resources/reporting/postgres/changelog-v2.yaml
index ef4635959b4..86628f66dad 100644
--- a/src/main/resources/reporting/postgres/changelog-v2.yaml
+++ b/src/main/resources/reporting/postgres/changelog-v2.yaml
@@ -55,3 +55,6 @@ databaseChangeLog:
 - include:
     file: add-additional-metric-spec-columns-to-metrics.sql
     relativeToChangeLogFile: true
+- include:
+    file: recreate-metric-reuse-index.sql
+    relativeToChangeLogFile: true
diff --git a/src/main/resources/reporting/postgres/recreate-metric-reuse-index.sql b/src/main/resources/reporting/postgres/recreate-metric-reuse-index.sql
new file mode 100644
index 00000000000..315e1c9c812
--- /dev/null
+++ b/src/main/resources/reporting/postgres/recreate-metric-reuse-index.sql
@@ -0,0 +1,45 @@
+-- liquibase formatted sql
+
+-- Copyright 2024 The Cross-Media Measurement Authors
+--
+-- Licensed under the Apache License, Version 2.0 (the "License");
+-- you may not use this file except in compliance with the License.
+-- You may obtain a copy of the License at
+--
+--      http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing, software
+-- distributed under the License is distributed on an "AS IS" BASIS,
+-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+-- See the License for the specific language governing permissions and
+-- limitations under the License.
+
+-- Postgres database schema for the Reporting server.
+--
+-- Table hierarchy:
+--   Root
+--   └── MeasurementConsumers
+--       ├── EventGroups
+--       ├── ReportingSets
+--       │   ├── ReportingSetEventGroups
+--       │   ├── PrimitiveReportingSetBases
+--       │   │   └── PrimitiveReportingSetBasisFilters
+--       │   ├── SetExpressions
+--       │   └── WeightedSubsetUnions
+--       │       └── WeightedSubsetUnionPrimitiveReportingSetBases
+--       ├── Metrics
+--       │   └── MetricMeasurements
+--       ├── Measurements
+--       │   └── MeasurementPrimitiveReportingSetBases
+--       ├── MetricCalculationSpecs
+--       ├── Reports
+--       │   ├── ReportTimeIntervals
+--       │   └── MetricCalculationSpecReportingMetrics
+--       └── ReportSchedules
+--           ├── ReportScheduleIterations
+--           └── ReportsReportSchedules
+
+-- changeset tristanvuong2021:recreate-metrics-comparison-index dbms:postgresl
+DROP INDEX metrics_comparison;
+CREATE INDEX metrics_comparison
+  ON Metrics (MeasurementConsumerId, ReportingSetId, TimeIntervalStart, TimeIntervalEndExclusive, MetricId);