Keep shared images on update (#5268)

* Test stub for keeping shared images after update

* Keep shared images on addon update

* ImageNotFound should only skip the one image not all

* Fix tests and nonetype error

* Normalize logic between two cleanup methods

supervisor/docker/addon.py

@@ -709,6 +709,28 @@ class DockerAddon(DockerInterface):
             with suppress(DockerError):
                 await self.cleanup()
 
+    @Job(name="docker_addon_cleanup", limit=JobExecutionLimit.GROUP_WAIT)
+    async def cleanup(
+        self,
+        old_image: str | None = None,
+        image: str | None = None,
+        version: AwesomeVersion | None = None,
+    ) -> None:
+        """Check if old version exists and cleanup other versions of image not in use."""
+        await self.sys_run_in_executor(
+            self.sys_docker.cleanup_old_images,
+            (image := image or self.image),
+            version or self.version,
+            {old_image} if old_image else None,
+            keep_images={
+                f"{addon.image}:{addon.version}"
+                for addon in self.sys_addons.installed
+                if addon.slug != self.addon.slug
+                and addon.image
+                and addon.image in {old_image, image}
+            },
+        )
+
     @Job(
         name="docker_addon_write_stdin",
         limit=JobExecutionLimit.GROUP_ONCE,

supervisor/docker/interface.py

@@ -512,14 +512,14 @@ class DockerInterface(JobGroup):
         return b""
 
     @Job(name="docker_interface_cleanup", limit=JobExecutionLimit.GROUP_WAIT)
-    def cleanup(
+    async def cleanup(
         self,
         old_image: str | None = None,
         image: str | None = None,
         version: AwesomeVersion | None = None,
-    ) -> Awaitable[None]:
+    ) -> None:
         """Check if old version exists and cleanup."""
-        return self.sys_run_in_executor(
+        await self.sys_run_in_executor(
             self.sys_docker.cleanup_old_images,
             image or self.image,
             version or self.version,

supervisor/docker/manager.py

@@ -548,10 +548,13 @@ class DockerAPI:
         current_image: str,
         current_version: AwesomeVersion,
         old_images: set[str] | None = None,
+        *,
+        keep_images: set[str] | None = None,
     ) -> None:
         """Clean up old versions of an image."""
+        image = f"{current_image}:{current_version!s}"
         try:
-            current: Image = self.images.get(f"{current_image}:{current_version!s}")
+            keep: set[str] = {self.images.get(image).id}
         except ImageNotFound:
             raise DockerNotFound(
                 f"{current_image} not found for cleanup", _LOGGER.warning
@@ -561,6 +564,19 @@ class DockerAPI:
                 f"Can't get {current_image} for cleanup", _LOGGER.warning
             ) from err
 
+        if keep_images:
+            keep_images -= {image}
+            try:
+                for image in keep_images:
+                    # If its not found, no need to preserve it from getting removed
+                    with suppress(ImageNotFound):
+                        keep.add(self.images.get(image).id)
+            except (DockerException, requests.RequestException) as err:
+                raise DockerError(
+                    f"Failed to get one or more images from {keep} during cleanup",
+                    _LOGGER.warning,
+                ) from err
+
         # Cleanup old and current
         image_names = list(
             old_images | {current_image} if old_images else {current_image}
@@ -573,7 +589,7 @@ class DockerAPI:
             ) from err
 
         for image in images_list:
-            if current.id == image.id:
+            if image.id in keep:
                 continue
 
             with suppress(DockerException, requests.RequestException):

tests/addons/test_manager.py

@@ -50,6 +50,19 @@ async def fixture_remove_wait_boot(coresys: CoreSys) -> None:
     coresys.config.wait_boot = 0
 
 
+@pytest.fixture(name="install_addon_example_image")
+def fixture_install_addon_example_image(coresys: CoreSys, repository) -> Addon:
+    """Install local_example add-on with image."""
+    store = coresys.addons.store["local_example_image"]
+    coresys.addons.data.install(store)
+    # pylint: disable-next=protected-access
+    coresys.addons.data._data = coresys.addons.data._schema(coresys.addons.data._data)
+
+    addon = Addon(coresys, store.slug)
+    coresys.addons.local[addon.slug] = addon
+    yield addon
+
+
 async def test_image_added_removed_on_update(
     coresys: CoreSys, install_addon_ssh: Addon
 ):
@@ -399,7 +412,10 @@ async def test_store_data_changes_during_update(
         ):
             await coresys.addons.update("local_ssh")
             cleanup.assert_called_once_with(
-                "test_image", AwesomeVersion("1.1.1"), {"local/amd64-addon-ssh"}
+                "test_image",
+                AwesomeVersion("1.1.1"),
+                {"local/amd64-addon-ssh"},
+                keep_images=set(),
             )
 
     update_task = coresys.create_task(simulate_update())
@@ -491,3 +507,44 @@ async def test_shared_image_kept_on_uninstall(
         "force": True,
     }
     assert not coresys.addons.get("local_example", local_only=True)
+
+
+async def test_shared_image_kept_on_update(
+    coresys: CoreSys, install_addon_example_image: Addon, docker: DockerAPI
+):
+    """Test if two addons share an image it is not removed on update."""
+    # Clone example to a new mock copy so two share an image
+    # But modify version in store so Supervisor sees an update
+    curr_store_data = deepcopy(coresys.store.data.addons["local_example_image"])
+    curr_store = AddonStore(coresys, "local_example2", curr_store_data)
+    install_addon_example_image.data_store["version"] = "1.3.0"
+    new_store_data = deepcopy(coresys.store.data.addons["local_example_image"])
+    new_store = AddonStore(coresys, "local_example2", new_store_data)
+
+    coresys.store.data.addons["local_example2"] = new_store_data
+    coresys.addons.store["local_example2"] = new_store
+    coresys.addons.data.install(curr_store)
+    # pylint: disable-next=protected-access
+    coresys.addons.data._data = coresys.addons.data._schema(coresys.addons.data._data)
+
+    example_2 = Addon(coresys, curr_store.slug)
+    coresys.addons.local[example_2.slug] = example_2
+
+    assert example_2.version == "1.2.0"
+    assert install_addon_example_image.version == "1.2.0"
+
+    image_new = MagicMock()
+    image_new.id = "image_new"
+    image_old = MagicMock()
+    image_old.id = "image_old"
+    docker.images.get.side_effect = [image_new, image_old]
+    docker.images.list.return_value = [image_new, image_old]
+
+    await coresys.addons.update("local_example2")
+    docker.images.remove.assert_not_called()
+    assert example_2.version == "1.3.0"
+
+    docker.images.get.side_effect = [image_new]
+    await coresys.addons.update("local_example_image")
+    docker.images.remove.assert_called_once_with("image_old", force=True)
+    assert install_addon_example_image.version == "1.3.0"

tests/fixtures/addons/local/example_image/build.yaml

@@ -0,0 +1,14 @@
+# https://developers.home-assistant.io/docs/add-ons/configuration#add-on-dockerfile
+build_from:
+  aarch64: "ghcr.io/home-assistant/aarch64-base:3.15"
+  amd64: "ghcr.io/home-assistant/amd64-base:3.15"
+  armhf: "ghcr.io/home-assistant/armhf-base:3.15"
+  armv7: "ghcr.io/home-assistant/armv7-base:3.15"
+  i386: "ghcr.io/home-assistant/i386-base:3.15"
+labels:
+  org.opencontainers.image.title: "Home Assistant Add-on: Example add-on"
+  org.opencontainers.image.description: "Example add-on to use as a blueprint for new add-ons."
+  org.opencontainers.image.source: "https://github.com/home-assistant/addons-example"
+  org.opencontainers.image.licenses: "Apache License 2.0"
+args:
+  TEMPIO_VERSION: "2021.09.0"

tests/fixtures/addons/local/example_image/config.yaml

@@ -0,0 +1,26 @@
+# https://developers.home-assistant.io/docs/add-ons/configuration#add-on-config
+name: Example add-on with image
+version: "1.2.0"
+slug: example_image
+description: Example add-on
+url: "https://github.com/home-assistant/addons-example/tree/main/example"
+arch:
+  - armhf
+  - armv7
+  - aarch64
+  - amd64
+  - i386
+init: false
+map:
+  - share:rw
+  - addon_config
+options:
+  message: "Hello world..."
+schema:
+  message: "str?"
+ingress: true
+ingress_port: 0
+breaking_versions:
+  - test
+  - 1.0
+image: example/with-image

tests/store/test_store_manager.py

@@ -56,7 +56,12 @@ async def test_default_load(coresys: CoreSys):
         in store_manager.repository_urls
     )
     # NOTE: When adding new stores, make sure to add it to tests/fixtures/addons/git/
-    assert refresh_cache_calls == {"local_ssh", "local_example", "core_samba"}
+    assert refresh_cache_calls == {
+        "local_ssh",
+        "local_example",
+        "core_samba",
+        "local_example_image",
+    }
 
 
 async def test_load_with_custom_repository(coresys: CoreSys):