Hide current context from NodeAgent

6 files changed, 41 insertions, 60 deletions

diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeadmin/NodeAdminImpl.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeadmin/NodeAdminImpl.java
index 5d021d329f4..5d2639d0a77 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeadmin/NodeAdminImpl.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeadmin/NodeAdminImpl.java
@@ -172,7 +172,7 @@ public class NodeAdminImpl implements NodeAdmin {
     @Override
     public void stop() {
         // Stop all node-agents in parallel, will block until the last NodeAgent is stopped
-        nodeAgentWithSchedulerByHostname.values().parallelStream().forEach(NodeAgent::stopForRemoval);
+        nodeAgentWithSchedulerByHostname.values().parallelStream().forEach(NodeAgentWithScheduler::stopForRemoval);
     }
 
     // Set-difference. Returns minuend minus subtrahend.
@@ -182,7 +182,7 @@ public class NodeAdminImpl implements NodeAdmin {
         return result;
     }
 
-    static class NodeAgentWithScheduler implements NodeAgent, NodeAgentScheduler {
+    static class NodeAgentWithScheduler implements NodeAgentScheduler {
         private final NodeAgent nodeAgent;
         private final NodeAgentScheduler nodeAgentScheduler;
 
@@ -191,14 +191,15 @@ public class NodeAdminImpl implements NodeAdmin {
             this.nodeAgentScheduler = nodeAgentScheduler;
         }
 
-        @Override public void start() { nodeAgent.start(); }
-        @Override public void stopForHostSuspension() { nodeAgent.stopForHostSuspension(); }
-        @Override public void stopForRemoval() { nodeAgent.stopForRemoval(); }
-        @Override public void updateContainerNodeMetrics() { nodeAgent.updateContainerNodeMetrics(); }
-        @Override public int getAndResetNumberOfUnhandledExceptions() { return nodeAgent.getAndResetNumberOfUnhandledExceptions(); }
+        void start() { nodeAgent.start(currentContext()); }
+        void stopForHostSuspension() { nodeAgent.stopForHostSuspension(currentContext()); }
+        void stopForRemoval() { nodeAgent.stopForRemoval(currentContext()); }
+        void updateContainerNodeMetrics() { nodeAgent.updateContainerNodeMetrics(currentContext()); }
+        int getAndResetNumberOfUnhandledExceptions() { return nodeAgent.getAndResetNumberOfUnhandledExceptions(); }
 
         @Override public void scheduleTickWith(NodeAgentContext context, Instant at) { nodeAgentScheduler.scheduleTickWith(context, at); }
         @Override public boolean setFrozen(boolean frozen, Duration timeout) { return nodeAgentScheduler.setFrozen(frozen, timeout); }
+        @Override public NodeAgentContext currentContext() { return nodeAgentScheduler.currentContext(); }
     }
 
     @FunctionalInterface
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgent.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgent.java
index de5ee1b69a4..f537884e708 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgent.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgent.java
@@ -13,24 +13,24 @@ public interface NodeAgent {
      * Starts the agent. After this method is called, the agent will asynchronously maintain the node, continuously
      * striving to make the current state equal to the wanted state.
      */
-    void start();
+    void start(NodeAgentContext context);
 
     /**
      * Stop the node in anticipation of host suspension, e.g. reboot or docker upgrade.
      */
-    void stopForHostSuspension();
+    void stopForHostSuspension(NodeAgentContext context);
 
     /**
      * Signals to the agent that the node is at the end of its lifecycle and no longer needs a managing agent.
      * Cleans up any resources the agent owns, such as threads, connections etc. Cleanup is synchronous; when this
      * method returns, no more actions will be taken by the agent.
      */
-    void stopForRemoval();
+    void stopForRemoval(NodeAgentContext context);
 
     /**
      * Updates metric receiver with the latest node-agent stats
      */
-    void updateContainerNodeMetrics();
+    default void updateContainerNodeMetrics(NodeAgentContext context) {}
 
     /**
      * Returns and resets number of unhandled exceptions
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentContextSupplier.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentContextSupplier.java
index 1fc730a3cb0..65611886f9c 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentContextSupplier.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentContextSupplier.java
@@ -13,9 +13,6 @@ public interface NodeAgentContextSupplier {
      */
     NodeAgentContext nextContext() throws InterruptedException;
 
-    /** @return the last context returned by {@link #nextContext()} or a default value */
-    NodeAgentContext currentContext();
-
     /** Interrupts the thread(s) currently waiting in {@link #nextContext()} */
     void interrupt();
 }
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 e9b48e89fd9..901a315ffb6 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
@@ -11,7 +11,6 @@ import com.yahoo.vespa.flags.FlagSource;
 import com.yahoo.vespa.flags.Flags;
 import com.yahoo.vespa.hosted.dockerapi.Container;
 import com.yahoo.vespa.hosted.dockerapi.ContainerResources;
-import com.yahoo.vespa.hosted.dockerapi.ContainerStats;
 import com.yahoo.vespa.hosted.dockerapi.exception.ContainerNotFoundException;
 import com.yahoo.vespa.hosted.dockerapi.exception.DockerException;
 import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeAttributes;
@@ -60,16 +59,16 @@ public class NodeAgentImpl implements NodeAgent {
     private final Optional<CredentialsMaintainer> credentialsMaintainer;
     private final Optional<AclMaintainer> aclMaintainer;
     private final Optional<HealthChecker> healthChecker;
-
     private final DoubleFlag containerCpuCap;
 
+    private Thread loopThread;
+    private ContainerState containerState = UNKNOWN;
+    private NodeSpec lastNode;
+
     private int numberOfUnhandledException = 0;
     private long currentRebootGeneration = 0;
     private Optional<Long> currentRestartGeneration = Optional.empty();
 
-    private final Thread loopThread;
-
-
     /**
      * ABSENT means container is definitely absent - A container that was absent will not suddenly appear without
      * NodeAgent explicitly starting it.
@@ -84,9 +83,6 @@ public class NodeAgentImpl implements NodeAgent {
         UNKNOWN
     }
 
-    private ContainerState containerState = UNKNOWN;
-
-    private NodeSpec lastNode = null;
 
     // Created in NodeAdminImpl
     public NodeAgentImpl(
@@ -107,11 +103,15 @@ public class NodeAgentImpl implements NodeAgent {
         this.credentialsMaintainer = credentialsMaintainer;
         this.aclMaintainer = aclMaintainer;
         this.healthChecker = healthChecker;
+        this.containerCpuCap = Flags.CONTAINER_CPU_CAP.bindTo(flagSource);
+    }
 
-        this.containerCpuCap = Flags.CONTAINER_CPU_CAP.bindTo(flagSource)
-                .with(FetchVector.Dimension.HOSTNAME, contextSupplier.currentContext().node().hostname());
+    @Override
+    public void start(NodeAgentContext initialContext) {
+        if (loopThread != null)
+            throw new IllegalStateException("Can not re-start a node agent.");
 
-        this.loopThread = new Thread(() -> {
+        loopThread = new Thread(() -> {
             while (!terminated.get()) {
                 try {
                     NodeAgentContext context = contextSupplier.nextContext();
@@ -119,19 +119,15 @@ public class NodeAgentImpl implements NodeAgent {
                 } catch (InterruptedException ignored) { }
             }
         });
-        this.loopThread.setName("tick-" + contextSupplier.currentContext().hostname());
-    }
-
-    @Override
-    public void start() {
+        loopThread.setName("tick-" + initialContext.hostname());
         loopThread.start();
     }
 
     @Override
-    public void stopForRemoval() {
-        if (!terminated.compareAndSet(false, true)) {
-            throw new RuntimeException("Can not re-stop a node agent.");
-        }
+    public void stopForRemoval(NodeAgentContext context) {
+        if (!terminated.compareAndSet(false, true))
+            throw new IllegalStateException("Can not re-stop a node agent.");
+
         contextSupplier.interrupt();
 
         do {
@@ -140,7 +136,7 @@ public class NodeAgentImpl implements NodeAgent {
             } catch (InterruptedException ignored) { }
         } while (loopThread.isAlive());
 
-        contextSupplier.currentContext().log(logger, "Stopped");
+        context.log(logger, "Stopped");
     }
 
     void startServicesIfNeeded(NodeAgentContext context) {
@@ -245,8 +241,7 @@ public class NodeAgentImpl implements NodeAgent {
         }
     }
 
-    private void stopServices() {
-        NodeAgentContext context = contextSupplier.currentContext();
+    private void stopServices(NodeAgentContext context) {
         context.log(logger, "Stopping services");
         if (containerState == ABSENT) return;
         try {
@@ -258,13 +253,11 @@ public class NodeAgentImpl implements NodeAgent {
     }
 
     @Override
-    public void stopForHostSuspension() {
-        NodeAgentContext context = contextSupplier.currentContext();
+    public void stopForHostSuspension(NodeAgentContext context) {
         getContainer(context).ifPresent(container -> removeContainer(context, container, "suspending host", true));
     }
 
-    public void suspend() {
-        NodeAgentContext context = contextSupplier.currentContext();
+    public void suspend(NodeAgentContext context) {
         context.log(logger, "Suspending services on node");
         if (containerState == ABSENT) return;
         try {
@@ -321,9 +314,9 @@ public class NodeAgentImpl implements NodeAgent {
 
             try {
                 if (context.node().state() != NodeState.dirty) {
-                    suspend();
+                    suspend(context);
                 }
-                stopServices();
+                stopServices(context);
             } catch (Exception e) {
                 context.log(logger, LogLevel.WARNING, "Failed stopping services, ignoring", e);
             }
@@ -355,6 +348,7 @@ public class NodeAgentImpl implements NodeAgent {
                         .map(NodeOwner::asApplicationId)
                         .map(appId -> containerCpuCap.with(FetchVector.Dimension.APPLICATION_ID, appId.serializedForm()))
                         .orElse(containerCpuCap)
+                        .with(FetchVector.Dimension.HOSTNAME, context.node().hostname())
                         .value() * context.node().vcpus();
 
         return ContainerResources.from(cpuCap, context.node().vcpus(), context.node().memoryGb());
@@ -503,14 +497,6 @@ public class NodeAgentImpl implements NodeAgent {
         }
     }
 
-    @Override
-    public void updateContainerNodeMetrics() {
-        NodeAgentContext context = contextSupplier.currentContext();
-        Optional<ContainerStats> containerStats = containerState != UNKNOWN ?
-                dockerOperations.getContainerStats(context) : Optional.empty();
-        updateContainerNodeMetrics(context, containerStats);
-    }
-
     private Optional<Container> getContainer(NodeAgentContext context) {
         if (containerState == ABSENT) return Optional.empty();
         Optional<Container> container = dockerOperations.getContainer(context);
@@ -570,8 +556,4 @@ public class NodeAgentImpl implements NodeAgent {
             }
         };
     }
-
-    protected void updateContainerNodeMetrics(NodeAgentContext context, Optional<ContainerStats> containerStats) {
-
-    }
 }
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentScheduler.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentScheduler.java
index a5daab8dcfd..956302dcdcc 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentScheduler.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentScheduler.java
@@ -19,4 +19,7 @@ public interface NodeAgentScheduler {
      * @return True if node agent has converged to the desired state
      */
     boolean setFrozen(boolean frozen, Duration timeout);
+
+    /** @return the last scheduled context or a default value */
+    NodeAgentContext currentContext();
 }
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 f26d0b57d72..f56669427ce 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
@@ -139,7 +139,7 @@ public class NodeAgentImplTest {
         inOrder.verify(dockerOperations, never()).startServices(eq(context));
         inOrder.verify(dockerOperations, times(1)).resumeNode(eq(context));
 
-        nodeAgent.stopForHostSuspension();
+        nodeAgent.stopForHostSuspension(context);
         nodeAgent.doConverge(context);
         inOrder.verify(dockerOperations, never()).startServices(eq(context));
         inOrder.verify(dockerOperations, times(1)).resumeNode(eq(context)); // Expect a resume, but no start services
@@ -149,7 +149,7 @@ public class NodeAgentImplTest {
         inOrder.verify(dockerOperations, never()).startServices(eq(context));
         inOrder.verify(dockerOperations, never()).resumeNode(eq(context));
 
-        nodeAgent.stopForHostSuspension();
+        nodeAgent.stopForHostSuspension(context);
         nodeAgent.doConverge(context);
         inOrder.verify(dockerOperations, times(1)).createContainer(eq(context), any(), any());
         inOrder.verify(dockerOperations, times(1)).startContainer(eq(context));
@@ -688,8 +688,6 @@ public class NodeAgentImplTest {
     }
     
     private NodeAgentContext createContext(NodeSpec nodeSpec) {
-        NodeAgentContext context = new NodeAgentContextImpl.Builder(nodeSpec).build();
-        when(contextSupplier.currentContext()).thenReturn(context);
-        return context;
+        return new NodeAgentContextImpl.Builder(nodeSpec).build();
     }
 }