Add support for blkio config keys

Signed-off-by: Joffrey F <[email protected]>

compose/config/config.py

@@ -16,6 +16,7 @@ from . import types
 from .. import const
 from ..const import COMPOSEFILE_V1 as V1
 from ..utils import build_string_dict
+from ..utils import parse_bytes
 from ..utils import parse_nanoseconds_int
 from ..utils import splitdrive
 from ..version import ComposeVersion
@@ -108,6 +109,7 @@ DOCKER_CONFIG_KEYS = [
 ]
 
 ALLOWED_KEYS = DOCKER_CONFIG_KEYS + [
+    'blkio_config',
     'build',
     'container_name',
     'credential_spec',
@@ -726,8 +728,9 @@ def process_service(service_config):
         if field in service_dict:
             service_dict[field] = to_list(service_dict[field])
 
-    service_dict = process_healthcheck(service_dict, service_config.name)
-    service_dict = process_ports(service_dict)
+    service_dict = process_blkio_config(process_ports(
+        process_healthcheck(service_dict, service_config.name)
+    ))
 
     return service_dict
 
@@ -754,6 +757,28 @@ def process_depends_on(service_dict):
     return service_dict
 
 
+def process_blkio_config(service_dict):
+    if not service_dict.get('blkio_config'):
+        return service_dict
+
+    for field in ['device_read_bps', 'device_write_bps']:
+        if field in service_dict['blkio_config']:
+            for v in service_dict['blkio_config'].get(field, []):
+                v['rate'] = parse_bytes(v.get('rate', 0))
+
+    for field in ['device_read_iops', 'device_write_iops']:
+        if field in service_dict['blkio_config']:
+            for v in service_dict['blkio_config'].get(field, []):
+                try:
+                    v['rate'] = int(v.get('rate', 0))
+                except ValueError:
+                    raise ConfigurationError(
+                        'Invalid IOPS value: "{}". Must be a positive integer.'.format(v.get('rate'))
+                    )
+
+    return service_dict
+
+
 def process_healthcheck(service_dict, service_name):
     if 'healthcheck' not in service_dict:
         return service_dict
@@ -940,6 +965,7 @@ def merge_service_dicts(base, override, version):
 
     md.merge_field('logging', merge_logging, default={})
     merge_ports(md, base, override)
+    md.merge_field('blkio_config', merge_blkio_config, default={})
 
     for field in set(ALLOWED_KEYS) - set(md):
         md.merge_scalar(field)
@@ -993,6 +1019,26 @@ def merge_build(output, base, override):
     return dict(md)
 
 
+def merge_blkio_config(base, override):
+    md = MergeDict(base, override)
+    md.merge_scalar('weight')
+
+    def merge_blkio_limits(base, override):
+        index = dict((b['path'], b) for b in base)
+        for o in override:
+            index[o['path']] = o
+
+        return sorted(list(index.values()), key=lambda x: x['path'])
+
+    for field in [
+            "device_read_bps", "device_read_iops", "device_write_bps",
+            "device_write_iops", "weight_device",
+    ]:
+        md.merge_field(field, merge_blkio_limits, default=[])
+
+    return dict(md)
+
+
 def merge_logging(base, override):
     md = MergeDict(base, override)
     md.merge_scalar('driver')

compose/config/config_schema_v2.0.json

@@ -50,6 +50,33 @@
       "type": "object",
 
       "properties": {
+        "blkio_config": {
+          "type": "object",
+          "properties": {
+            "device_read_bps": {
+              "type": "array",
+              "items": {"$ref": "#/definitions/blkio_limit"}
+            },
+            "device_read_iops": {
+              "type": "array",
+              "items": {"$ref": "#/definitions/blkio_limit"}
+            },
+            "device_write_bps": {
+              "type": "array",
+              "items": {"$ref": "#/definitions/blkio_limit"}
+            },
+            "device_write_iops": {
+              "type": "array",
+              "items": {"$ref": "#/definitions/blkio_limit"}
+            },
+            "weight": {"type": "integer"},
+            "weight_device": {
+              "type": "array",
+              "items": {"$ref": "#/definitions/blkio_weight"}
+            }
+          },
+          "additionalProperties": false
+        },
         "build": {
           "oneOf": [
             {"type": "string"},
@@ -326,6 +353,23 @@
       ]
     },
 
+    "blkio_limit": {
+      "type": "object",
+      "properties": {
+        "path": {"type": "string"},
+        "rate": {"type": ["integer", "string"]}
+      },
+      "additionalProperties": false
+    },
+    "blkio_weight": {
+      "type": "object",
+      "properties": {
+        "path": {"type": "string"},
+        "weight": {"type": "integer"}
+      },
+      "additionalProperties": false
+    },
+
     "constraints": {
       "service": {
         "id": "#/definitions/constraints/service",

compose/config/config_schema_v2.1.json

@@ -50,6 +50,34 @@
       "type": "object",
 
       "properties": {
+        "blkio_config": {
+          "type": "object",
+          "properties": {
+            "device_read_bps": {
+              "type": "array",
+              "items": {"$ref": "#/definitions/blkio_limit"}
+            },
+            "device_read_iops": {
+              "type": "array",
+              "items": {"$ref": "#/definitions/blkio_limit"}
+            },
+            "device_write_bps": {
+              "type": "array",
+              "items": {"$ref": "#/definitions/blkio_limit"}
+            },
+            "device_write_iops": {
+              "type": "array",
+              "items": {"$ref": "#/definitions/blkio_limit"}
+            },
+            "weight": {"type": "integer"},
+            "weight_device": {
+              "type": "array",
+              "items": {"$ref": "#/definitions/blkio_weight"}
+            }
+          },
+          "additionalProperties": false
+        },
+
         "build": {
           "oneOf": [
             {"type": "string"},
@@ -376,6 +404,23 @@
       ]
     },
 
+    "blkio_limit": {
+      "type": "object",
+      "properties": {
+        "path": {"type": "string"},
+        "rate": {"type": ["integer", "string"]}
+      },
+      "additionalProperties": false
+    },
+    "blkio_weight": {
+      "type": "object",
+      "properties": {
+        "path": {"type": "string"},
+        "weight": {"type": "integer"}
+      },
+      "additionalProperties": false
+    },
+
     "constraints": {
       "service": {
         "id": "#/definitions/constraints/service",

compose/config/config_schema_v2.2.json

@@ -50,6 +50,34 @@
       "type": "object",
 
       "properties": {
+        "blkio_config": {
+          "type": "object",
+          "properties": {
+            "device_read_bps": {
+              "type": "array",
+              "items": {"$ref": "#/definitions/blkio_limit"}
+            },
+            "device_read_iops": {
+              "type": "array",
+              "items": {"$ref": "#/definitions/blkio_limit"}
+            },
+            "device_write_bps": {
+              "type": "array",
+              "items": {"$ref": "#/definitions/blkio_limit"}
+            },
+            "device_write_iops": {
+              "type": "array",
+              "items": {"$ref": "#/definitions/blkio_limit"}
+            },
+            "weight": {"type": "integer"},
+            "weight_device": {
+              "type": "array",
+              "items": {"$ref": "#/definitions/blkio_weight"}
+            }
+          },
+          "additionalProperties": false
+        },
+
         "build": {
           "oneOf": [
             {"type": "string"},
@@ -383,6 +411,23 @@
       ]
     },
 
+    "blkio_limit": {
+      "type": "object",
+      "properties": {
+        "path": {"type": "string"},
+        "rate": {"type": ["integer", "string"]}
+      },
+      "additionalProperties": false
+    },
+    "blkio_weight": {
+      "type": "object",
+      "properties": {
+        "path": {"type": "string"},
+        "weight": {"type": "integer"}
+      },
+      "additionalProperties": false
+    },
+
     "constraints": {
       "service": {
         "id": "#/definitions/constraints/service",

compose/config/config_schema_v2.3.json

@@ -50,6 +50,34 @@
       "type": "object",
 
       "properties": {
+        "blkio_config": {
+          "type": "object",
+          "properties": {
+            "device_read_bps": {
+              "type": "array",
+              "items": {"$ref": "#/definitions/blkio_limit"}
+            },
+            "device_read_iops": {
+              "type": "array",
+              "items": {"$ref": "#/definitions/blkio_limit"}
+            },
+            "device_write_bps": {
+              "type": "array",
+              "items": {"$ref": "#/definitions/blkio_limit"}
+            },
+            "device_write_iops": {
+              "type": "array",
+              "items": {"$ref": "#/definitions/blkio_limit"}
+            },
+            "weight": {"type": "integer"},
+            "weight_device": {
+              "type": "array",
+              "items": {"$ref": "#/definitions/blkio_weight"}
+            }
+          },
+          "additionalProperties": false
+        },
+
         "build": {
           "oneOf": [
             {"type": "string"},
@@ -384,6 +412,23 @@
       ]
     },
 
+    "blkio_limit": {
+      "type": "object",
+      "properties": {
+        "path": {"type": "string"},
+        "rate": {"type": ["integer", "string"]}
+      },
+      "additionalProperties": false
+    },
+    "blkio_weight": {
+      "type": "object",
+      "properties": {
+        "path": {"type": "string"},
+        "weight": {"type": "integer"}
+      },
+      "additionalProperties": false
+    },
+
     "constraints": {
       "service": {
         "id": "#/definitions/constraints/service",

compose/service.py

@@ -813,6 +813,7 @@ class Service(object):
         options = dict(self.options, **override_options)
 
         logging_dict = options.get('logging', None)
+        blkio_config = convert_blkio_config(options.get('blkio_config', None))
         log_config = get_log_config(logging_dict)
         init_path = None
         if isinstance(options.get('init'), six.string_types):
@@ -868,7 +869,13 @@ class Service(object):
             volume_driver=options.get('volume_driver'),
             cpuset_cpus=options.get('cpuset'),
             cpu_shares=options.get('cpu_shares'),
-            storage_opt=options.get('storage_opt')
+            storage_opt=options.get('storage_opt'),
+            blkio_weight=blkio_config.get('weight'),
+            blkio_weight_device=blkio_config.get('weight_device'),
+            device_read_bps=blkio_config.get('device_read_bps'),
+            device_read_iops=blkio_config.get('device_read_iops'),
+            device_write_bps=blkio_config.get('device_write_bps'),
+            device_write_iops=blkio_config.get('device_write_iops'),
         )
 
     def get_secret_volumes(self):
@@ -1395,3 +1402,22 @@ def build_container_ports(container_ports, options):
                 port = tuple(port.split('/'))
             ports.append(port)
     return ports
+
+
+def convert_blkio_config(blkio_config):
+    result = {}
+    if blkio_config is None:
+        return result
+
+    result['weight'] = blkio_config.get('weight')
+    for field in [
+        "device_read_bps", "device_read_iops", "device_write_bps",
+        "device_write_iops", "weight_device",
+    ]:
+        if field not in blkio_config:
+            continue
+        arr = []
+        for item in blkio_config[field]:
+            arr.append(dict([(k.capitalize(), v) for k, v in item.items()]))
+        result[field] = arr
+    return result

compose/utils.py

@@ -9,7 +9,10 @@ import logging
 import ntpath
 
 import six
+from docker.errors import DockerException
+from docker.utils import parse_bytes as sdk_parse_bytes
 
+from .config.errors import ConfigurationError
 from .errors import StreamParseError
 from .timeparse import timeparse
 
@@ -133,3 +136,10 @@ def splitdrive(path):
     if path[0] in ['.', '\\', '/', '~']:
         return ('', path)
     return ntpath.splitdrive(path)
+
+
+def parse_bytes(n):
+    try:
+        return sdk_parse_bytes(n)
+    except DockerException:
+        raise ConfigurationError('Invalid format for bytes value: {}'.format(n))

tests/integration/service_test.py

@@ -203,6 +203,34 @@ class ServiceTest(DockerClientTestCase):
         service.start_container(container)
         assert container.get('HostConfig.ReadonlyRootfs') == read_only
 
+    def test_create_container_with_blkio_config(self):
+        blkio_config = {
+            'weight': 300,
+            'weight_device': [{'path': '/dev/sda', 'weight': 200}],
+            'device_read_bps': [{'path': '/dev/sda', 'rate': 1024 * 1024 * 100}],
+            'device_read_iops': [{'path': '/dev/sda', 'rate': 1000}],
+            'device_write_bps': [{'path': '/dev/sda', 'rate': 1024 * 1024}],
+            'device_write_iops': [{'path': '/dev/sda', 'rate': 800}]
+        }
+        service = self.create_service('web', blkio_config=blkio_config)
+        container = service.create_container()
+        assert container.get('HostConfig.BlkioWeight') == 300
+        assert container.get('HostConfig.BlkioWeightDevice') == [{
+            'Path': '/dev/sda', 'Weight': 200
+        }]
+        assert container.get('HostConfig.BlkioDeviceReadBps') == [{
+            'Path': '/dev/sda', 'Rate': 1024 * 1024 * 100
+        }]
+        assert container.get('HostConfig.BlkioDeviceWriteBps') == [{
+            'Path': '/dev/sda', 'Rate': 1024 * 1024
+        }]
+        assert container.get('HostConfig.BlkioDeviceReadIOps') == [{
+            'Path': '/dev/sda', 'Rate': 1000
+        }]
+        assert container.get('HostConfig.BlkioDeviceWriteIOps') == [{
+            'Path': '/dev/sda', 'Rate': 800
+        }]
+
     def test_create_container_with_security_opt(self):
         security_opt = ['label:disable']
         service = self.create_service('db', security_opt=security_opt)

tests/unit/config/config_test.py

@@ -2151,6 +2151,53 @@ class ConfigTest(unittest.TestCase):
         actual = config.merge_service_dicts(base, override, V2_2)
         assert actual == {'image': 'bar', 'scale': 4}
 
+    def test_merge_blkio_config(self):
+        base = {
+            'image': 'bar',
+            'blkio_config': {
+                'weight': 300,
+                'weight_device': [
+                    {'path': '/dev/sda1', 'weight': 200}
+                ],
+                'device_read_iops': [
+                    {'path': '/dev/sda1', 'rate': 300}
+                ],
+                'device_write_iops': [
+                    {'path': '/dev/sda1', 'rate': 1000}
+                ]
+            }
+        }
+
+        override = {
+            'blkio_config': {
+                'weight': 450,
+                'weight_device': [
+                    {'path': '/dev/sda2', 'weight': 400}
+                ],
+                'device_read_iops': [
+                    {'path': '/dev/sda1', 'rate': 2000}
+                ],
+                'device_read_bps': [
+                    {'path': '/dev/sda1', 'rate': 1024}
+                ]
+            }
+        }
+
+        actual = config.merge_service_dicts(base, override, V2_2)
+        assert actual == {
+            'image': 'bar',
+            'blkio_config': {
+                'weight': override['blkio_config']['weight'],
+                'weight_device': (
+                    base['blkio_config']['weight_device'] +
+                    override['blkio_config']['weight_device']
+                ),
+                'device_read_iops': override['blkio_config']['device_read_iops'],
+                'device_read_bps': override['blkio_config']['device_read_bps'],
+                'device_write_iops': base['blkio_config']['device_write_iops']
+            }
+        }
+
     def test_external_volume_config(self):
         config_details = build_config_details({
             'version': '2',