From cfddb9ad578a7be4b08a92bb3e29632642aeac12 Mon Sep 17 00:00:00 2001
From: courtneyeh <courtney.hunter.nz@gmail.com>
Date: Wed, 25 Oct 2023 13:18:43 +1100
Subject: [PATCH] Clean up ValidatorDutyMetrics

---
 .../validator/client/duties/BlockProductionDuty.java |  9 +++------
 .../client/duties/ValidatorDutyMetrics.java          | 12 ++++++------
 2 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/validator/client/src/main/java/tech/pegasys/teku/validator/client/duties/BlockProductionDuty.java b/validator/client/src/main/java/tech/pegasys/teku/validator/client/duties/BlockProductionDuty.java
index 6377219e7d1..a88569a2102 100644
--- a/validator/client/src/main/java/tech/pegasys/teku/validator/client/duties/BlockProductionDuty.java
+++ b/validator/client/src/main/java/tech/pegasys/teku/validator/client/duties/BlockProductionDuty.java
@@ -77,20 +77,17 @@ public SafeFuture<DutyResult> performDuty() {
   }
 
   public SafeFuture<DutyResult> produceBlock(final ForkInfo forkInfo) {
-    final String dutyType = getType().getType();
     return createRandaoReveal(forkInfo)
         .thenCompose(
             signature ->
-                validatorDutyMetrics.record(
-                    () -> createUnsignedBlock(signature), dutyType, "create"))
+                validatorDutyMetrics.record(() -> createUnsignedBlock(signature), this, "create"))
         .thenCompose(
             unsignedBlock ->
                 validatorDutyMetrics.record(
-                    () -> signBlockContainer(forkInfo, unsignedBlock), dutyType, "sign"))
+                    () -> signBlockContainer(forkInfo, unsignedBlock), this, "sign"))
         .thenCompose(
             signedBlockContainer ->
-                validatorDutyMetrics.record(
-                    () -> sendBlock(signedBlockContainer), dutyType, "send"))
+                validatorDutyMetrics.record(() -> sendBlock(signedBlockContainer), this, "send"))
         .exceptionally(error -> DutyResult.forError(validator.getPublicKey(), error));
   }
 
diff --git a/validator/client/src/main/java/tech/pegasys/teku/validator/client/duties/ValidatorDutyMetrics.java b/validator/client/src/main/java/tech/pegasys/teku/validator/client/duties/ValidatorDutyMetrics.java
index 6f5f3049e17..49935441a34 100644
--- a/validator/client/src/main/java/tech/pegasys/teku/validator/client/duties/ValidatorDutyMetrics.java
+++ b/validator/client/src/main/java/tech/pegasys/teku/validator/client/duties/ValidatorDutyMetrics.java
@@ -43,6 +43,12 @@ public SafeFuture<DutyResult> performDutyWithMetrics(final Duty duty) {
     return duty.performDuty().alwaysRun(context::stopTimer);
   }
 
+  public <T> SafeFuture<T> record(
+          final Supplier<SafeFuture<T>> dutyStepFutureSupplier, final Duty duty, final String step) {
+    final OperationTimer.TimingContext context = startTimer(getDutyType(duty), step);
+    return dutyStepFutureSupplier.get().alwaysRun(context::stopTimer);
+  }
+
   private OperationTimer.TimingContext startTimer(final String dutyType, final String step) {
     final OperationTimer timer = dutyMetric.labels(dutyType, step);
     return timer.startTimer();
@@ -51,10 +57,4 @@ private OperationTimer.TimingContext startTimer(final String dutyType, final Str
   private static String getDutyType(final Duty duty) {
     return duty.getType().getType();
   }
-
-  public <T> SafeFuture<T> record(
-      final Supplier<SafeFuture<T>> supplier, final String dutyType, final String step) {
-    final OperationTimer.TimingContext context = startTimer(dutyType, step);
-    return supplier.get().alwaysRun(context::stopTimer);
-  }
 }