Listening to host changes to the service model

12 files changed, 323 insertions, 34 deletions

diff --git a/service-monitor/src/main/java/com/yahoo/vespa/service/duper/DuperModel.java b/service-monitor/src/main/java/com/yahoo/vespa/service/duper/DuperModel.java
index 1065ed29dc4..2df807e9510 100644
--- a/service-monitor/src/main/java/com/yahoo/vespa/service/duper/DuperModel.java
+++ b/service-monitor/src/main/java/com/yahoo/vespa/service/duper/DuperModel.java
@@ -34,6 +34,9 @@ public class DuperModel {
 
     public void registerListener(DuperModelListener listener) {
         applicationsById.values().forEach(listener::applicationActivated);
+        if (isComplete) {
+            listener.bootstrapComplete();
+        }
         listeners.add(listener);
     }
 
@@ -82,7 +85,7 @@ public class DuperModel {
         } else {
             logPrefix = isComplete ? "Reactivated application " : "Rebootstrapped application ";
         }
-        logger.log(LogLevel.INFO, logPrefix + id);
+        logger.log(LogLevel.INFO, logPrefix + id.toFullString());
 
         Set<HostName> hostnames = hostnamesById.computeIfAbsent(id, k -> new HashSet<>());
         Set<HostName> removedHosts = new HashSet<>(hostnames);
@@ -100,7 +103,7 @@ public class DuperModel {
                             // different application we will patch up our data structures to remain
                             // internally consistent. But listeners may be fooled.
                             logger.log(LogLevel.WARNING, hostname + " has been reassigned from " +
-                                    previousId + " to " + id);
+                                    previousId.toFullString() + " to " + id.toFullString());
 
                             Set<HostName> previousHostnames = hostnamesById.get(previousId);
                             if (previousHostnames != null) {
@@ -123,7 +126,7 @@ public class DuperModel {
 
         ApplicationInfo application = applicationsById.remove(applicationId);
         if (application != null) {
-            logger.log(LogLevel.INFO, "Removed application " + applicationId);
+            logger.log(LogLevel.INFO, "Removed application " + applicationId.toFullString());
             listeners.forEach(listener -> listener.applicationRemoved(applicationId));
         }
     }
diff --git a/service-monitor/src/main/java/com/yahoo/vespa/service/model/ApplicationInstanceGenerator.java b/service-monitor/src/main/java/com/yahoo/vespa/service/model/ApplicationInstanceGenerator.java
index 0116b992e23..d0ecad5f27a 100644
--- a/service-monitor/src/main/java/com/yahoo/vespa/service/model/ApplicationInstanceGenerator.java
+++ b/service-monitor/src/main/java/com/yahoo/vespa/service/model/ApplicationInstanceGenerator.java
@@ -58,8 +58,12 @@ public class ApplicationInstanceGenerator {
     }
 
     public ApplicationInstanceReference toApplicationInstanceReference() {
-        TenantId tenantId = new TenantId(applicationInfo.getApplicationId().tenant().toString());
-        ApplicationInstanceId applicationInstanceId = toApplicationInstanceId(applicationInfo.getApplicationId(), zone);
+        return toApplicationInstanceReference(applicationInfo.getApplicationId(), zone);
+    }
+
+    static ApplicationInstanceReference toApplicationInstanceReference(ApplicationId applicationId, Zone zone) {
+        TenantId tenantId = new TenantId(applicationId.tenant().toString());
+        ApplicationInstanceId applicationInstanceId = toApplicationInstanceId(applicationId, zone);
         return new ApplicationInstanceReference(tenantId, applicationInstanceId);
     }
 
diff --git a/service-monitor/src/main/java/com/yahoo/vespa/service/model/ModelGenerator.java b/service-monitor/src/main/java/com/yahoo/vespa/service/model/ModelGenerator.java
index d9378da957e..ad57e5e3565 100644
--- a/service-monitor/src/main/java/com/yahoo/vespa/service/model/ModelGenerator.java
+++ b/service-monitor/src/main/java/com/yahoo/vespa/service/model/ModelGenerator.java
@@ -2,6 +2,7 @@
 package com.yahoo.vespa.service.model;
 
 import com.yahoo.config.model.api.ApplicationInfo;
+import com.yahoo.config.provision.ApplicationId;
 import com.yahoo.config.provision.Zone;
 import com.yahoo.vespa.applicationmodel.ApplicationInstance;
 import com.yahoo.vespa.applicationmodel.ApplicationInstanceReference;
@@ -66,4 +67,8 @@ public class ModelGenerator {
         var generator = new ApplicationInstanceGenerator(applicationInfo, zone);
         return generator.makeApplicationInstanceLimitedTo(hostname, serviceStatusProvider);
     }
+
+    public ApplicationInstanceReference toApplicationInstanceReference(ApplicationId applicationId) {
+        return ApplicationInstanceGenerator.toApplicationInstanceReference(applicationId, zone);
+    }
 }
diff --git a/service-monitor/src/main/java/com/yahoo/vespa/service/model/ServiceHostListenerAdapter.java b/service-monitor/src/main/java/com/yahoo/vespa/service/model/ServiceHostListenerAdapter.java
new file mode 100644
index 00000000000..73e2573fc58
--- /dev/null
+++ b/service-monitor/src/main/java/com/yahoo/vespa/service/model/ServiceHostListenerAdapter.java
@@ -0,0 +1,73 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+package com.yahoo.vespa.service.model;
+
+import com.yahoo.config.model.api.ApplicationInfo;
+import com.yahoo.config.model.api.HostInfo;
+import com.yahoo.config.provision.ApplicationId;
+import com.yahoo.vespa.applicationmodel.ApplicationInstanceReference;
+import com.yahoo.vespa.applicationmodel.HostName;
+import com.yahoo.vespa.service.monitor.DuperModelListener;
+import com.yahoo.vespa.service.monitor.ServiceHostListener;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/**
+ * Allows a {@link ServiceHostListener} to register with the duper model as a {@link DuperModelListener}.
+ *
+ * <p>This class is not thread-safe: As with the DuperModelListener, events from the duper model
+ * happens within an exclusive duper model lock.</p>
+ */
+public class ServiceHostListenerAdapter implements DuperModelListener {
+    private final ServiceHostListener listener;
+    private final ModelGenerator modelGenerator;
+
+    private final Map<ApplicationInstanceReference, Set<HostName>> hostnamesByReference = new HashMap<>();
+
+    public static ServiceHostListenerAdapter asDuperModelListener(ServiceHostListener listener,
+                                                                  ModelGenerator generator) {
+        return new ServiceHostListenerAdapter(listener, generator);
+    }
+
+    ServiceHostListenerAdapter(ServiceHostListener listener, ModelGenerator generator) {
+        this.listener = listener;
+        this.modelGenerator = generator;
+    }
+
+    @Override
+    public void applicationActivated(ApplicationInfo application) {
+        Set<HostName> newHostnames = application.getModel().getHosts().stream()
+                .map(HostInfo::getHostname)
+                .map(HostName::new)
+                .collect(Collectors.toSet());
+
+        var reference = toApplicationInstanceReference(application.getApplicationId());
+
+        Set<HostName> oldHostnames = hostnamesByReference.get(reference);
+        if (!Objects.equals(newHostnames, oldHostnames)) {
+            hostnamesByReference.put(reference, newHostnames);
+            listener.onApplicationActivate(reference, newHostnames);
+        }
+    }
+
+    @Override
+    public void applicationRemoved(ApplicationId applicationId) {
+        var reference = toApplicationInstanceReference(applicationId);
+
+        if (hostnamesByReference.remove(reference) != null) {
+            listener.onApplicationRemove(reference);
+        }
+    }
+
+    @Override
+    public void bootstrapComplete() {
+    }
+
+    private ApplicationInstanceReference toApplicationInstanceReference(ApplicationId applicationId) {
+        return modelGenerator.toApplicationInstanceReference(applicationId);
+    }
+}
diff --git a/service-monitor/src/main/java/com/yahoo/vespa/service/model/ServiceModelCache.java b/service-monitor/src/main/java/com/yahoo/vespa/service/model/ServiceModelCache.java
index 29b2832efd0..cc0304e996a 100644
--- a/service-monitor/src/main/java/com/yahoo/vespa/service/model/ServiceModelCache.java
+++ b/service-monitor/src/main/java/com/yahoo/vespa/service/model/ServiceModelCache.java
@@ -3,11 +3,10 @@
 package com.yahoo.vespa.service.model;
 
 import com.yahoo.jdisc.Timer;
+import com.yahoo.vespa.service.monitor.ServiceHostListener;
 import com.yahoo.vespa.service.monitor.ServiceModel;
 import com.yahoo.vespa.service.monitor.ServiceMonitor;
 
-import java.util.function.Supplier;
-
 /**
  * Adds caching of a supplier of ServiceModel.
  *
@@ -16,7 +15,7 @@ import java.util.function.Supplier;
 public class ServiceModelCache implements ServiceMonitor {
     public static final long EXPIRY_MILLIS = 10000;
 
-    private final Supplier<ServiceModel> expensiveSupplier;
+    private final ServiceMonitor expensiveServiceMonitor;
     private final Timer timer;
 
     private volatile ServiceModel snapshot;
@@ -25,8 +24,8 @@ public class ServiceModelCache implements ServiceMonitor {
     private final Object updateMonitor = new Object();
     private long snapshotMillis;
 
-    public ServiceModelCache(Supplier<ServiceModel> expensiveSupplier, Timer timer) {
-        this.expensiveSupplier = expensiveSupplier;
+    public ServiceModelCache(ServiceMonitor expensiveServiceMonitor, Timer timer) {
+        this.expensiveServiceMonitor = expensiveServiceMonitor;
         this.timer = timer;
     }
 
@@ -59,8 +58,13 @@ public class ServiceModelCache implements ServiceMonitor {
         return snapshot;
     }
 
+    @Override
+    public void registerListener(ServiceHostListener listener) {
+        expensiveServiceMonitor.registerListener(listener);
+    }
+
     private void takeSnapshot() {
-        snapshot = expensiveSupplier.get();
+        snapshot = expensiveServiceMonitor.getServiceModelSnapshot();
         snapshotMillis = timer.currentTimeMillis();
     }
 
diff --git a/service-monitor/src/main/java/com/yahoo/vespa/service/model/ServiceModelProvider.java b/service-monitor/src/main/java/com/yahoo/vespa/service/model/ServiceModelProvider.java
index 45ee38ab560..d4a46bd59a9 100644
--- a/service-monitor/src/main/java/com/yahoo/vespa/service/model/ServiceModelProvider.java
+++ b/service-monitor/src/main/java/com/yahoo/vespa/service/model/ServiceModelProvider.java
@@ -4,11 +4,16 @@ package com.yahoo.vespa.service.model;
 import com.yahoo.config.model.api.ApplicationInfo;
 import com.yahoo.config.provision.ApplicationId;
 import com.yahoo.config.provision.Zone;
+import com.yahoo.jdisc.Metric;
+import com.yahoo.jdisc.Timer;
 import com.yahoo.vespa.applicationmodel.ApplicationInstance;
 import com.yahoo.vespa.applicationmodel.ApplicationInstanceReference;
 import com.yahoo.vespa.applicationmodel.HostName;
 import com.yahoo.vespa.applicationmodel.ServiceInstance;
 import com.yahoo.vespa.service.duper.DuperModelManager;
+import com.yahoo.vespa.service.manager.MonitorManager;
+import com.yahoo.vespa.service.manager.UnionMonitorManager;
+import com.yahoo.vespa.service.monitor.ServiceHostListener;
 import com.yahoo.vespa.service.monitor.ServiceModel;
 import com.yahoo.vespa.service.monitor.ServiceMonitor;
 import com.yahoo.vespa.service.monitor.ServiceStatusProvider;
@@ -27,19 +32,30 @@ public class ServiceModelProvider implements ServiceMonitor {
     private final ServiceMonitorMetrics metrics;
     private final DuperModelManager duperModelManager;
     private final ModelGenerator modelGenerator;
-    private final Zone zone;
     private final ServiceStatusProvider serviceStatusProvider;
 
-    public ServiceModelProvider(ServiceStatusProvider serviceStatusProvider,
-                                ServiceMonitorMetrics metrics,
-                                DuperModelManager duperModelManager,
-                                ModelGenerator modelGenerator,
+    public ServiceModelProvider(DuperModelManager duperModelManager,
+                                UnionMonitorManager monitorManager,
+                                Metric metric,
+                                Timer timer,
                                 Zone zone) {
-        this.serviceStatusProvider = serviceStatusProvider;
+        this(monitorManager,
+                new ServiceMonitorMetrics(metric, timer),
+                duperModelManager,
+                new ModelGenerator(zone)
+        );
+    }
+
+    ServiceModelProvider(MonitorManager monitorManager,
+                         ServiceMonitorMetrics metrics,
+                         DuperModelManager duperModelManager,
+                         ModelGenerator modelGenerator) {
+        this.serviceStatusProvider = monitorManager;
         this.metrics = metrics;
         this.duperModelManager = duperModelManager;
         this.modelGenerator = modelGenerator;
-        this.zone = zone;
+
+        duperModelManager.registerListener(monitorManager);
     }
 
     @Override
@@ -86,6 +102,12 @@ public class ServiceModelProvider implements ServiceMonitor {
         return getServiceModelSnapshot().getServiceInstancesByHostName();
     }
 
+    @Override
+    public void registerListener(ServiceHostListener listener) {
+        var duperModelListener = ServiceHostListenerAdapter.asDuperModelListener(listener, modelGenerator);
+        duperModelManager.registerListener(duperModelListener);
+    }
+
     private Optional<ApplicationInfo> getApplicationInfo(ApplicationInstanceReference reference) {
         ApplicationId applicationId = ApplicationInstanceGenerator.toApplicationId(reference);
         return duperModelManager.getApplicationInfo(applicationId);
diff --git a/service-monitor/src/main/java/com/yahoo/vespa/service/model/ServiceMonitorImpl.java b/service-monitor/src/main/java/com/yahoo/vespa/service/model/ServiceMonitorImpl.java
index d3297d711ff..7fd50e2f941 100644
--- a/service-monitor/src/main/java/com/yahoo/vespa/service/model/ServiceMonitorImpl.java
+++ b/service-monitor/src/main/java/com/yahoo/vespa/service/model/ServiceMonitorImpl.java
@@ -13,6 +13,7 @@ import com.yahoo.vespa.flags.FlagSource;
 import com.yahoo.vespa.flags.Flags;
 import com.yahoo.vespa.service.duper.DuperModelManager;
 import com.yahoo.vespa.service.manager.UnionMonitorManager;
+import com.yahoo.vespa.service.monitor.ServiceHostListener;
 import com.yahoo.vespa.service.monitor.ServiceModel;
 import com.yahoo.vespa.service.monitor.ServiceMonitor;
 
@@ -38,11 +39,10 @@ public class ServiceMonitorImpl implements ServiceMonitor {
                 monitorManager,
                 new ServiceMonitorMetrics(metric, timer),
                 duperModelManager,
-                new ModelGenerator(zone),
-                zone);
+                new ModelGenerator(zone));
 
         if (Flags.SERVICE_MODEL_CACHE.bindTo(flagSource).value()) {
-            delegate = new ServiceModelCache(serviceMonitor::getServiceModelSnapshot, timer);
+            delegate = new ServiceModelCache(serviceMonitor, timer);
         } else {
             delegate = serviceMonitor;
         }
@@ -77,4 +77,9 @@ public class ServiceMonitorImpl implements ServiceMonitor {
     public Map<HostName, List<ServiceInstance>> getServicesByHostname() {
         return delegate.getServicesByHostname();
     }
+
+    @Override
+    public void registerListener(ServiceHostListener listener) {
+        delegate.registerListener(listener);
+    }
 }
diff --git a/service-monitor/src/main/java/com/yahoo/vespa/service/monitor/ServiceHostListener.java b/service-monitor/src/main/java/com/yahoo/vespa/service/monitor/ServiceHostListener.java
new file mode 100644
index 00000000000..17d6fcf0d75
--- /dev/null
+++ b/service-monitor/src/main/java/com/yahoo/vespa/service/monitor/ServiceHostListener.java
@@ -0,0 +1,20 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+package com.yahoo.vespa.service.monitor;
+
+import com.yahoo.vespa.applicationmodel.ApplicationInstanceReference;
+import com.yahoo.vespa.applicationmodel.HostName;
+
+import java.util.Set;
+
+/**
+ * Interface for listening to changes to the set of applications, or the set of hosts
+ * assigned to each application, in the service model.
+ *
+ * <p>This is equivalent to listening to the duper model, since no health information leaks through
+ * from the service model, but the exposed types are those of the service model.</p>
+ */
+public interface ServiceHostListener {
+    void onApplicationActivate(ApplicationInstanceReference reference, Set<HostName> hostnames);
+    void onApplicationRemove(ApplicationInstanceReference reference);
+}
diff --git a/service-monitor/src/main/java/com/yahoo/vespa/service/monitor/ServiceMonitor.java b/service-monitor/src/main/java/com/yahoo/vespa/service/monitor/ServiceMonitor.java
index deafaaace0e..0d162fee182 100644
--- a/service-monitor/src/main/java/com/yahoo/vespa/service/monitor/ServiceMonitor.java
+++ b/service-monitor/src/main/java/com/yahoo/vespa/service/monitor/ServiceMonitor.java
@@ -47,4 +47,14 @@ public interface ServiceMonitor {
         return getServiceModelSnapshot().getServiceInstancesByHostName();
     }
 
+    /**
+     * Get notified of changes to the set of applications, or set of hosts assigned to an application.
+     *
+     * <p>When notified of model changes, the new model can be accessed through this interface
+     * by the listener. The model changes are visible to other threads strictly after the listener
+     * has been notified.</p>
+     *
+     * <p>WARNING: Methods on the listener may be invoked before returning from this method.</p>
+     */
+    default void registerListener(ServiceHostListener listener) { }
 }
