Turn metric consumer into a typed class

6 files changed, 75 insertions, 34 deletions

diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/maintenance/ConfigServerMaintainer.java b/configserver/src/main/java/com/yahoo/vespa/config/server/maintenance/ConfigServerMaintainer.java
index 4938f34131e..0a42ed8a384 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/maintenance/ConfigServerMaintainer.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/maintenance/ConfigServerMaintainer.java
@@ -35,14 +35,23 @@ public abstract class ConfigServerMaintainer extends Maintainer {
     ConfigServerMaintainer(ApplicationRepository applicationRepository, Curator curator, FlagSource flagSource,
                            Instant now, Duration interval) {
         super(null, interval, now, new JobControl(new JobControlFlags(curator, flagSource)),
-              jobMetrics(applicationRepository.metric()), cluster(curator), false);
+              new ConfigServerJobMetrics(applicationRepository.metric()), cluster(curator), false);
         this.applicationRepository = applicationRepository;
     }
 
-    private static JobMetrics jobMetrics(Metric metric) {
-        return new JobMetrics((job, consecutiveFailures) -> {
+    private static class ConfigServerJobMetrics extends JobMetrics {
+
+        private final Metric metric;
+
+        public ConfigServerJobMetrics(Metric metric) {
+            this.metric = metric;
+        }
+
+        @Override
+        protected void consume(String job, Long consecutiveFailures) {
             metric.set("maintenance.consecutiveFailures", consecutiveFailures, metric.createContext(Map.of("job", job)));
-        });
+        }
+
     }
 
     private static class JobControlFlags implements JobControlState {
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/ControllerMaintainer.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/ControllerMaintainer.java
index 03a6268397e..f8c0dcb1e14 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/ControllerMaintainer.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/ControllerMaintainer.java
@@ -34,7 +34,7 @@ public abstract class ControllerMaintainer extends Maintainer {
 
     public ControllerMaintainer(Controller controller, Duration interval, String name, Set<SystemName> activeSystems) {
         super(name, interval, controller.clock().instant(), controller.jobControl(),
-              jobMetrics(controller.metric()), controller.curator().cluster(), true);
+              new ControllerJobMetrics(controller.metric()), controller.curator().cluster(), true);
         this.controller = controller;
         this.activeSystems = Set.copyOf(Objects.requireNonNull(activeSystems));
     }
@@ -47,10 +47,19 @@ public abstract class ControllerMaintainer extends Maintainer {
         super.run();
     }
 
-    private static JobMetrics jobMetrics(Metric metric) {
-        return new JobMetrics((job, consecutiveFailures) -> {
-            metric.set("maintenance.consecutiveFailures", consecutiveFailures, metric.createContext(Map.of("job", job)));
-        });
+    private static class ControllerJobMetrics extends JobMetrics {
+
+        private final Metric metric;
+
+        public ControllerJobMetrics(Metric metric) {
+            this.metric = metric;
+        }
+
+        @Override
+        protected void consume(String job, Long incompleteRuns) {
+            metric.set("maintenance.consecutiveFailures", incompleteRuns, metric.createContext(Map.of("job", job)));
+        }
+
     }
 
 }
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/NodeRepositoryMaintainer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/NodeRepositoryMaintainer.java
index 0a1f6961f9f..3c0ed529ee9 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/NodeRepositoryMaintainer.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/NodeRepositoryMaintainer.java
@@ -25,7 +25,7 @@ public abstract class NodeRepositoryMaintainer extends Maintainer {
 
     public NodeRepositoryMaintainer(NodeRepository nodeRepository, Duration interval, Metric metric) {
         super(null, interval, nodeRepository.clock().instant(), nodeRepository.jobControl(),
-              jobMetrics(metric), nodeRepository.database().cluster(), true);
+              new NodeRepositoryJobMetrics(metric), nodeRepository.database().cluster(), true);
         this.nodeRepository = nodeRepository;
     }
 
@@ -48,10 +48,19 @@ public abstract class NodeRepositoryMaintainer extends Maintainer {
                                .groupingBy(node -> node.allocation().get().owner());
     }
 
-    private static JobMetrics jobMetrics(Metric metric) {
-        return new JobMetrics((job, consecutiveFailures) -> {
+    private static class NodeRepositoryJobMetrics extends JobMetrics {
+
+        private final Metric metric;
+
+        public NodeRepositoryJobMetrics(Metric metric) {
+            this.metric = metric;
+        }
+
+        @Override
+        protected void consume(String job, Long consecutiveFailures) {
             metric.set("maintenance.consecutiveFailures", consecutiveFailures, metric.createContext(Map.of("job", job)));
-        });
+        }
+
     }
 
 }
diff --git a/vespajlib/src/main/java/com/yahoo/concurrent/maintenance/JobMetrics.java b/vespajlib/src/main/java/com/yahoo/concurrent/maintenance/JobMetrics.java
index fcc5b8e57a2..bb1bd0085ef 100644
--- a/vespajlib/src/main/java/com/yahoo/concurrent/maintenance/JobMetrics.java
+++ b/vespajlib/src/main/java/com/yahoo/concurrent/maintenance/JobMetrics.java
@@ -9,16 +9,10 @@ import java.util.function.BiConsumer;
  *
  * @author mpolden
  */
-public class JobMetrics {
-
-    private final BiConsumer<String, Long> metricConsumer;
+public abstract class JobMetrics {
 
     private final ConcurrentHashMap<String, Long> incompleteRuns = new ConcurrentHashMap<>();
 
-    public JobMetrics(BiConsumer<String, Long> metricConsumer) {
-        this.metricConsumer = metricConsumer;
-    }
-
     /** Record a run for given job */
     public void recordRunOf(String job) {
         incompleteRuns.merge(job, 1L, Long::sum);
@@ -33,8 +27,10 @@ public class JobMetrics {
     public void forward(String job) {
         Long incompleteRuns = this.incompleteRuns.get(job);
         if (incompleteRuns != null) {
-            metricConsumer.accept(job, incompleteRuns);
+            consume(job, incompleteRuns);
         }
     }
 
+    protected abstract void consume(String job, Long incompleteRuns);
+
 }
diff --git a/vespajlib/src/test/java/com/yahoo/concurrent/maintenance/JobControlTest.java b/vespajlib/src/test/java/com/yahoo/concurrent/maintenance/JobControlTest.java
index 139a2901cd3..fe66e61a1ba 100644
--- a/vespajlib/src/test/java/com/yahoo/concurrent/maintenance/JobControlTest.java
+++ b/vespajlib/src/test/java/com/yahoo/concurrent/maintenance/JobControlTest.java
@@ -3,6 +3,8 @@ package com.yahoo.concurrent.maintenance;
 
 import org.junit.Test;
 
+import java.util.concurrent.atomic.AtomicLong;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
@@ -19,9 +21,8 @@ public class JobControlTest {
 
         String job1 = "Job1";
         String job2 = "Job2";
-        JobMetrics metrics = new JobMetrics((job, instant) -> {});
-        TestMaintainer maintainer1 = new TestMaintainer(job1, jobControl, metrics);
-        TestMaintainer maintainer2 = new TestMaintainer(job2, jobControl, metrics);
+        TestMaintainer maintainer1 = new TestMaintainer(job1, jobControl, new NoopJobMetrics());
+        TestMaintainer maintainer2 = new TestMaintainer(job2, jobControl, new NoopJobMetrics());
         assertEquals(2, jobControl.jobs().size());
         assertTrue(jobControl.jobs().contains(job1));
         assertTrue(jobControl.jobs().contains(job2));
@@ -62,7 +63,7 @@ public class JobControlTest {
     public void testJobControlMayDeactivateJobs() {
         JobControlStateMock state = new JobControlStateMock();
         JobControl jobControl = new JobControl(state);
-        TestMaintainer mockMaintainer = new TestMaintainer(null, jobControl, new JobMetrics((job, instant) -> {}));
+        TestMaintainer mockMaintainer = new TestMaintainer(null, jobControl, new NoopJobMetrics());
 
         assertTrue(jobControl.jobs().contains("TestMaintainer"));
 
@@ -80,4 +81,11 @@ public class JobControlTest {
         assertEquals(2, mockMaintainer.totalRuns());
     }
 
+    private static class NoopJobMetrics extends JobMetrics {
+
+        @Override
+        protected void consume(String job, Long incompleteRuns) { }
+
+    }
+
 }
diff --git a/vespajlib/src/test/java/com/yahoo/concurrent/maintenance/MaintainerTest.java b/vespajlib/src/test/java/com/yahoo/concurrent/maintenance/MaintainerTest.java
index e881d4b3ff6..871199c70d9 100644
--- a/vespajlib/src/test/java/com/yahoo/concurrent/maintenance/MaintainerTest.java
+++ b/vespajlib/src/test/java/com/yahoo/concurrent/maintenance/MaintainerTest.java
@@ -42,35 +42,45 @@ public class MaintainerTest {
 
     @Test
     public void success_metric() {
-        AtomicLong consecutiveFailures = new AtomicLong();
-        JobMetrics jobMetrics = new JobMetrics((job, count) -> consecutiveFailures.set(count));
+        TestJobMetrics jobMetrics = new TestJobMetrics();
         TestMaintainer maintainer = new TestMaintainer(null, jobControl, jobMetrics);
 
         // Maintainer fails twice in a row
         maintainer.successOnNextRun(false).run();
-        assertEquals(1, consecutiveFailures.get());
+        assertEquals(1, jobMetrics.consecutiveFailures.get());
         maintainer.successOnNextRun(false).run();
-        assertEquals(2, consecutiveFailures.get());
+        assertEquals(2, jobMetrics.consecutiveFailures.get());
 
         // Maintainer runs successfully
         maintainer.successOnNextRun(true).run();
-        assertEquals(0, consecutiveFailures.get());
+        assertEquals(0, jobMetrics.consecutiveFailures.get());
 
         // Maintainer runs successfully again
         maintainer.run();
-        assertEquals(0, consecutiveFailures.get());
+        assertEquals(0, jobMetrics.consecutiveFailures.get());
 
         // Maintainer throws
         maintainer.throwOnNextRun(new RuntimeException()).run();
-        assertEquals(1, consecutiveFailures.get());
+        assertEquals(1, jobMetrics.consecutiveFailures.get());
 
         // Maintainer recovers
         maintainer.throwOnNextRun(null).run();
-        assertEquals(0, consecutiveFailures.get());
+        assertEquals(0, jobMetrics.consecutiveFailures.get());
 
         // Lock exception is treated as a failure
         maintainer.throwOnNextRun(new UncheckedTimeoutException()).run();
-        assertEquals(1, consecutiveFailures.get());
+        assertEquals(1, jobMetrics.consecutiveFailures.get());
+    }
+
+    private static class TestJobMetrics extends JobMetrics {
+
+        AtomicLong consecutiveFailures = new AtomicLong();
+
+        @Override
+        protected void consume(String job, Long incompleteRuns) {
+            consecutiveFailures.set(incompleteRuns);
+        }
+
     }
 
 }