make saveImage cancellable (podman-desktop#8232)

* feat: cancel saveImage Signed-off-by: Philippe Martin <[email protected]> * test: unit tests Signed-off-by: Philippe Martin <[email protected]>
feloy · Jul 29, 2024 · 110c796 · 110c796
1 parent 25245a6
commit 110c796
Show file tree

Hide file tree

Showing 4 changed files with 196 additions and 8 deletions.
diff --git a/packages/extension-api/src/extension-api.d.ts b/packages/extension-api/src/extension-api.d.ts
@@ -3565,8 +3565,9 @@ declare module '@podman-desktop/api' {
      * @param engineId the id of the engine managing the image, obtained from the result of {@link containerEngine.listImages}
      * @param id the id or name of the image on this engine, obtained from the result of {@link containerEngine.listImages}
      * @param filename the file on which to save the container image content
+     * @param token an optional cancellation token which will cancel saving the image on disk when the token is canceled
      */
-    export function saveImage(engineId: string, id: string, filename: string): Promise<void>;
+    export function saveImage(engineId: string, id: string, filename: string, token?: CancellationToken): Promise<void>;
 
     /**
      * List the container images. Only images from a final layer (no children) are returned.

diff --git a/packages/main/src/plugin/container-registry.spec.ts b/packages/main/src/plugin/container-registry.spec.ts
@@ -18,14 +18,16 @@
 
 import { EventEmitter } from 'node:events';
 import * as fs from 'node:fs';
-import { PassThrough } from 'node:stream';
+import os from 'node:os';
+import path from 'node:path';
+import { PassThrough, Readable } from 'node:stream';
 import * as streamPromises from 'node:stream/promises';
 
 import type * as podmanDesktopAPI from '@podman-desktop/api';
 import Dockerode from 'dockerode';
 import moment from 'moment';
 import nock from 'nock';
-import { beforeEach, describe, expect, test, vi } from 'vitest';
+import { afterEach, beforeEach, describe, expect, test, vi } from 'vitest';
 
 import type { ApiSenderType } from '/@/plugin/api.js';
 import type { Certificates } from '/@/plugin/certificates.js';
@@ -39,6 +41,7 @@ import type { ImageInfo } from '/@api/image-info.js';
 import type { ProviderContainerConnectionInfo } from '/@api/provider-info.js';
 
 import * as util from '../util.js';
+import { CancellationTokenRegistry } from './cancellation-token-registry.js';
 import type { ConfigurationRegistry } from './configuration-registry.js';
 import type { ContainerCreateOptions as PodmanContainerCreateOptions, LibPod } from './dockerode/libpod-dockerode.js';
 import { LibpodDockerode } from './dockerode/libpod-dockerode.js';
@@ -50,6 +53,7 @@ const tar: { pack: (dir: string) => NodeJS.ReadableStream } = require('tar-fs');
 /* eslint-disable @typescript-eslint/no-empty-function */
 /* eslint-disable @typescript-eslint/no-explicit-any */
 /* eslint-disable no-null/no-null */
+/* eslint-disable @typescript-eslint/consistent-type-imports */
 
 const fakeContainerWithComposeProject: Dockerode.ContainerInfo = {
   Id: '1234567890',
@@ -4916,3 +4920,161 @@ test('test removeManifest', async () => {
   expect(removeManifestMock).toBeCalledWith('manifestId');
   expect(result).toBeUndefined();
 });
+
+test('saveImage succeeds', async () => {
+  const dockerode = new Dockerode({ protocol: 'http', host: 'localhost' });
+  const stream: Dockerode.Image = {
+    get: vi.fn(),
+  } as unknown as Dockerode.Image;
+  const getImageMock = vi.fn().mockReturnValue(stream);
+  const api = {
+    ...dockerode,
+    getImage: getImageMock,
+  };
+
+  const pipelineMock = vi
+    .spyOn(streamPromises, 'pipeline')
+    .mockImplementation((_source: NodeJS.ReadableStream, _destination: NodeJS.WritableStream) => {
+      return Promise.resolve();
+    });
+
+  containerRegistry.addInternalProvider('podman1', {
+    name: 'podman-1',
+    id: 'podman1',
+    connection: {
+      type: 'podman',
+    },
+    api,
+  } as unknown as InternalContainerProvider);
+  await containerRegistry.saveImage('podman1', 'an-image', '/path/to/file');
+
+  expect(pipelineMock).toHaveBeenCalledOnce();
+});
+
+test('saveImage succeeds when a passing a cancellable token never canceled', async () => {
+  const cancellationTokenRegistry = new CancellationTokenRegistry();
+  const cancellableTokenId = cancellationTokenRegistry.createCancellationTokenSource();
+  const token = cancellationTokenRegistry.getCancellationTokenSource(cancellableTokenId)!.token;
+  const dockerode = new Dockerode({ protocol: 'http', host: 'localhost' });
+  const stream: Dockerode.Image = {
+    get: vi.fn(),
+  } as unknown as Dockerode.Image;
+  const getImageMock = vi.fn().mockReturnValue(stream);
+  const api = {
+    ...dockerode,
+    getImage: getImageMock,
+  };
+
+  const pipelineMock = vi
+    .spyOn(streamPromises, 'pipeline')
+    .mockImplementation((_source: NodeJS.ReadableStream, _destination: NodeJS.WritableStream) => {
+      return Promise.resolve();
+    });
+
+  containerRegistry.addInternalProvider('podman1', {
+    name: 'podman-1',
+    id: 'podman1',
+    connection: {
+      type: 'podman',
+    },
+    api,
+  } as unknown as InternalContainerProvider);
+  await containerRegistry.saveImage('podman1', 'an-image', '/path/to/file', token);
+
+  expect(pipelineMock).toHaveBeenCalledOnce();
+});
+
+describe('using fake timers', () => {
+  beforeEach(() => {
+    vi.useFakeTimers();
+  });
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  test('saveImage canceled during image download', async () => {
+    const cancellationTokenRegistry = new CancellationTokenRegistry();
+    const cancellableTokenId = cancellationTokenRegistry.createCancellationTokenSource();
+    const tokenSource = cancellationTokenRegistry.getCancellationTokenSource(cancellableTokenId)!;
+    const token = tokenSource.token;
+    const dockerode = new Dockerode({ protocol: 'http', host: 'localhost' });
+    const imageObjectGetMock = vi.fn().mockImplementation(() => {
+      return new Promise(resolve => {
+        setTimeout(() => {
+          resolve(undefined);
+        }, 1000);
+      });
+    });
+    const stream: Dockerode.Image = {
+      get: imageObjectGetMock,
+    } as unknown as Dockerode.Image;
+    const getImageMock = vi.fn().mockReturnValue(stream);
+    const api = {
+      ...dockerode,
+      getImage: getImageMock,
+    };
+
+    containerRegistry.addInternalProvider('podman1', {
+      name: 'podman-1',
+      id: 'podman1',
+      connection: {
+        type: 'podman',
+      },
+      api,
+    } as unknown as InternalContainerProvider);
+    setTimeout(() => {
+      tokenSource.cancel();
+    }, 500);
+
+    const savePromise = containerRegistry.saveImage('podman1', 'an-image', '/path/to/file', token);
+    vi.advanceTimersByTime(2000);
+
+    await expect(savePromise).rejects.toThrowError('saveImage operation canceled');
+  });
+});
+
+test('saveImage canceled during image saving on filesystem', async () => {
+  const streamModule = await vi.importActual<typeof import('node:stream/promises')>('node:stream/promises');
+  const fsModule = await vi.importActual<typeof import('node:fs')>('node:fs');
+  vi.mocked(streamPromises.pipeline).mockImplementation(streamModule.pipeline);
+  vi.mocked(fs.createWriteStream).mockImplementation(fsModule.createWriteStream);
+  const cancellationTokenRegistry = new CancellationTokenRegistry();
+  const cancellableTokenId = cancellationTokenRegistry.createCancellationTokenSource();
+  const tokenSource = cancellationTokenRegistry.getCancellationTokenSource(cancellableTokenId)!;
+  const token = tokenSource.token;
+  const dockerode = new Dockerode({ protocol: 'http', host: 'localhost' });
+  const imageObjectGetMock = vi.fn().mockResolvedValue(() => {
+    const stream = Readable.from(Buffer.from('a content'));
+    stream.on('readable', () => {
+      setTimeout(() => {
+        // too late
+        stream.read();
+      }, 300);
+    });
+    return stream;
+  });
+  const stream: Dockerode.Image = {
+    get: imageObjectGetMock,
+  } as unknown as Dockerode.Image;
+  const getImageMock = vi.fn().mockReturnValue(stream);
+  const api = {
+    ...dockerode,
+    getImage: getImageMock,
+  };
+
+  containerRegistry.addInternalProvider('podman1', {
+    name: 'podman-1',
+    id: 'podman1',
+    connection: {
+      type: 'podman',
+    },
+    api,
+  } as unknown as InternalContainerProvider);
+  setTimeout(() => {
+    tokenSource.cancel();
+  }, 50);
+
+  const tmpdir = os.tmpdir();
+  const savePromise = containerRegistry.saveImage('podman1', 'an-image', path.join(tmpdir, 'image-to-save'), token);
+  await expect(savePromise).rejects.toThrowError('The operation was aborted');
+});
diff --git a/packages/main/src/plugin/container-registry.ts b/packages/main/src/plugin/container-registry.ts
@@ -19,6 +19,7 @@
 import * as crypto from 'node:crypto';
 import { EventEmitter } from 'node:events';
 import * as fs from 'node:fs';
+import { rm } from 'node:fs/promises';
 import path from 'node:path';
 import { Writable } from 'node:stream';
 import { pipeline } from 'node:stream/promises';
@@ -2239,7 +2240,12 @@ export class ContainerProviderRegistry {
     }
   }
 
-  async saveImage(engineId: string, id: string, filename: string): Promise<void> {
+  async saveImage(
+    engineId: string,
+    id: string,
+    filename: string,
+    token?: containerDesktopAPI.CancellationToken,
+  ): Promise<void> {
     let telemetryOptions = {};
     try {
       // need to find the container engine of the container
@@ -2253,8 +2259,27 @@ export class ContainerProviderRegistry {
 
       const imageObject = provider.api.getImage(id);
       if (imageObject) {
-        const imageStream = await imageObject.get();
-        return pipeline(imageStream, fs.createWriteStream(filename));
+        // make the download of image cancellable
+        const getImageObjectPromise = imageObject.get();
+        const cancelPromise = new Promise<NodeJS.ReadableStream>((_, reject) => {
+          token?.onCancellationRequested(() => {
+            reject(new Error('saveImage operation canceled'));
+          });
+        });
+        const imageStream = await Promise.race<NodeJS.ReadableStream>([getImageObjectPromise, cancelPromise]);
+
+        // make the saving on filesystem cancellable
+        const ac = new AbortController();
+        const signal = ac.signal;
+        token?.onCancellationRequested(() => {
+          ac.abort();
+        });
+        try {
+          return await pipeline(imageStream, fs.createWriteStream(filename), { signal });
+        } catch (err: unknown) {
+          await rm(filename, { force: true });
+          throw err;
+        }
       }
     } catch (error) {
       telemetryOptions = { error: error };

diff --git a/packages/main/src/plugin/extension-loader.ts b/packages/main/src/plugin/extension-loader.ts
@@ -1099,8 +1099,8 @@ export class ExtensionLoader {
       listImages(options?: containerDesktopAPI.ListImagesOptions): Promise<containerDesktopAPI.ImageInfo[]> {
         return containerProviderRegistry.podmanListImages(options);
       },
-      saveImage(engineId: string, id: string, filename: string) {
-        return containerProviderRegistry.saveImage(engineId, id, filename);
+      saveImage(engineId: string, id: string, filename: string, token?: containerDesktopAPI.CancellationToken) {
+        return containerProviderRegistry.saveImage(engineId, id, filename, token);
       },
       pushImage(
         engineId: string,