diff --git a/service-monitor/src/test/java/com/yahoo/vespa/service/model/ServiceHostListenerAdapterTest.java b/service-monitor/src/test/java/com/yahoo/vespa/service/model/ServiceHostListenerAdapterTest.java
new file mode 100644
index 00000000000..6b1f0dfd5b6
--- /dev/null
+++ b/service-monitor/src/test/java/com/yahoo/vespa/service/model/ServiceHostListenerAdapterTest.java
@@ -0,0 +1,148 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.service.model;
+
+import com.yahoo.config.model.api.ApplicationInfo;
+import com.yahoo.config.model.api.HostInfo;
+import com.yahoo.config.model.api.Model;
+import com.yahoo.config.provision.ApplicationId;
+import com.yahoo.config.provision.CloudName;
+import com.yahoo.config.provision.Environment;
+import com.yahoo.config.provision.RegionName;
+import com.yahoo.config.provision.SystemName;
+import com.yahoo.config.provision.Zone;
+import com.yahoo.vespa.applicationmodel.ApplicationInstanceReference;
+import com.yahoo.vespa.applicationmodel.HostName;
+import com.yahoo.vespa.service.duper.DuperModel;
+import com.yahoo.vespa.service.monitor.ServiceHostListener;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import static org.junit.Assert.assertEquals;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+import static org.mockito.Mockito.when;
+
+public class ServiceHostListenerAdapterTest {
+    private final Zone zone = new Zone(CloudName.from("AWS"), SystemName.cd, Environment.dev, RegionName.from("us-east-1"));
+    private final ModelGenerator generator = new ModelGenerator(zone);
+    private final ServiceHostListener listener = mock(ServiceHostListener.class);
+    private final ServiceHostListenerAdapter adapter = new ServiceHostListenerAdapter(listener, generator);
+    private final DuperModel duperModel = new DuperModel();
+
+    @Before
+    public void setUp() {
+        duperModel.registerListener(adapter);
+    }
+
+    @Test
+    public void test() {
+        var applicationId1 = ApplicationId.from("tnt", "app", "default");
+        var applicationId2 = ApplicationId.from("tnt2", "app2", "default2");
+
+        verifyNoMoreInteractions(listener);
+
+        activate(applicationId1, "host1", "host2");
+        verifyActivate(applicationId1, "host1", "host2");
+
+        activate(applicationId1, "host1", "host3");
+        verifyActivate(applicationId1, "host1", "host3");
+
+        activate(applicationId1, "host1", "host3");
+        verifyNoActivate();
+
+        activate(applicationId2, "host4");
+        verifyActivate(applicationId2, "host4");
+
+        activate(applicationId1, "host1", "host3");
+        verifyNoActivate();
+
+        removeAndVerify(applicationId1, true);
+
+        activate(applicationId1, "host1", "host5");
+        verifyActivate(applicationId1, "host1", "host5");
+    }
+
+    @Test
+    public void documentDuplicateHostnameStrangeness() {
+        var applicationId1 = ApplicationId.from("tnt", "app", "default");
+        var applicationInfo1 = makeApplicationInfo(applicationId1, "host1", "host2");
+        duperModel.add(applicationInfo1);
+        verifyActivate(applicationId1, "host1", "host2");
+
+        var applicationId2 = ApplicationId.from("tnt2", "app2", "default2");
+        var applicationInfo2 = makeApplicationInfo(applicationId2, "host2", "host3");
+        duperModel.add(applicationInfo2);
+        verifyActivate(applicationId2, "host2", "host3");
+
+        // Duplicate hosts doesn't affect the ServiceHostListener.
+
+        duperModel.add(applicationInfo1);
+        verifyNoMoreInteractions(listener);
+
+        duperModel.add(applicationInfo2);
+        verifyNoMoreInteractions(listener);
+
+        // But do affect host lookup in duper model.
+
+        assertEquals(Optional.of(applicationInfo1), getDuperModelApplicationInfo("host1"));
+        assertEquals(Optional.of(applicationInfo2), getDuperModelApplicationInfo("host2")); // <--
+        assertEquals(Optional.of(applicationInfo2), getDuperModelApplicationInfo("host3"));
+    }
+
+    private Optional<ApplicationInfo> getDuperModelApplicationInfo(String hostname) {
+        return duperModel.getApplicationInfo(com.yahoo.config.provision.HostName.from(hostname));
+    }
+
+    private void removeAndVerify(ApplicationId id, boolean listenerInvoked) {
+        duperModel.remove(id);
+
+        if (listenerInvoked) {
+            ApplicationInstanceReference reference = generator.toApplicationInstanceReference(id);
+            verify(listener, times(1)).onApplicationRemove(reference);
+        }
+
+        verifyNoMoreInteractions(listener);
+    }
+
+    private void verifyActivate(ApplicationId id, String... hostnames) {
+        Set<HostName> hostnameSet = Stream.of(hostnames)
+                .map(HostName::new)
+                .collect(Collectors.toSet());
+
+        ApplicationInstanceReference reference = generator.toApplicationInstanceReference(id);
+
+        verify(listener, times(1)).onApplicationActivate(reference, hostnameSet);
+        verifyNoMoreInteractions(listener);
+    }
+
+    private void verifyNoActivate() {
+        verifyNoMoreInteractions(listener);
+    }
+
+    private void activate(ApplicationId id, String... hostnames) {
+        duperModel.add(makeApplicationInfo(id, hostnames));
+    }
+
+    private ApplicationInfo makeApplicationInfo(ApplicationId applicationId, String... hostnames) {
+        var applicationInfo = mock(ApplicationInfo.class);
+        when(applicationInfo.getApplicationId()).thenReturn(applicationId);
+
+        var model = mock(Model.class);
+        when(applicationInfo.getModel()).thenReturn(model);
+
+        List<HostInfo> hostnameList = Stream.of(hostnames)
+                .map(hostname -> new HostInfo(hostname, List.of()))
+                .collect(Collectors.toList());
+        when(model.getHosts()).thenReturn(hostnameList);
+
+        return applicationInfo;
+    }
+}
\ No newline at end of file
diff --git a/service-monitor/src/test/java/com/yahoo/vespa/service/model/ServiceModelCacheTest.java b/service-monitor/src/test/java/com/yahoo/vespa/service/model/ServiceModelCacheTest.java
index 21c6a49cc18..8649fec8e7e 100644
--- a/service-monitor/src/test/java/com/yahoo/vespa/service/model/ServiceModelCacheTest.java
+++ b/service-monitor/src/test/java/com/yahoo/vespa/service/model/ServiceModelCacheTest.java
@@ -4,10 +4,9 @@ package com.yahoo.vespa.service.model;
 
 import com.yahoo.jdisc.Timer;
 import com.yahoo.vespa.service.monitor.ServiceModel;
