From 58a6c69855126143825687230578157c36b86a56 Mon Sep 17 00:00:00 2001
From: Farid Zakaria <farid.m.zakaria@gmail.com>
Date: Mon, 26 Nov 2018 11:41:19 -0800
Subject: [PATCH] Allow to set Shared Memory size (shmSize) for
 GenericContainer (#957)

---
 .../containers/GenericContainer.java          | 33 ++++++++++++++++++-
 .../junit/GenericContainerRuleTest.java       | 16 +++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/core/src/main/java/org/testcontainers/containers/GenericContainer.java b/core/src/main/java/org/testcontainers/containers/GenericContainer.java
index 32cae84951a..0554fd0e2de 100644
--- a/core/src/main/java/org/testcontainers/containers/GenericContainer.java
+++ b/core/src/main/java/org/testcontainers/containers/GenericContainer.java
@@ -6,6 +6,7 @@
 import com.github.dockerjava.api.model.Bind;
 import com.github.dockerjava.api.model.ContainerNetwork;
 import com.github.dockerjava.api.model.ExposedPort;
+import com.github.dockerjava.api.model.HostConfig;
 import com.github.dockerjava.api.model.Info;
 import com.github.dockerjava.api.model.Link;
 import com.github.dockerjava.api.model.PortBinding;
@@ -153,6 +154,13 @@ public class GenericContainer<SELF extends GenericContainer<SELF>>
     @Nullable
     private String workingDirectory = null;
 
+    /**
+     * The shared memory size to use when starting the container.
+     * This value is in bytes.
+     */
+    @Nullable
+    private Long shmSize;
+
     private Map<MountableFile, String> copyToFileContainerPathMap = new HashMap<>();
 
     /*
@@ -318,6 +326,17 @@ private void tryStart(Profiler profiler) {
         }
     }
 
+    /**
+     * Set any custom settings for the create command such as shared memory size.
+     */
+    private HostConfig buildHostConfig() {
+        HostConfig config = new HostConfig();
+        if (shmSize != null) {
+            config.withShmSize(shmSize);
+        }
+        return config;
+    }
+
     private void connectToPortForwardingNetwork(String networkMode) {
         PortForwardingContainer.INSTANCE.getNetwork().map(ContainerNetwork::getNetworkID).ifPresent(networkId -> {
             if (!Arrays.asList(networkId, "none", "host").contains(networkMode)) {
@@ -434,7 +453,9 @@ public Set<Integer> getLivenessCheckPortNumbers() {
     }
 
     private void applyConfiguration(CreateContainerCmd createCommand) {
-
+        HostConfig hostConfig = buildHostConfig();
+        createCommand.withHostConfig(hostConfig);
+        
         // Set up exposed ports (where there are no host port bindings defined)
         ExposedPort[] portArray = exposedPorts.stream()
                 .map(ExposedPort::new)
@@ -1157,6 +1178,16 @@ public SELF withCreateContainerCmdModifier(Consumer<CreateContainerCmd> modifier
         return self();
     }
 
+    /**
+     * Size of /dev/shm
+     * @param bytes The number of megabytes to assign the shared memory. If null, it will apply the Docker default which is 64 MB.
+     * @return this
+     */
+    public SELF withSharedMemorySize(Long bytes) {
+        this.shmSize = bytes;
+        return self();
+    }
+
     /**
      * Convenience class with access to non-public members of GenericContainer.
      *
diff --git a/core/src/test/java/org/testcontainers/junit/GenericContainerRuleTest.java b/core/src/test/java/org/testcontainers/junit/GenericContainerRuleTest.java
index deb13453287..f1aa5bd86b1 100644
--- a/core/src/test/java/org/testcontainers/junit/GenericContainerRuleTest.java
+++ b/core/src/test/java/org/testcontainers/junit/GenericContainerRuleTest.java
@@ -1,11 +1,13 @@
 package org.testcontainers.junit;
 
+import com.github.dockerjava.api.model.HostConfig;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.util.concurrent.Uninterruptibles;
 import com.mongodb.MongoClient;
 import com.mongodb.client.MongoCollection;
 import com.mongodb.client.MongoDatabase;
 import com.rabbitmq.client.*;
+import org.apache.commons.io.FileUtils;
 import org.bson.Document;
 import org.junit.*;
 import org.rnorth.ducttape.RetryCountExceededException;
@@ -375,4 +377,18 @@ public void addExposedPortAfterWithExposedPortsTest() {
         assertTrue("withExposedPort should be exposed", redis.getExposedPorts().contains(REDIS_PORT));
         assertTrue("addExposedPort should be exposed", redis.getExposedPorts().contains(8987));
     }
+
+    @Test
+    public void sharedMemorySetTest() {
+        try (GenericContainer containerWithSharedMemory = new GenericContainer("busybox:1.29")
+            .withSharedMemorySize(1024L * FileUtils.ONE_MB)) {
+
+            containerWithSharedMemory.start();
+
+            HostConfig hostConfig =
+                containerWithSharedMemory.getDockerClient().inspectContainerCmd(containerWithSharedMemory.getContainerId())
+                    .exec().getHostConfig();
+            assertEquals("Shared memory not set on container", hostConfig.getShmSize(), 1024 * FileUtils.ONE_MB);
+        }
+    }
 }