From 6dfc562e5f3d28b521168e11b09739ce8ae94b6c Mon Sep 17 00:00:00 2001
From: Rieman Li <riemanli@google.com>
Date: Sat, 22 Apr 2023 00:13:28 +0000
Subject: [PATCH] Use an extension property to get the state from an internal
 metric.

---
 .../service/api/v2alpha/MetricsService.kt     | 32 +++++++++----------
 .../internal/reporting/v2/metric.proto        | 11 -------
 2 files changed, 16 insertions(+), 27 deletions(-)

diff --git a/src/main/kotlin/org/wfanet/measurement/reporting/service/api/v2alpha/MetricsService.kt b/src/main/kotlin/org/wfanet/measurement/reporting/service/api/v2alpha/MetricsService.kt
index 692b7376718..01045e9b85f 100644
--- a/src/main/kotlin/org/wfanet/measurement/reporting/service/api/v2alpha/MetricsService.kt
+++ b/src/main/kotlin/org/wfanet/measurement/reporting/service/api/v2alpha/MetricsService.kt
@@ -900,7 +900,7 @@ class MetricsService(
     // Only syncs pending measurements which can only be in metrics that are still running.
     val toBeSyncedInternalMeasurements: List<InternalMeasurement> =
       subResults
-        .filter { internalMetric -> determineMetricState(internalMetric) == Metric.State.RUNNING }
+        .filter { internalMetric -> internalMetric.state == Metric.State.RUNNING }
         .flatMap { internalMetric -> internalMetric.weightedMeasurementsList }
         .map { weightedMeasurement -> weightedMeasurement.measurement }
         .filter { internalMeasurement ->
@@ -979,7 +979,7 @@ class MetricsService(
         throw Exception("Unable to create the metric in the reporting database.", e)
       }
 
-    if (determineMetricState(internalMetric) == Metric.State.RUNNING) {
+    if (internalMetric.state == Metric.State.RUNNING) {
       measurementSupplier.createCmmsMeasurements(listOf(internalMetric), principal)
     }
 
@@ -1029,7 +1029,7 @@ class MetricsService(
             }
           )
           .metricsList
-          .filter { internalMetric -> determineMetricState(internalMetric) == Metric.State.RUNNING }
+          .filter { internalMetric -> internalMetric.state == Metric.State.RUNNING }
       } catch (e: StatusException) {
         throw Exception("Unable to create the metric in the reporting database.", e)
       }
@@ -1373,7 +1373,7 @@ private fun InternalMetric.toMetric(): Metric {
     timeInterval = source.timeInterval.toTimeInterval()
     metricSpec = source.metricSpec.toMetricSpec()
     filters += source.details.filtersList
-    state = determineMetricState(source)
+    state = source.state
     createTime = source.createTime
     if (state == Metric.State.SUCCEEDED) {
       result = buildMetricResult(source)
@@ -1381,18 +1381,6 @@ private fun InternalMetric.toMetric(): Metric {
   }
 }
 
-/** Determines the [Metric.State] based on the given [InternalMetric]. */
-private fun determineMetricState(metric: InternalMetric): Metric.State {
-  val measurementStates = metric.weightedMeasurementsList.map { it.measurement.state }
-  return if (measurementStates.all { it == InternalMeasurement.State.SUCCEEDED }) {
-    Metric.State.SUCCEEDED
-  } else if (measurementStates.any { it == InternalMeasurement.State.FAILED }) {
-    Metric.State.FAILED
-  } else {
-    Metric.State.RUNNING
-  }
-}
-
 /** Builds a [MetricResult] from the given [InternalMetric]. */
 private fun buildMetricResult(metric: InternalMetric): MetricResult {
   return metricResult {
@@ -1528,3 +1516,15 @@ private operator fun Duration.times(weight: Int): Duration {
 private operator fun Duration.plus(other: Duration): Duration {
   return Durations.add(this, other)
 }
+
+private val InternalMetric.state: Metric.State
+  get() {
+    val measurementStates = weightedMeasurementsList.map { it.measurement.state }
+    return if (measurementStates.all { it == InternalMeasurement.State.SUCCEEDED }) {
+      Metric.State.SUCCEEDED
+    } else if (measurementStates.any { it == InternalMeasurement.State.FAILED }) {
+      Metric.State.FAILED
+    } else {
+      Metric.State.RUNNING
+    }
+  }
diff --git a/src/main/proto/wfa/measurement/internal/reporting/v2/metric.proto b/src/main/proto/wfa/measurement/internal/reporting/v2/metric.proto
index eff7fee11ca..1cfa7954693 100644
--- a/src/main/proto/wfa/measurement/internal/reporting/v2/metric.proto
+++ b/src/main/proto/wfa/measurement/internal/reporting/v2/metric.proto
@@ -110,17 +110,6 @@ message Metric {
 
   MetricSpec metric_spec = 6;
 
-  enum State {
-    // Default value. This value is unused.
-    STATE_UNSPECIFIED = 0;
-    // Computation is running.
-    RUNNING = 1;
-    // Completed successfully. Terminal state.
-    SUCCEEDED = 2;
-    // Completed with failure. Terminal state.
-    FAILED = 3;
-  }
-
   message WeightedMeasurement {
     int32 weight = 1;
     Measurement measurement = 2;