+import com.yahoo.vespa.service.monitor.ServiceMonitor;
 import org.junit.Test;
 
-import java.util.function.Supplier;
-
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
@@ -15,15 +14,14 @@ import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 public class ServiceModelCacheTest {
-    @SuppressWarnings("unchecked")
-    private final Supplier<ServiceModel> expensiveServiceMonitor = mock(Supplier.class);
+    private final ServiceMonitor expensiveServiceMonitor = mock(ServiceMonitor.class);
     private final Timer timer = mock(Timer.class);
     private final ServiceModelCache cache = new ServiceModelCache(expensiveServiceMonitor, timer);
 
     @Test
     public void sanityCheck() {
         ServiceModel serviceModel = mock(ServiceModel.class);
-        when(expensiveServiceMonitor.get()).thenReturn(serviceModel);
+        when(expensiveServiceMonitor.getServiceModelSnapshot()).thenReturn(serviceModel);
 
         long timeMillis = 0;
         when(timer.currentTimeMillis()).thenReturn(timeMillis);
@@ -31,7 +29,7 @@ public class ServiceModelCacheTest {
         // Will always populate cache the first time
         ServiceModel actualServiceModel = cache.getServiceModelSnapshot();
         assertTrue(actualServiceModel == serviceModel);
-        verify(expensiveServiceMonitor, times(1)).get();
+        verify(expensiveServiceMonitor, times(1)).getServiceModelSnapshot();
 
         // Cache hit
         timeMillis += ServiceModelCache.EXPIRY_MILLIS / 2;
@@ -44,12 +42,12 @@ public class ServiceModelCacheTest {
         when(timer.currentTimeMillis()).thenReturn(timeMillis);
 
         ServiceModel serviceModel2 = mock(ServiceModel.class);
-        when(expensiveServiceMonitor.get()).thenReturn(serviceModel2);
+        when(expensiveServiceMonitor.getServiceModelSnapshot()).thenReturn(serviceModel2);
 
         actualServiceModel = cache.getServiceModelSnapshot();
         assertTrue(actualServiceModel == serviceModel2);
         // '2' because it's cumulative with '1' from the first times(1).
-        verify(expensiveServiceMonitor, times(2)).get();
+        verify(expensiveServiceMonitor, times(2)).getServiceModelSnapshot();
 
         // Cache hit #2
         timeMillis += 1;
diff --git a/service-monitor/src/test/java/com/yahoo/vespa/service/model/ServiceModelProviderTest.java b/service-monitor/src/test/java/com/yahoo/vespa/service/model/ServiceModelProviderTest.java
index f34d61970ef..3b98e50bf9a 100644
--- a/service-monitor/src/test/java/com/yahoo/vespa/service/model/ServiceModelProviderTest.java
+++ b/service-monitor/src/test/java/com/yahoo/vespa/service/model/ServiceModelProviderTest.java
@@ -2,9 +2,8 @@
 package com.yahoo.vespa.service.model;
 
 import com.yahoo.config.model.api.ApplicationInfo;
-import com.yahoo.config.provision.Zone;
-import com.yahoo.vespa.service.monitor.ServiceModel;
 import com.yahoo.vespa.service.duper.DuperModelManager;
+import com.yahoo.vespa.service.monitor.ServiceModel;
 import com.yahoo.vespa.service.slobrok.SlobrokMonitorManagerImpl;
 import org.junit.Test;
 
@@ -23,13 +22,11 @@ public class ServiceModelProviderTest {
         SlobrokMonitorManagerImpl slobrokMonitorManager = mock(SlobrokMonitorManagerImpl.class);
         DuperModelManager duperModelManager = mock(DuperModelManager.class);
         ModelGenerator modelGenerator = mock(ModelGenerator.class);
-        Zone zone = mock(Zone.class);
         ServiceModelProvider provider = new ServiceModelProvider(
                 slobrokMonitorManager,
                 mock(ServiceMonitorMetrics.class),
                 duperModelManager,
-                modelGenerator,
-                zone);
+                modelGenerator);
 
         ApplicationInfo application1 = mock(ApplicationInfo.class);
         ApplicationInfo application2 = mock(ApplicationInfo.class);