Merge pull request #554 from docker/ulimit_support

ulimits support

10 files changed, 145 insertions, 9 deletions

diff --git a/docker/client.py b/docker/client.py
index 78ca03a..a94849d 100644
--- a/docker/client.py
+++ b/docker/client.py
@@ -933,7 +933,7 @@ class Client(requests.Session):
               dns=None, dns_search=None, volumes_from=None, network_mode=None,
               restart_policy=None, cap_add=None, cap_drop=None, devices=None,
               extra_hosts=None, read_only=None, pid_mode=None, ipc_mode=None,
-              security_opt=None):
+              security_opt=None, ulimits=None):
 
         if utils.compare_version('1.10', self._version) < 0:
             if dns is not None:
@@ -965,6 +965,12 @@ class Client(requests.Session):
                     'pid_mode is only supported for API version >= 1.17'
                 )
 
+        if utils.compare_version('1.18', self._version) < 0:
+            if ulimits is not None:
+                raise errors.InvalidVersion(
+                    'ulimits is only supported for API version >= 1.18'
+                )
+
         start_config = utils.create_host_config(
             binds=binds, port_bindings=port_bindings, lxc_conf=lxc_conf,
             publish_all_ports=publish_all_ports, links=links, dns=dns,
@@ -972,7 +978,7 @@ class Client(requests.Session):
             cap_drop=cap_drop, volumes_from=volumes_from, devices=devices,
             network_mode=network_mode, restart_policy=restart_policy,
             extra_hosts=extra_hosts, read_only=read_only, pid_mode=pid_mode,
-            ipc_mode=ipc_mode, security_opt=security_opt
+            ipc_mode=ipc_mode, security_opt=security_opt, ulimits=ulimits
         )
 
         if isinstance(container, dict):
diff --git a/docker/utils/__init__.py b/docker/utils/__init__.py
index 0c5934f..411d278 100644
--- a/docker/utils/__init__.py
+++ b/docker/utils/__init__.py
@@ -4,3 +4,5 @@ from .utils import (
     kwargs_from_env, convert_filters, create_host_config,
     create_container_config, parse_bytes, ping_registry
 ) # flake8: noqa
+
+from .types import Ulimit # flake8: noqa
\ No newline at end of file
diff --git a/docker/utils/types.py b/docker/utils/types.py
new file mode 100644
index 0000000..f960afb
--- /dev/null
+++ b/docker/utils/types.py
@@ -0,0 +1,49 @@
+import six
+
+
+class DictType(dict):
+    def __init__(self, init):
+        for k, v in six.iteritems(init):
+            self[k] = v
+
+
+class Ulimit(DictType):
+    def __init__(self, **kwargs):
+        name = kwargs.get('name', kwargs.get('Name'))
+        soft = kwargs.get('soft', kwargs.get('Soft'))
+        hard = kwargs.get('hard', kwargs.get('Hard'))
+        if not isinstance(name, six.string_types):
+            raise ValueError("Ulimit.name must be a string")
+        if soft and not isinstance(soft, int):
+            raise ValueError("Ulimit.soft must be an integer")
+        if hard and not isinstance(hard, int):
+            raise ValueError("Ulimit.hard must be an integer")
+        super(Ulimit, self).__init__({
+            'Name': name,
+            'Soft': soft,
+            'Hard': hard
+        })
+
+    @property
+    def name(self):
+        return self['Name']
+
+    @name.setter
+    def name(self, value):
+        self['Name'] = value
+
+    @property
+    def soft(self):
+        return self.get('Soft')
+
+    @soft.setter
+    def soft(self, value):
+        self['Soft'] = value
+
+    @property
+    def hard(self):
+        return self.get('Hard')
+
+    @hard.setter
+    def hard(self, value):
+        self['Hard'] = value
diff --git a/docker/utils/utils.py b/docker/utils/utils.py
index 5c729a0..fb32904 100644
--- a/docker/utils/utils.py
+++ b/docker/utils/utils.py
@@ -28,6 +28,7 @@ import six
 
 from .. import errors
 from .. import tls
+from .types import Ulimit
 
 
 DEFAULT_HTTP_HOST = "127.0.0.1"
