From 110c796e85d0720052db87e86714df274fecf233 Mon Sep 17 00:00:00 2001 From: Philippe Martin Date: Mon, 29 Jul 2024 15:21:19 +0200 Subject: [PATCH] make saveImage cancellable (#8232) * feat: cancel saveImage Signed-off-by: Philippe Martin * test: unit tests Signed-off-by: Philippe Martin --- packages/extension-api/src/extension-api.d.ts | 3 +- .../src/plugin/container-registry.spec.ts | 166 +++++++++++++++++- .../main/src/plugin/container-registry.ts | 31 +++- packages/main/src/plugin/extension-loader.ts | 4 +- 4 files changed, 196 insertions(+), 8 deletions(-) diff --git a/packages/extension-api/src/extension-api.d.ts b/packages/extension-api/src/extension-api.d.ts index 14d17ef7a44e4..ba2c6a3386e67 100644 --- 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; + export function saveImage(engineId: string, id: string, filename: string, token?: CancellationToken): Promise; /** * 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 index 1e86acefb74c7..f789a0c725821 100644 --- 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('node:stream/promises'); + const fsModule = await vi.importActual('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 index df55a290d205e..7d9f550dcfdf1 100644 --- 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 { + async saveImage( + engineId: string, + id: string, + filename: string, + token?: containerDesktopAPI.CancellationToken, + ): Promise { 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((_, reject) => { + token?.onCancellationRequested(() => { + reject(new Error('saveImage operation canceled')); + }); + }); + const imageStream = await Promise.race([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 index cfbbd7fca25ba..8d4c164c2b461 100644 --- 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 { 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,