Set CPUs from feature flag

4 files changed, 60 insertions, 6 deletions

diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/NodeSpec.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/NodeSpec.java
index 071d4b549de..c00caef5587 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/NodeSpec.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/configserver/noderepository/NodeSpec.java
@@ -1,6 +1,7 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.node.admin.configserver.noderepository;
 
+import com.yahoo.config.provision.ApplicationId;
 import com.yahoo.config.provision.NodeType;
 import com.yahoo.vespa.hosted.dockerapi.DockerImage;
 import com.yahoo.vespa.hosted.provision.Node;
@@ -350,6 +351,10 @@ public class NodeSpec {
             return instance;
         }
 
+        public ApplicationId asApplicationId() {
+            return ApplicationId.from(tenant, application, instance);
+        }
+
         @Override
         public boolean equals(Object o) {
             if (this == o) return true;
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 96c6fed63de..7b01f9e3379 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
@@ -4,6 +4,10 @@ package com.yahoo.vespa.hosted.node.admin.nodeagent;
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.yahoo.concurrent.ThreadFactoryFactory;
 import com.yahoo.log.LogLevel;
+import com.yahoo.vespa.flags.DoubleFlag;
+import com.yahoo.vespa.flags.FetchVector;
+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;
@@ -69,6 +73,8 @@ public class NodeAgentImpl implements NodeAgent {
     private final Optional<AclMaintainer> aclMaintainer;
     private final Optional<HealthChecker> healthChecker;
 
+    private final DoubleFlag containerCpuCap;
+
     private int numberOfUnhandledException = 0;
     private DockerImage imageBeingDownloaded = null;
 
@@ -108,6 +114,7 @@ public class NodeAgentImpl implements NodeAgent {
             final Orchestrator orchestrator,
             final DockerOperations dockerOperations,
             final StorageMaintainer storageMaintainer,
+            final FlagSource flagSource,
             final Optional<AthenzCredentialsMaintainer> athenzCredentialsMaintainer,
             final Optional<AclMaintainer> aclMaintainer,
             final Optional<HealthChecker> healthChecker) {
@@ -120,6 +127,9 @@ public class NodeAgentImpl implements NodeAgent {
         this.aclMaintainer = aclMaintainer;
         this.healthChecker = healthChecker;
 
+        this.containerCpuCap = Flags.CONTAINER_CPU_CAP.bindTo(flagSource)
+                .with(FetchVector.Dimension.HOSTNAME, contextSupplier.currentContext().node().getHostname());
+
         this.loopThread = new Thread(() -> {
             while (!terminated.get()) {
                 try {
@@ -350,8 +360,14 @@ public class NodeAgentImpl implements NodeAgent {
     }
 
     private void updateContainerIfNeeded(NodeAgentContext context, Container existingContainer) {
+        double cpuCap = context.node().getOwner()
+                .map(NodeSpec.Owner::asApplicationId)
+                .map(appId -> containerCpuCap.with(FetchVector.Dimension.APPLICATION_ID, appId.serializedForm()))
+                .orElse(containerCpuCap)
+                .value();
+
         ContainerResources wantedContainerResources = ContainerResources.from(
-                0, context.node().getMinCpuCores(), context.node().getMinMainMemoryAvailableGb());
+                cpuCap, context.node().getMinCpuCores(), context.node().getMinMainMemoryAvailableGb());
 
         if (wantedContainerResources.equals(existingContainer.resources)) return;
         context.log(logger, "Container should be running with different resource allocation, wanted: %s, current: %s",
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/DockerTester.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/DockerTester.java
index 109bce4c13f..7ac74b190ad 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/DockerTester.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integrationTests/DockerTester.java
@@ -6,6 +6,8 @@ import com.yahoo.config.provision.HostName;
 import com.yahoo.config.provision.NodeType;
 import com.yahoo.metrics.simple.MetricReceiver;
 import com.yahoo.system.ProcessExecuter;
+import com.yahoo.vespa.flags.Flags;
+import com.yahoo.vespa.flags.InMemoryFlagSource;
 import com.yahoo.vespa.hosted.dockerapi.Docker;
 import com.yahoo.vespa.hosted.dockerapi.metrics.MetricReceiverWrapper;
 import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeSpec;
@@ -58,6 +60,7 @@ public class DockerTester implements AutoCloseable {
     final Orchestrator orchestrator = mock(Orchestrator.class);
     final StorageMaintainer storageMaintainer = mock(StorageMaintainer.class);
     final InOrder inOrder = Mockito.inOrder(docker, nodeRepository, orchestrator, storageMaintainer);
+    final InMemoryFlagSource flagSource = new InMemoryFlagSource().withFlag(Flags.CONTAINER_CPU_CAP.id());
 
     final NodeAdminStateUpdater nodeAdminStateUpdater;
     final NodeAdminImpl nodeAdmin;
@@ -92,8 +95,8 @@ public class DockerTester implements AutoCloseable {
 
         MetricReceiverWrapper mr = new MetricReceiverWrapper(MetricReceiver.nullImplementation);
         NodeAgentFactory nodeAgentFactory = contextSupplier -> new NodeAgentImpl(
-                contextSupplier, nodeRepository,
-                orchestrator, dockerOperations, storageMaintainer, Optional.empty(), Optional.empty(), Optional.empty());
+                contextSupplier, nodeRepository, orchestrator, dockerOperations, storageMaintainer, flagSource,
+                Optional.empty(), Optional.empty(), Optional.empty());
         NodeAgentContextFactory nodeAgentContextFactory = nodeSpec ->
                 new NodeAgentContextImpl.Builder(nodeSpec).fileSystem(fileSystem).build();
         nodeAdmin = new NodeAdminImpl(nodeAgentFactory, nodeAgentContextFactory, Optional.empty(), mr, Clock.systemUTC());
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 59a40fe3ccc..78fa53edb56 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
@@ -4,6 +4,11 @@ package com.yahoo.vespa.hosted.node.admin.nodeagent;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.yahoo.config.provision.NodeType;
 import com.yahoo.metrics.simple.MetricReceiver;
+import com.yahoo.vespa.flags.FetchVector;
+import com.yahoo.vespa.flags.Flags;
+import com.yahoo.vespa.flags.InMemoryFlagSource;
+import com.yahoo.vespa.flags.JsonNodeRawFlag;
+import com.yahoo.vespa.flags.json.Rule;
 import com.yahoo.vespa.hosted.dockerapi.Container;
 import com.yahoo.vespa.hosted.dockerapi.ContainerName;
 import com.yahoo.vespa.hosted.dockerapi.ContainerResources;
@@ -80,6 +85,8 @@ public class NodeAgentImplTest {
     private final ContainerStats emptyContainerStats = new ContainerStats(Collections.emptyMap(),
             Collections.emptyMap(), Collections.emptyMap(), Collections.emptyMap());
     private final AthenzCredentialsMaintainer athenzCredentialsMaintainer = mock(AthenzCredentialsMaintainer.class);
+    private final InMemoryFlagSource flagSource = new InMemoryFlagSource()
+            .withFlag(Flags.CONTAINER_CPU_CAP.id());
 
 
     @Test
@@ -252,15 +259,34 @@ public class NodeAgentImplTest {
         nodeAgent.doConverge(secondContext);
         NodeAgentContext thirdContext = createContext(specBuilder.minCpuCores(4).build());
         nodeAgent.doConverge(thirdContext);
+        ContainerResources resourcesAfterThird = ContainerResources.from(0, 4, 16);
+        mockGetContainer(dockerImage, resourcesAfterThird, true);
 
         InOrder inOrder = inOrder(orchestrator, dockerOperations);
         inOrder.verify(orchestrator).resume(any(String.class));
         inOrder.verify(orchestrator).resume(any(String.class));
         inOrder.verify(orchestrator).suspend(any(String.class));
-        inOrder.verify(dockerOperations).updateContainer(eq(thirdContext), any());
+        inOrder.verify(dockerOperations).updateContainer(eq(thirdContext), eq(resourcesAfterThird));
         inOrder.verify(dockerOperations, never()).removeContainer(any(), any());
         inOrder.verify(dockerOperations, never()).startContainer(any());
         inOrder.verify(orchestrator).resume(any(String.class));
+
+        // No changes
+        nodeAgent.converge(thirdContext);
+        inOrder.verify(orchestrator, never()).suspend(any(String.class));
+        inOrder.verify(dockerOperations, never()).updateContainer(eq(thirdContext), any());
+        inOrder.verify(dockerOperations, never()).removeContainer(any(), any());
+        inOrder.verify(orchestrator).resume(any(String.class));
+
+        // Set the feature flag
+        flagSource.withFlag(
+                Flags.CONTAINER_CPU_CAP.id(),
+                new FetchVector(),
+                new Rule(Optional.of(JsonNodeRawFlag.fromJson("2.3"))));
+
+        nodeAgent.converge(thirdContext);
+        inOrder.verify(dockerOperations).updateContainer(eq(thirdContext), eq(ContainerResources.from(2.3, 4, 16)));
+        inOrder.verify(orchestrator).resume(any(String.class));
     }
 
     @Test
@@ -702,11 +728,15 @@ public class NodeAgentImplTest {
         doNothing().when(storageMaintainer).writeMetricsConfig(any());
 
         return new NodeAgentImpl(contextSupplier, nodeRepository, orchestrator, dockerOperations,
-                storageMaintainer, Optional.of(athenzCredentialsMaintainer), Optional.of(aclMaintainer),
+                storageMaintainer, flagSource, Optional.of(athenzCredentialsMaintainer), Optional.of(aclMaintainer),
                 Optional.of(healthChecker));
     }
 
     private void mockGetContainer(DockerImage dockerImage, boolean isRunning) {
+        mockGetContainer(dockerImage, ContainerResources.from(0, MIN_CPU_CORES, MIN_MAIN_MEMORY_AVAILABLE_GB), isRunning);
+    }
+
+    private void mockGetContainer(DockerImage dockerImage, ContainerResources containerResources, boolean isRunning) {
         doAnswer(invoc -> {
             NodeAgentContext context = invoc.getArgument(0);
             if (!hostName.equals(context.hostname().value()))
@@ -715,7 +745,7 @@ public class NodeAgentImplTest {
                     Optional.of(new Container(
                             hostName,
                             dockerImage,
-                            ContainerResources.from(0, MIN_CPU_CORES, MIN_MAIN_MEMORY_AVAILABLE_GB),
+                            containerResources,
                             ContainerName.fromHostname(hostName),
                             isRunning ? Container.State.RUNNING : Container.State.EXITED,
                             isRunning ? 1 : 0)) :