From 0812d0b4955ab6d1a9b09a63e1686a6e2125b205 Mon Sep 17 00:00:00 2001
From: Felix Fontein <felix@fontein.de>
Date: Thu, 28 Dec 2023 21:42:55 +0100
Subject: [PATCH] Support labels and shm_size for image build. Allow to specify
 (swap) memory limits in other units than bytes. (#727)

---
 .../fragments/727-docker_image-build.yml      |  3 +
 plugins/modules/docker_container.py           |  3 +-
 plugins/modules/docker_image.py               | 63 ++++++++++++++++--
 .../docker_image/tasks/tests/basic.yml        |  2 +-
 .../docker_image/tasks/tests/options.yml      | 64 ++++++++++++++++++-
 5 files changed, 123 insertions(+), 12 deletions(-)
 create mode 100644 changelogs/fragments/727-docker_image-build.yml

diff --git a/changelogs/fragments/727-docker_image-build.yml b/changelogs/fragments/727-docker_image-build.yml
new file mode 100644
index 00000000..59750b90
--- /dev/null
+++ b/changelogs/fragments/727-docker_image-build.yml
@@ -0,0 +1,3 @@
+minor_changes:
+  - "docker_image - allow to specify labels and ``/dev/shm`` size when building images (https://github.com/ansible-collections/community.docker/issues/726, https://github.com/ansible-collections/community.docker/pull/727)."
+  - "docker_image - allow to specify memory size and swap memory size in other units than bytes (https://github.com/ansible-collections/community.docker/pull/727)."
diff --git a/plugins/modules/docker_container.py b/plugins/modules/docker_container.py
index f168d103..5590483a 100644
--- a/plugins/modules/docker_container.py
+++ b/plugins/modules/docker_container.py
@@ -156,7 +156,8 @@ options:
     type: float
   cpuset_cpus:
     description:
-      - CPUs in which to allow execution V(1,3) or V(1-3).
+      - CPUs in which to allow execution.
+      - For example V(1,3) or V(1-3).
     type: str
   cpuset_mems:
     description:
diff --git a/plugins/modules/docker_image.py b/plugins/modules/docker_image.py
index 02ae3838..613e8ea4 100644
--- a/plugins/modules/docker_image.py
+++ b/plugins/modules/docker_image.py
@@ -112,13 +112,21 @@ options:
         suboptions:
           memory:
             description:
-              - Set memory limit for build.
-            type: int
+              - "Memory limit for build in format C(<number>[<unit>]). Number is a positive integer.
+                Unit can be V(B) (byte), V(K) (kibibyte, 1024B), V(M) (mebibyte), V(G) (gibibyte),
+                V(T) (tebibyte), or V(P) (pebibyte)."
+              - Omitting the unit defaults to bytes.
+              - Before community.docker 3.6.0, no units were allowed.
+            type: str
           memswap:
             description:
-              - Total memory (memory + swap).
-              - Use V(-1) to disable swap.
-            type: int
+              - "Total memory limit (memory + swap) for build in format C(<number>[<unit>]), or
+                the special values V(unlimited) or V(-1) for unlimited swap usage.
+                Number is a positive integer. Unit can be V(B) (byte), V(K) (kibibyte, 1024B),
+                V(M) (mebibyte), V(G) (gibibyte), V(T) (tebibyte), or V(P) (pebibyte)."
+              - Omitting the unit defaults to bytes.
+              - Before community.docker 3.6.0, no units were allowed, and neither was the special value V(unlimited).
+            type: str
           cpushares:
             description:
               - CPU shares (relative weight).
@@ -144,6 +152,19 @@ options:
           - Platform in the format C(os[/arch[/variant]]).
         type: str
         version_added: 1.1.0
+      shm_size:
+        description:
+          - "Size of C(/dev/shm) in format C(<number>[<unit>]). Number is positive integer.
+            Unit can be V(B) (byte), V(K) (kibibyte, 1024B), V(M) (mebibyte), V(G) (gibibyte),
+            V(T) (tebibyte), or V(P) (pebibyte)."
+          - Omitting the unit defaults to bytes. If you omit the size entirely, Docker daemon uses V(64M).
+        type: str
+        version_added: 3.6.0
+      labels:
+        description:
+          - Dictionary of key value pairs.
+        type: dict
+        version_added: 3.6.0
   archive_path:
     description:
       - Use with O(state=present) to archive an image to a C(.tar) file.
@@ -338,6 +359,7 @@ import os
 import traceback
 
 from ansible.module_utils.common.text.converters import to_native
+from ansible.module_utils.common.text.formatters import human_to_bytes
 
 from ansible_collections.community.docker.plugins.module_utils.common_api import (
     AnsibleDockerClient,
@@ -377,6 +399,17 @@ from ansible_collections.community.docker.plugins.module_utils._api.utils.utils
 )
 
 
+def convert_to_bytes(value, module, name, unlimited_value=None):
+    if value is None:
+        return value
+    try:
+        if unlimited_value is not None and value in ('unlimited', str(unlimited_value)):
+            return unlimited_value
+        return human_to_bytes(value)
+    except ValueError as exc:
+        module.fail_json(msg='Failed to convert %s to bytes: %s' % (name, to_native(exc)))
+
+
 class ImageManager(DockerBaseClass):
 
     def __init__(self, client, results):
@@ -402,6 +435,12 @@ class ImageManager(DockerBaseClass):
         self.archive_path = parameters['archive_path']
         self.cache_from = build.get('cache_from')
         self.container_limits = build.get('container_limits')
+        if self.container_limits and 'memory' in self.container_limits:
+            self.container_limits['memory'] = convert_to_bytes(
+                self.container_limits['memory'], self.client.module, 'build.container_limits.memory')
+        if self.container_limits and 'memswap' in self.container_limits:
+            self.container_limits['memswap'] = convert_to_bytes(
+                self.container_limits['memswap'], self.client.module, 'build.container_limits.memswap', unlimited_value=-1)
         self.dockerfile = build.get('dockerfile')
         self.force_source = parameters['force_source']
         self.force_absent = parameters['force_absent']
@@ -424,6 +463,8 @@ class ImageManager(DockerBaseClass):
         self.buildargs = build.get('args')
         self.build_platform = build.get('platform')
         self.use_config_proxy = build.get('use_config_proxy')
+        self.shm_size = convert_to_bytes(build.get('shm_size'), self.client.module, 'build.shm_size')
+        self.labels = clean_dict_booleans_for_docker_api(build.get('labels'))
 
         # If name contains a tag, it takes precedence over tag parameter.
         if not is_image_name_id(self.name):
@@ -825,6 +866,12 @@ class ImageManager(DockerBaseClass):
         if self.build_platform is not None:
             params['platform'] = self.build_platform
 
+        if self.shm_size is not None:
+            params['shmsize'] = self.shm_size
+
+        if self.labels:
+            params['labels'] = json.dumps(self.labels)
+
         if context is not None:
             headers['Content-Type'] = 'application/tar'
 
@@ -945,8 +992,8 @@ def main():
         build=dict(type='dict', options=dict(
             cache_from=dict(type='list', elements='str'),
             container_limits=dict(type='dict', options=dict(
-                memory=dict(type='int'),
-                memswap=dict(type='int'),
+                memory=dict(type='str'),
+                memswap=dict(type='str'),
                 cpushares=dict(type='int'),
                 cpusetcpus=dict(type='str'),
             )),
@@ -962,6 +1009,8 @@ def main():
             target=dict(type='str'),
             etc_hosts=dict(type='dict'),
             platform=dict(type='str'),
+            shm_size=dict(type='str'),
+            labels=dict(type='dict'),
         )),
         archive_path=dict(type='path'),
         force_source=dict(type='bool', default=False),
diff --git a/tests/integration/targets/docker_image/tasks/tests/basic.yml b/tests/integration/targets/docker_image/tasks/tests/basic.yml
index 78b4f773..068e92f4 100644
--- a/tests/integration/targets/docker_image/tasks/tests/basic.yml
+++ b/tests/integration/targets/docker_image/tasks/tests/basic.yml
@@ -117,7 +117,7 @@
   register: fail_3
   ignore_errors: true
 
-- name: buildargs
+- name: Build image ID (must fail)
   docker_image:
     source: build
     name: "{{ present_1.image.Id }}"
diff --git a/tests/integration/targets/docker_image/tasks/tests/options.yml b/tests/integration/targets/docker_image/tasks/tests/options.yml
index 0670f133..b8f6c700 100644
--- a/tests/integration/targets/docker_image/tasks/tests/options.yml
+++ b/tests/integration/targets/docker_image/tasks/tests/options.yml
@@ -76,7 +76,7 @@
     build:
       path: "{{ remote_tmp_dir }}/files"
       container_limits:
-        memory: 4000
+        memory: 4KB
       pull: false
     source: build
   ignore_errors: true
@@ -88,8 +88,8 @@
     build:
       path: "{{ remote_tmp_dir }}/files"
       container_limits:
-        memory: 7000000
-        memswap: 8000000
+        memory: 7MB
+        memswap: 8MB
       pull: false
     source: build
   register: container_limits_2
@@ -444,3 +444,61 @@
 - assert:
     that:
       - path_1 is changed
+
+####################################################################
+## build.shm_size ##################################################
+####################################################################
+
+- name: Build image with custom shm_size
+  docker_image:
+    name: "{{ iname }}"
+    build:
+      path: "{{ remote_tmp_dir }}/files"
+      dockerfile: "MyDockerfile"
+      pull: false
+      shm_size: 128MB
+    source: build
+  register: path_1
+
+- name: cleanup
+  docker_image:
+    name: "{{ iname }}"
+    state: absent
+    force_absent: true
+
+- assert:
+    that:
+      - path_1 is changed
+
+####################################################################
+## build.labels ####################################################
+####################################################################
+
+- name: Build image with labels
+  docker_image:
+    name: "{{ iname }}"
+    build:
+      path: "{{ remote_tmp_dir }}/files"
+      dockerfile: "MyDockerfile"
+      pull: false
+      labels:
+        FOO: BAR
+        this is a label: this is the label's value
+    source: build
+  register: labels_1
+
+- name: cleanup
+  docker_image:
+    name: "{{ iname }}"
+    state: absent
+    force_absent: true
+
+- name: Show image information
+  debug:
+    var: labels_1.image
+
+- assert:
+    that:
+      - labels_1 is changed
+      - labels_1.image.Config.Labels.FOO == 'BAR'
+      - labels_1.image.Config.Labels["this is a label"] == "this is the label's value"