Add context parameter to all container operations

6 files changed, 31 insertions, 23 deletions

diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/docker/ContainerOperations.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/docker/ContainerOperations.java
index c8fd4c077ce..255c04e0072 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/docker/ContainerOperations.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/docker/ContainerOperations.java
@@ -8,6 +8,7 @@ import com.yahoo.vespa.hosted.dockerapi.ContainerResources;
 import com.yahoo.vespa.hosted.dockerapi.ContainerStats;
 import com.yahoo.vespa.hosted.dockerapi.ProcessResult;
 import com.yahoo.vespa.hosted.dockerapi.RegistryCredentials;
+import com.yahoo.vespa.hosted.node.admin.component.TaskContext;
 import com.yahoo.vespa.hosted.node.admin.nodeagent.ContainerData;
 import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgentContext;
 import com.yahoo.vespa.hosted.node.admin.task.util.process.CommandResult;
@@ -20,6 +21,7 @@ import java.util.Set;
 /**
  * @author hakonhall
  */
+// TODO(mpolden): Clean up this interface when ContainerOperationsImpl, DockerEngine and friends can be removed
 public interface ContainerOperations {
 
     void createContainer(NodeAgentContext context, ContainerData containerData, ContainerResources containerResources);
@@ -32,7 +34,7 @@ public interface ContainerOperations {
 
     Optional<Container> getContainer(NodeAgentContext context);
 
-    boolean pullImageAsyncIfNeeded(DockerImage dockerImage, RegistryCredentials registryCredentials);
+    boolean pullImageAsyncIfNeeded(TaskContext context, DockerImage dockerImage, RegistryCredentials registryCredentials);
 
     ProcessResult executeCommandInContainerAsRoot(NodeAgentContext context, String... command);
 
@@ -59,10 +61,14 @@ public interface ContainerOperations {
 
     Optional<ContainerStats> getContainerStats(NodeAgentContext context);
 
-    boolean noManagedContainersRunning();
+    boolean noManagedContainersRunning(TaskContext context);
 
-    /** Stops and removes all managed containers except the ones given in {@code containerNames} */
-    boolean retainManagedContainers(Set<ContainerName> containerNames);
+    /**
+     * Stops and removes all managed containers except the ones given in {@code containerNames}.
+     *
+     * @return true if any containers were removed
+     */
+    boolean retainManagedContainers(TaskContext context, Set<ContainerName> containerNames);
 
     /** Deletes the local images that are currently not in use by any container and not recently used. */
     boolean deleteUnusedContainerImages(List<DockerImage> excludes, Duration minImageAgeToDelete);
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/docker/ContainerOperationsImpl.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/docker/ContainerOperationsImpl.java
index 40e1622fa0d..bcdeaf7db68 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/docker/ContainerOperationsImpl.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/docker/ContainerOperationsImpl.java
@@ -13,6 +13,7 @@ import com.yahoo.vespa.hosted.dockerapi.ContainerResources;
 import com.yahoo.vespa.hosted.dockerapi.ContainerStats;
 import com.yahoo.vespa.hosted.dockerapi.ProcessResult;
 import com.yahoo.vespa.hosted.dockerapi.RegistryCredentials;
+import com.yahoo.vespa.hosted.node.admin.component.TaskContext;
 import com.yahoo.vespa.hosted.node.admin.nodeadmin.ConvergenceException;
 import com.yahoo.vespa.hosted.node.admin.nodeagent.ContainerData;
 import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgentContext;
@@ -211,7 +212,7 @@ public class ContainerOperationsImpl implements ContainerOperations {
     }
 
     @Override
-    public boolean pullImageAsyncIfNeeded(DockerImage dockerImage, RegistryCredentials registryCredentials) {
+    public boolean pullImageAsyncIfNeeded(TaskContext context, DockerImage dockerImage, RegistryCredentials registryCredentials) {
         return containerEngine.pullImageAsyncIfNeeded(dockerImage, registryCredentials);
     }
 
@@ -321,12 +322,12 @@ public class ContainerOperationsImpl implements ContainerOperations {
     }
 
     @Override
-    public boolean noManagedContainersRunning() {
+    public boolean noManagedContainersRunning(TaskContext context) {
         return containerEngine.noManagedContainersRunning(MANAGER_NAME);
     }
 
     @Override
-    public boolean retainManagedContainers(Set<ContainerName> containerNames) {
+    public boolean retainManagedContainers(TaskContext context, Set<ContainerName> containerNames) {
         return containerEngine.listManagedContainers(MANAGER_NAME).stream()
                 .filter(containerName -> ! containerNames.contains(containerName))
                 .peek(containerName -> {
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentContextImpl.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentContextImpl.java
index f93cd005fae..5ea0a5d12c3 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentContextImpl.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentContextImpl.java
@@ -244,7 +244,7 @@ public class NodeAgentContextImpl implements NodeAgentContext {
             return this;
         }
 
-        public Builder dockerNetworking(ContainerNetworkMode containerNetworkMode) {
+        public Builder networkMode(ContainerNetworkMode containerNetworkMode) {
             this.containerNetworkMode = containerNetworkMode;
             return this;
         }
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImpl.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImpl.java
index 9fa21e5a676..dc0b6dc9d85 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImpl.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImpl.java
@@ -393,12 +393,13 @@ public class NodeAgentImpl implements NodeAgent {
         return zone.getEnvironment() == Environment.dev || zone.getSystemName().isCd();
     }
 
-    private boolean downloadImageIfNeeded(NodeSpec node, Optional<Container> container) {
+    private boolean downloadImageIfNeeded(NodeAgentContext context, Optional<Container> container) {
+        NodeSpec node = context.node();
         if (node.wantedDockerImage().equals(container.map(c -> c.image))) return false;
 
         RegistryCredentials credentials = registryCredentialsProvider.get();
         return node.wantedDockerImage()
-                   .map(image -> containerOperations.pullImageAsyncIfNeeded(image, credentials))
+                   .map(image -> containerOperations.pullImageAsyncIfNeeded(context, image, credentials))
                    .orElse(false);
     }
 
@@ -454,7 +455,7 @@ public class NodeAgentImpl implements NodeAgent {
                 storageMaintainer.cleanDiskIfFull(context);
                 storageMaintainer.handleCoreDumpsForContainer(context, container);
 
-                if (downloadImageIfNeeded(node, container)) {
+                if (downloadImageIfNeeded(context, container)) {
                     context.log(logger, "Waiting for image to download " + context.node().wantedDockerImage().get().asString());
                     return;
                 }
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/docker/ContainerOperationsImplTest.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/docker/ContainerOperationsImplTest.java
index 5f870664970..240fb492aff 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/docker/ContainerOperationsImplTest.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/docker/ContainerOperationsImplTest.java
@@ -130,7 +130,7 @@ public class ContainerOperationsImplTest {
     public void retainContainersTest() {
         when(containerEngine.listManagedContainers(ContainerOperationsImpl.MANAGER_NAME))
                 .thenReturn(List.of(new ContainerName("cnt1"), new ContainerName("cnt2"), new ContainerName("cnt3")));
-        containerOperations.retainManagedContainers(Set.of(new ContainerName("cnt2"), new ContainerName("cnt4")));
+        containerOperations.retainManagedContainers(any(), Set.of(new ContainerName("cnt2"), new ContainerName("cnt4")));
 
         verify(containerEngine).stopContainer(eq(new ContainerName("cnt1")));
         verify(containerEngine).deleteContainer(eq(new ContainerName("cnt1")));
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImplTest.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImplTest.java
index eea775f9a63..fcd5e8cc187 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImplTest.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImplTest.java
@@ -85,7 +85,7 @@ public class NodeAgentImplTest {
 
         verify(containerOperations, never()).removeContainer(eq(context), any());
         verify(orchestrator, never()).suspend(any(String.class));
-        verify(containerOperations, never()).pullImageAsyncIfNeeded(any(), any());
+        verify(containerOperations, never()).pullImageAsyncIfNeeded(any(), any(), any());
 
         final InOrder inOrder = inOrder(containerOperations, orchestrator, nodeRepository);
         // TODO: Verify this isn't run unless 1st time
@@ -155,7 +155,7 @@ public class NodeAgentImplTest {
         NodeAgentImpl nodeAgent = makeNodeAgent(null, false);
 
         when(nodeRepository.getOptionalNode(hostName)).thenReturn(Optional.of(node));
-        when(containerOperations.pullImageAsyncIfNeeded(eq(dockerImage), any())).thenReturn(false);
+        when(containerOperations.pullImageAsyncIfNeeded(any(), eq(dockerImage), any())).thenReturn(false);
 
         nodeAgent.doConverge(context);
 
@@ -164,7 +164,7 @@ public class NodeAgentImplTest {
         verify(orchestrator, never()).suspend(any(String.class));
 
         final InOrder inOrder = inOrder(containerOperations, orchestrator, nodeRepository, aclMaintainer, healthChecker);
-        inOrder.verify(containerOperations, times(1)).pullImageAsyncIfNeeded(eq(dockerImage), any());
+        inOrder.verify(containerOperations, times(1)).pullImageAsyncIfNeeded(any(), eq(dockerImage), any());
         inOrder.verify(containerOperations, times(1)).createContainer(eq(context), any(), any());
         inOrder.verify(containerOperations, times(1)).startContainer(eq(context));
         inOrder.verify(aclMaintainer, times(1)).converge(eq(context));
@@ -187,7 +187,7 @@ public class NodeAgentImplTest {
         NodeAgentImpl nodeAgent = makeNodeAgent(dockerImage, true);
 
         when(nodeRepository.getOptionalNode(hostName)).thenReturn(Optional.of(node));
-        when(containerOperations.pullImageAsyncIfNeeded(any(), any())).thenReturn(true);
+        when(containerOperations.pullImageAsyncIfNeeded(any(), any(), any())).thenReturn(true);
 
         nodeAgent.doConverge(context);
 
@@ -196,7 +196,7 @@ public class NodeAgentImplTest {
         verify(containerOperations, never()).removeContainer(eq(context), any());
 
         final InOrder inOrder = inOrder(containerOperations);
-        inOrder.verify(containerOperations, times(1)).pullImageAsyncIfNeeded(eq(newDockerImage), any());
+        inOrder.verify(containerOperations, times(1)).pullImageAsyncIfNeeded(any(), eq(newDockerImage), any());
     }
 
     @Test
@@ -209,7 +209,7 @@ public class NodeAgentImplTest {
         NodeAgentContext firstContext = createContext(specBuilder.build());
         NodeAgentImpl nodeAgent = makeNodeAgent(dockerImage, true);
 
-        when(containerOperations.pullImageAsyncIfNeeded(any(), any())).thenReturn(true);
+        when(containerOperations.pullImageAsyncIfNeeded(any(), any(), any())).thenReturn(true);
 
         InOrder inOrder = inOrder(orchestrator, containerOperations);
 
@@ -256,7 +256,7 @@ public class NodeAgentImplTest {
         NodeAgentContext firstContext = createContext(specBuilder.build());
         NodeAgentImpl nodeAgent = makeNodeAgent(dockerImage, true);
 
-        when(containerOperations.pullImageAsyncIfNeeded(any(), any())).thenReturn(true);
+        when(containerOperations.pullImageAsyncIfNeeded(any(), any(), any())).thenReturn(true);
 
         nodeAgent.doConverge(firstContext);
         NodeAgentContext secondContext = createContext(specBuilder.memoryGb(20).build());
@@ -316,7 +316,7 @@ public class NodeAgentImplTest {
         NodeAgentImpl nodeAgent = makeNodeAgent(dockerImage, true);
 
         when(nodeRepository.getOptionalNode(hostName)).thenReturn(Optional.of(node));
-        when(containerOperations.pullImageAsyncIfNeeded(eq(dockerImage), any())).thenReturn(false);
+        when(containerOperations.pullImageAsyncIfNeeded(any(), eq(dockerImage), any())).thenReturn(false);
         doThrow(new ConvergenceException("Connection refused")).doNothing()
                 .when(healthChecker).verifyHealth(eq(context));
 
@@ -543,7 +543,7 @@ public class NodeAgentImplTest {
         NodeAgentContext context = createContext(node);
         NodeAgentImpl nodeAgent = spy(makeNodeAgent(null, false));
 
-        when(containerOperations.pullImageAsyncIfNeeded(eq(dockerImage), any())).thenReturn(false);
+        when(containerOperations.pullImageAsyncIfNeeded(any(), eq(dockerImage), any())).thenReturn(false);
         doThrow(new DockerException("Failed to set up network")).doNothing().when(containerOperations).startContainer(eq(context));
 
         try {
@@ -580,7 +580,7 @@ public class NodeAgentImplTest {
         NodeAgentImpl nodeAgent = makeNodeAgent(null, false);
 
         when(nodeRepository.getOptionalNode(hostName)).thenReturn(Optional.of(node));
-        when(containerOperations.pullImageAsyncIfNeeded(eq(dockerImage), any())).thenReturn(false);
+        when(containerOperations.pullImageAsyncIfNeeded(any(), eq(dockerImage), any())).thenReturn(false);
 
         nodeAgent.doConverge(context);
 
@@ -588,7 +588,7 @@ public class NodeAgentImplTest {
         verify(orchestrator, never()).suspend(any(String.class));
 
         final InOrder inOrder = inOrder(containerOperations, orchestrator, nodeRepository, aclMaintainer);
-        inOrder.verify(containerOperations, times(1)).pullImageAsyncIfNeeded(eq(dockerImage), any());
+        inOrder.verify(containerOperations, times(1)).pullImageAsyncIfNeeded(any(), eq(dockerImage), any());
         inOrder.verify(containerOperations, times(1)).createContainer(eq(context), any(), any());
         inOrder.verify(containerOperations, times(1)).startContainer(eq(context));
         inOrder.verify(aclMaintainer, times(1)).converge(eq(context));