@@ -358,7 +359,7 @@ def create_host_config(
     dns=None, dns_search=None, volumes_from=None, network_mode=None,
     restart_policy=None, cap_add=None, cap_drop=None, devices=None,
     extra_hosts=None, read_only=None, pid_mode=None, ipc_mode=None,
-    security_opt=None
+    security_opt=None, ulimits=None
 ):
     host_config = {}
 
@@ -451,6 +452,17 @@ def create_host_config(
     if lxc_conf is not None:
         host_config['LxcConf'] = lxc_conf
 
+    if ulimits is not None:
+        if not isinstance(ulimits, list):
+            raise errors.DockerException(
+                'Invalid type for ulimits param: expected list but found'
+                ' {0}'.format(type(ulimits)))
+        host_config['Ulimits'] = []
+        for l in ulimits:
+            if not isinstance(l, Ulimit):
+                l = Ulimit(**l)
+            host_config['Ulimits'].append(l)
+
     return host_config
 
 
diff --git a/docs/api.md b/docs/api.md
index aa5faa9..5afe533 100644
--- a/docs/api.md
+++ b/docs/api.md
@@ -749,6 +749,7 @@ from. Optionally a single string joining container id's with commas
 * pid_mode (str): if set to "host", use the host PID namespace inside the
   container
 * security_opt (list): A list of string values to customize labels for MLS systems, such as SELinux.
+* ulimits (list): A list of dicts or `docker.utils.Ulimit` objects.
 
 ```python
 >>> from docker import Client
diff --git a/docs/hostconfig.md b/docs/hostconfig.md
index 3c11031..1d82c83 100644
--- a/docs/hostconfig.md
+++ b/docs/hostconfig.md
@@ -85,6 +85,8 @@ for example:
 * read_only (bool): mount the container's root filesystem as read only
 * pid_mode (str): if set to "host", use the host PID namespace inside the
   container
+* security_opt (list): A list of string values to customize labels for MLS systems, such as SELinux.
+* ulimits (list): A list of dicts or `docker.utils.Ulimit` objects.
 
 **Returns** (dict) HostConfig dictionary
 
diff --git a/tests/base.py b/tests/base.py
new file mode 100644
index 0000000..1d5a370
--- /dev/null
+++ b/tests/base.py
@@ -0,0 +1,11 @@
+import sys
+import unittest
+
+import six
+
+
+class BaseTestCase(unittest.TestCase):
+    def assertIn(self, object, collection):
+        if six.PY2 and sys.version_info[1] <= 6:
+            return self.assertTrue(object in collection)
+        return super(BaseTestCase, self).assertIn(object, collection)
diff --git a/tests/integration_test.py b/tests/integration_test.py
index f35a7ac..ba69e22 100644
--- a/tests/integration_test.py
+++ b/tests/integration_test.py
@@ -927,6 +927,19 @@ class TestStartContainerWithVolumesFrom(BaseTestCase):
         self.assertCountEqual(info['HostConfig']['VolumesFrom'], vol_names)
 
 
+class TestStartContainerWithUlimits(BaseTestCase):
+    def runTest(self):
+        ulimit = docker.utils.Ulimit('nofile', 4096, 4096)
+
+        res0 = self.client.create_container('busybox', 'true')
+        container1_id = res0['Id']
+        self.tmp_containers.append(container1_id)
+        self.client.start(container1_id, ulimits=[ulimit])
+
+        info = self.client.inspect_container(container1_id)
+        self.assertCountEqual(info['HostConfig']['Ulimits'], [ulimit])
+
+
 class TestStartContainerWithLinks(BaseTestCase):
     def runTest(self):
         res0 = self.client.create_container(
diff --git a/tests/test.py b/tests/test.py
index 04b5a11..59bf4cc 100644
--- a/tests/test.py
+++ b/tests/test.py
@@ -35,6 +35,7 @@ import docker
 import requests
 import six
 
+import base
 import fake_api
 
 try:
@@ -102,7 +103,7 @@ class Cleanup(object):
 
 @mock.patch.multiple('docker.Client', get=fake_request, post=fake_request,
                      put=fake_request, delete=fake_request)
-class DockerClientTest(Cleanup, unittest.TestCase):
+class DockerClientTest(Cleanup, base.BaseTestCase):
     def setUp(self):
         self.client = docker.Client()
         # Force-clear authconfig to avoid tampering with the tests
@@ -2295,7 +2296,7 @@ class DockerClientTest(Cleanup, unittest.TestCase):
         )
 
 
-class StreamTest(Cleanup, unittest.TestCase):
+class StreamTest(Cleanup, base.BaseTestCase):
 
     def setUp(self):
         socket_dir = tempfile.mkdtemp()
diff --git a/tests/utils_test.py b/tests/utils_test.py
index 852c4ca..1f18512 100644
--- a/tests/utils_test.py
+++ b/tests/utils_test.py
@@ -6,13 +6,15 @@ from docker.client import Client
 from docker.errors import DockerException
 from docker.utils import (
     parse_repository_tag, parse_host, convert_filters, kwargs_from_env,
-    create_host_config
+    create_host_config, Ulimit
 )
 from docker.utils.ports import build_port_bindings, split_port
 from docker.auth import resolve_authconfig
 
+import base
 
-class UtilsTest(unittest.TestCase):
+
+class UtilsTest(base.BaseTestCase):
     longMessage = True
 
     def setUp(self):
@@ -98,10 +100,47 @@ class UtilsTest(unittest.TestCase):
         for filters, expected in tests:
             self.assertEqual(convert_filters(filters), expected)
 
-    def test_create_host_config(self):
+    def test_create_empty_host_config(self):
         empty_config = create_host_config()
         self.assertEqual(empty_config, {})
 
+    def test_create_host_config_dict_ulimit(self):
+        ulimit_dct = {'name': 'nofile', 'soft': 8096}
+        config = create_host_config(ulimits=[ulimit_dct])
+        self.assertIn('Ulimits', config)
+        self.assertEqual(len(config['Ulimits']), 1)
+        ulimit_obj = config['Ulimits'][0]
+        self.assertTrue(isinstance(ulimit_obj, Ulimit))
+        self.assertEqual(ulimit_obj.name, ulimit_dct['name'])
+        self.assertEqual(ulimit_obj.soft, ulimit_dct['soft'])
+        self.assertEqual(ulimit_obj['Soft'], ulimit_obj.soft)
+
+    def test_create_host_config_dict_ulimit_capitals(self):
+        ulimit_dct = {'Name': 'nofile', 'Soft': 8096, 'Hard': 8096 * 4}
+        config = create_host_config(ulimits=[ulimit_dct])
+        self.assertIn('Ulimits', config)
+        self.assertEqual(len(config['Ulimits']), 1)
+        ulimit_obj = config['Ulimits'][0]
+        self.assertTrue(isinstance(ulimit_obj, Ulimit))
+        self.assertEqual(ulimit_obj.name, ulimit_dct['Name'])
+        self.assertEqual(ulimit_obj.soft, ulimit_dct['Soft'])
+        self.assertEqual(ulimit_obj.hard, ulimit_dct['Hard'])
+        self.assertEqual(ulimit_obj['Soft'], ulimit_obj.soft)
+
+    def test_create_host_config_obj_ulimit(self):
+        ulimit_dct = Ulimit(name='nofile', soft=8096)
+        config = create_host_config(ulimits=[ulimit_dct])
+        self.assertIn('Ulimits', config)
+        self.assertEqual(len(config['Ulimits']), 1)
+        ulimit_obj = config['Ulimits'][0]
+        self.assertTrue(isinstance(ulimit_obj, Ulimit))
+        self.assertEqual(ulimit_obj, ulimit_dct)
+
+    def test_ulimit_invalid_type(self):
+        self.assertRaises(ValueError, lambda: Ulimit(name=None))
+        self.assertRaises(ValueError, lambda: Ulimit(name='hello', soft='123'))
+        self.assertRaises(ValueError, lambda: Ulimit(name='hello', hard='456'))
+
     def test_resolve_authconfig(self):
         auth_config = {
             'https://index.docker.io/v1/': {'auth': 'indexuser'},
@@ -222,7 +261,7 @@ class UtilsTest(unittest.TestCase):
         self.assertRaises(
             ValueError,
             lambda: split_port("0.0.0.0:1000-1010:2000-2002/tcp")
-            )
+        )
 
     def test_port_and_range_invalid(self):
         self.assertRaises(ValueError,