Add support for setting cpu quota with ContainerResources

5 files changed, 88 insertions, 46 deletions

diff --git a/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/ContainerResources.java b/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/ContainerResources.java
index 346223d0e7e..1a643eafacf 100644
--- a/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/ContainerResources.java
+++ b/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/ContainerResources.java
@@ -1,26 +1,42 @@
 // Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.dockerapi;
 
+import java.util.Objects;
+
 /**
  * @author valerijf
  */
 public class ContainerResources {
-    public static final ContainerResources UNLIMITED = ContainerResources.from(0, 0);
+    public static final ContainerResources UNLIMITED = ContainerResources.from(0, 0, 0);
 
+    private final double cpus;
     private final int cpuShares;
     private final long memoryBytes;
 
-    ContainerResources(int cpuShares, long memoryBytes) {
+    ContainerResources(double cpus, int cpuShares, long memoryBytes) {
+        this.cpus = cpus;
         this.cpuShares = cpuShares;
         this.memoryBytes = memoryBytes;
+
+        if (cpus < 0)
+            throw new IllegalArgumentException("CPUs must be a positive number or 0 for unlimited, was " + cpus);
+        if (cpuShares < 0)
+            throw new IllegalArgumentException("CPU shares must be a positive integer or 0 for unlimited, was " + cpuShares);
+        if (memoryBytes < 0)
+            throw new IllegalArgumentException("memoryBytes must be a positive integer or 0 for unlimited, was " + memoryBytes);
     }
 
-    public static ContainerResources from(double cpuCores, double memoryGb) {
+    public static ContainerResources from(double cpus, double cpuCores, double memoryGb) {
         return new ContainerResources(
+                cpus,
                 (int) Math.round(10 * cpuCores),
                 (long) ((1L << 30) * memoryGb));
     }
 
+    public double cpus() {
+        return cpus;
+    }
+
     public int cpuShares() {
         return cpuShares;
     }
@@ -33,22 +49,21 @@ public class ContainerResources {
     public boolean equals(Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
-
         ContainerResources that = (ContainerResources) o;
-
-        if (cpuShares != that.cpuShares) return false;
-        return memoryBytes == that.memoryBytes;
+        return Double.compare(that.cpus, cpus) == 0 &&
+                cpuShares == that.cpuShares &&
+                memoryBytes == that.memoryBytes;
     }
 
     @Override
     public int hashCode() {
-        int result = cpuShares;
-        result = 31 * result + (int) (memoryBytes ^ (memoryBytes >>> 32));
-        return result;
+        return Objects.hash(cpus, cpuShares, memoryBytes);
     }
 
     @Override
     public String toString() {
-        return cpuShares + " CPU Shares, " + memoryBytes + "B memory";
+        return (cpus > 0 ? cpus : "unlimited") +" CPUs, " +
+                (cpuShares > 0 ? cpuShares : "unlimited") + " CPU Shares, " +
+                (memoryBytes > 0 ? memoryBytes + "B" : "unlimited") + " memory";
     }
 }
diff --git a/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/CreateContainerCommandImpl.java b/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/CreateContainerCommandImpl.java
index 04aad562711..36b71b66158 100644
--- a/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/CreateContainerCommandImpl.java
+++ b/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/CreateContainerCommandImpl.java
@@ -5,6 +5,7 @@ import com.github.dockerjava.api.DockerClient;
 import com.github.dockerjava.api.command.CreateContainerCmd;
 import com.github.dockerjava.api.model.Bind;
 import com.github.dockerjava.api.model.Capability;
+import com.github.dockerjava.api.model.HostConfig;
 import com.github.dockerjava.api.model.Ulimit;
 import com.yahoo.vespa.hosted.dockerapi.exception.DockerException;
 
@@ -25,13 +26,14 @@ import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
 
+import static com.yahoo.vespa.hosted.dockerapi.DockerImpl.CPU_PERIOD;
+import static com.yahoo.vespa.hosted.dockerapi.DockerImpl.LABEL_NAME_MANAGEDBY;
+
 class CreateContainerCommandImpl implements Docker.CreateContainerCommand {
 
     private final DockerClient docker;
     private final DockerImage dockerImage;
-    private final ContainerResources containerResources;
     private final ContainerName containerName;
-    private final String hostName;
     private final Map<String, String> labels = new HashMap<>();
     private final List<String> environmentAssignments = new ArrayList<>();
     private final List<String> volumeBindSpecs = new ArrayList<>();
@@ -39,22 +41,31 @@ class CreateContainerCommandImpl implements Docker.CreateContainerCommand {
     private final Set<Capability> addCapabilities = new HashSet<>();
     private final Set<Capability> dropCapabilities = new HashSet<>();
 
+    private Optional<String> hostName = Optional.empty();
+    private Optional<ContainerResources> containerResources = Optional.empty();
     private Optional<String> networkMode = Optional.empty();
     private Optional<String> ipv4Address = Optional.empty();
     private Optional<String> ipv6Address = Optional.empty();
     private Optional<String[]> entrypoint = Optional.empty();
     private boolean privileged = false;
 
-    CreateContainerCommandImpl(DockerClient docker,
-                               DockerImage dockerImage,
-                               ContainerResources containerResources,
-                               ContainerName containerName,
-                               String hostName) {
+    CreateContainerCommandImpl(DockerClient docker, DockerImage dockerImage, ContainerName containerName) {
         this.docker = docker;
         this.dockerImage = dockerImage;
-        this.containerResources = containerResources;
         this.containerName = containerName;
-        this.hostName = hostName;
+    }
+
+
+    @Override
+    public Docker.CreateContainerCommand withHostName(String hostName) {
+        this.hostName = Optional.of(hostName);
+        return this;
+    }
+
+    @Override
+    public Docker.CreateContainerCommand withResources(ContainerResources containerResources) {
+        this.containerResources = Optional.of(containerResources);
+        return this;
     }
 
     @Override
@@ -65,7 +76,7 @@ class CreateContainerCommandImpl implements Docker.CreateContainerCommand {
     }
 
     public Docker.CreateContainerCommand withManagedBy(String manager) {
-        return withLabel(DockerImpl.LABEL_NAME_MANAGEDBY, manager);
+        return withLabel(LABEL_NAME_MANAGEDBY, manager);
     }
 
     @Override
@@ -140,19 +151,25 @@ class CreateContainerCommandImpl implements Docker.CreateContainerCommand {
         try {
             createCreateContainerCmd().exec();
         } catch (RuntimeException e) {
-            throw new DockerException("Failed to create container " + containerName.asString(), e);
+            throw new DockerException("Failed to create container " + toString(), e);
         }
     }
 
     private CreateContainerCmd createCreateContainerCmd() {
         List<Bind> volumeBinds = volumeBindSpecs.stream().map(Bind::parse).collect(Collectors.toList());
 
+        final HostConfig hostConfig = new HostConfig();
+
+        containerResources.ifPresent(cr -> hostConfig
+                .withCpuShares(cr.cpuShares())
+                .withMemory(cr.memoryBytes())
+                .withCpuPeriod(cr.cpus() > 0 ? CPU_PERIOD : null)
+                .withCpuQuota(cr.cpus() > 0 ? (int) (CPU_PERIOD * cr.cpus()) : null));
+
         final CreateContainerCmd containerCmd = docker
                 .createContainerCmd(dockerImage.asString())
-                .withCpuShares(containerResources.cpuShares())
-                .withMemory(containerResources.memoryBytes())
+                .withHostConfig(hostConfig) // MUST BE FIRST (some of the later setters are simply proxied to HostConfig)
                 .withName(containerName.asString())
-                .withHostName(hostName)
                 .withLabels(labels)
                 .withEnv(environmentAssignments)
                 .withBinds(volumeBinds)
@@ -165,6 +182,7 @@ class CreateContainerCommandImpl implements Docker.CreateContainerCommand {
                 .filter(mode -> ! mode.toLowerCase().equals("host"))
                 .ifPresent(mode -> containerCmd.withMacAddress(generateMACAddress(hostName, ipv4Address, ipv6Address)));
 
+        hostName.ifPresent(containerCmd::withHostName);
         networkMode.ifPresent(containerCmd::withNetworkMode);
         ipv4Address.ifPresent(containerCmd::withIpv4Address);
         ipv6Address.ifPresent(containerCmd::withIpv6Address);
@@ -174,17 +192,17 @@ class CreateContainerCommandImpl implements Docker.CreateContainerCommand {
     }
 
     /** Maps ("--env", {"A", "B", "C"}) to "--env A --env B --env C" */
-    private String toRepeatedOption(String option, List<String> optionValues) {
+    private static String toRepeatedOption(String option, List<String> optionValues) {
         return optionValues.stream()
                 .map(optionValue -> option + " " + optionValue)
                 .collect(Collectors.joining(" "));
     }
 
-    private String toOptionalOption(String option, Optional<String> value) {
+    private static String toOptionalOption(String option, Optional<?> value) {
         return value.map(o -> option + " " + o).orElse("");
     }
 
-    private String toFlagOption(String option, boolean value) {
+    private static String toFlagOption(String option, boolean value) {
         return value ? option : "";
     }
 
@@ -205,9 +223,10 @@ class CreateContainerCommandImpl implements Docker.CreateContainerCommand {
 
         return Stream.of(
                 "--name " + containerName.asString(),
-                "--hostname " + hostName,
-                "--cpu-shares " + containerResources.cpuShares(),
-                "--memory " + containerResources.memoryBytes(),
+                toOptionalOption("--hostname", hostName),
+                toOptionalOption("--cpu-shares", containerResources.map(ContainerResources::cpuShares)),
+                toOptionalOption("--cpus", containerResources.map(ContainerResources::cpus)),
+                toOptionalOption("--memory", containerResources.map(ContainerResources::memoryBytes)),
                 toRepeatedOption("--label", labelList),
                 toRepeatedOption("--ulimit", ulimitList),
                 toRepeatedOption("--env", environmentAssignments),
@@ -228,8 +247,8 @@ class CreateContainerCommandImpl implements Docker.CreateContainerCommand {
     /**
      * Generates a pseudo-random MAC address based on the hostname, IPv4- and IPv6-address.
      */
-    static String generateMACAddress(String hostname, Optional<String> ipv4Address, Optional<String> ipv6Address) {
-        final String seed = hostname + ipv4Address.orElse("") + ipv6Address.orElse("");
+    static String generateMACAddress(Optional<String> hostname, Optional<String> ipv4Address, Optional<String> ipv6Address) {
+        final String seed = hostname.orElse("") + ipv4Address.orElse("") + ipv6Address.orElse("");
         Random rand = getPRNG(seed);
         byte[] macAddr = new byte[6];
         rand.nextBytes(macAddr);
diff --git a/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/Docker.java b/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/Docker.java
index 1713b4570b8..e422e633966 100644
--- a/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/Docker.java
+++ b/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/Docker.java
@@ -15,6 +15,8 @@ import java.util.OptionalLong;
 public interface Docker {
 
     interface CreateContainerCommand {
+        CreateContainerCommand withHostName(String hostname);
+        CreateContainerCommand withResources(ContainerResources containerResources);
         CreateContainerCommand withLabel(String name, String value);
         CreateContainerCommand withEnvironment(String name, String value);
 
@@ -53,11 +55,7 @@ public interface Docker {
         void create();
     }
 
-    CreateContainerCommand createContainerCommand(
-            DockerImage dockerImage,
-            ContainerResources containerResources,
-            ContainerName containerName,
-            String hostName);
+    CreateContainerCommand createContainerCommand(DockerImage dockerImage, ContainerName containerName);
 
     Optional<ContainerStats> getContainerStats(ContainerName containerName);
 
diff --git a/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/DockerImpl.java b/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/DockerImpl.java
index 2ac68be2b2e..7ea46b1e361 100644
--- a/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/DockerImpl.java
+++ b/docker-api/src/main/java/com/yahoo/vespa/hosted/dockerapi/DockerImpl.java
@@ -6,9 +6,11 @@ import com.github.dockerjava.api.command.ExecCreateCmdResponse;
 import com.github.dockerjava.api.command.InspectContainerResponse;
 import com.github.dockerjava.api.command.InspectExecResponse;
 import com.github.dockerjava.api.command.InspectImageResponse;
+import com.github.dockerjava.api.command.UpdateContainerCmd;
 import com.github.dockerjava.api.exception.DockerClientException;
 import com.github.dockerjava.api.exception.NotFoundException;
 import com.github.dockerjava.api.exception.NotModifiedException;
+import com.github.dockerjava.api.model.HostConfig;
 import com.github.dockerjava.api.model.Image;
 import com.github.dockerjava.api.model.Statistics;
 import com.github.dockerjava.core.DefaultDockerClientConfig;
@@ -46,6 +48,7 @@ public class DockerImpl implements Docker {
     private static final Logger logger = Logger.getLogger(DockerImpl.class.getName());
 
     static final String LABEL_NAME_MANAGEDBY = "com.yahoo.vespa.managedby";
+    static final int CPU_PERIOD = 100_000; // 100 µs
     private static final String FRAMEWORK_CONTAINER_PREFIX = "/";
     private static final Duration WAIT_BEFORE_KILLING = Duration.ofSeconds(10);
 
@@ -114,9 +117,8 @@ public class DockerImpl implements Docker {
     }
 
     @Override
-    public CreateContainerCommand createContainerCommand(DockerImage image, ContainerResources containerResources,
-                                                         ContainerName name, String hostName) {
-        return new CreateContainerCommandImpl(dockerClient, image, containerResources, name, hostName);
+    public CreateContainerCommand createContainerCommand(DockerImage image, ContainerName containerName) {
+        return new CreateContainerCommandImpl(dockerClient, image, containerName);
     }
 
 
@@ -251,8 +253,7 @@ public class DockerImpl implements Docker {
                         new Container(
                                 response.getConfig().getHostName(),
                                 new DockerImage(response.getConfig().getImage()),
-                                new ContainerResources(response.getHostConfig().getCpuShares(),
-                                        response.getHostConfig().getMemory()),
+                                containerResourcesFromHostConfig(response.getHostConfig()),
                                 new ContainerName(decode(response.getName())),
                                 Container.State.valueOf(response.getState().getStatus().toUpperCase()),
                                 response.getState().getPid()
@@ -261,6 +262,12 @@ public class DockerImpl implements Docker {
                 .orElse(Stream.empty());
     }
 
+    private static ContainerResources containerResourcesFromHostConfig(HostConfig hostConfig) {
+        final double cpus = hostConfig.getCpuPeriod() > 0 ?
+                (double) hostConfig.getCpuQuota() / hostConfig.getCpuPeriod() : 0;
+        return new ContainerResources(cpus, hostConfig.getCpuShares(), hostConfig.getMemory());
+    }
+
     private boolean isManagedBy(com.github.dockerjava.api.model.Container container, String manager) {
         final Map<String, String> labels = container.getLabels();
         return labels != null && manager.equals(labels.get(LABEL_NAME_MANAGEDBY));
diff --git a/docker-api/src/test/java/com/yahoo/vespa/hosted/dockerapi/CreateContainerCommandImplTest.java b/docker-api/src/test/java/com/yahoo/vespa/hosted/dockerapi/CreateContainerCommandImplTest.java
index dc041b159b5..ccaf8fb652d 100644
--- a/docker-api/src/test/java/com/yahoo/vespa/hosted/dockerapi/CreateContainerCommandImplTest.java
+++ b/docker-api/src/test/java/com/yahoo/vespa/hosted/dockerapi/CreateContainerCommandImplTest.java
@@ -17,12 +17,14 @@ public class CreateContainerCommandImplTest {
     @Test
     public void testToString() throws UnknownHostException {
         DockerImage dockerImage = new DockerImage("docker.registry.domain.tld/my/image:1.2.3");
-        ContainerResources containerResources = new ContainerResources(100, 1024);
+        ContainerResources containerResources = new ContainerResources(2.5, 100, 1024);
         String hostname = "docker-1.region.domain.tld";
         ContainerName containerName = ContainerName.fromHostname(hostname);
 
         Docker.CreateContainerCommand createContainerCommand = new CreateContainerCommandImpl(
-                null, dockerImage, containerResources, containerName, hostname)
+                null, dockerImage, containerName)
+                .withHostName(hostname)
+                .withResources(containerResources)
                 .withLabel("my-label", "test-label")
                 .withUlimit("nofile", 1, 2)
                 .withUlimit("nproc", 10, 20)
@@ -41,6 +43,7 @@ public class CreateContainerCommandImplTest {
         assertEquals("--name docker-1 " +
                 "--hostname docker-1.region.domain.tld " +
                 "--cpu-shares 100 " +
+                "--cpus 2.5 " +
                 "--memory 1024 " +
                 "--label my-label=test-label " +
                 "--ulimit nofile=1:2 " +
@@ -71,7 +74,7 @@ public class CreateContainerCommandImplTest {
 
         Stream.of(addresses).forEach(address -> {
             String generatedMac = CreateContainerCommandImpl.generateMACAddress(
-                    address[0], Optional.ofNullable(address[1]), Optional.ofNullable(address[2]));
+                    Optional.of(address[0]), Optional.ofNullable(address[1]), Optional.ofNullable(address[2]));
             String expectedMac = address[3];
             assertEquals(expectedMac, generatedMac);
         });