@aws-sdk/lib-transfer-manager 0.0.0-stage → 0.1.0

package/dist-es/submodules/transfer-manager/worker-http-handler.js

@@ -0,0 +1,453 @@
+import { HttpResponse } from "@smithy/core/protocols";
+import { NodeHttpHandler } from "@smithy/node-http-handler";
+import { existsSync } from "node:fs";
+import { cpus } from "node:os";
+import * as path from "node:path";
+import { Readable } from "node:stream";
+import { Worker } from "node:worker_threads";
+function buildDownloadError(msg) {
+    const error = new Error(msg.error);
+    if (msg.name)
+        error.name = msg.name;
+    if (msg.code)
+        error.code = msg.code;
+    return error;
+}
+export function createEmptyReadable() {
+    return new Readable({
+        read() {
+            this.push(null);
+        },
+    });
+}
+export function defaultWorkerCount() {
+    return cpus().length;
+}
+export class WorkerHttpHandler {
+    workers = [];
+    inflightRequests = new Map();
+    inflightDownloads = new Map();
+    inflightStreamDownloads = new Map();
+    nextRequestId = 0;
+    workerInflightCounts = [];
+    workerThreadCount;
+    maxConcurrentUploads;
+    maxConcurrentDownloads;
+    useODirect;
+    logger;
+    initialized = false;
+    initPromise;
+    fallbackHandler;
+    completedDownloads = new Map();
+    completedTransfers = new Map();
+    metadata = { handlerProtocol: "http/1.1" };
+    constructor(options) {
+        this.workerThreadCount = options?.workerThreadCount ?? defaultWorkerCount();
+        this.maxConcurrentUploads = options?.maxConcurrentUploads ?? 32;
+        this.maxConcurrentDownloads = options?.maxConcurrentDownloads ?? 32;
+        this.useODirect = options?.useODirect ?? false;
+        this.logger = options?.logger;
+        this.fallbackHandler = new NodeHttpHandler();
+    }
+    getDownloadResult(token) {
+        const result = this.completedDownloads.get(token);
+        if (result) {
+            this.completedDownloads.delete(token);
+        }
+        return result;
+    }
+    getStreamDownloadResult(token) {
+        const result = this.completedTransfers.get(token);
+        if (result) {
+            this.completedTransfers.delete(token);
+        }
+        return result;
+    }
+    resolveWorkerConfig() {
+        const submodulePath = path.join(__dirname, "..", "worker", "index.js");
+        if (existsSync(submodulePath)) {
+            return { workerPath: submodulePath };
+        }
+        const tsPath = path.join(__dirname, "..", "worker", "index.ts");
+        if (existsSync(tsPath)) {
+            return { workerPath: tsPath, workerOptions: { execArgv: ["--require", "tsx/cjs"] } };
+        }
+        return { workerPath: submodulePath };
+    }
+    ensureInitialized() {
+        if (this.initialized)
+            return Promise.resolve();
+        if (this.initPromise)
+            return this.initPromise;
+        this.initPromise = new Promise((resolve, reject) => {
+            const { workerPath, workerOptions } = this.resolveWorkerConfig();
+            let readyCount = 0;
+            let errorCount = 0;
+            let settled = false;
+            this.workerInflightCounts = new Array(this.workerThreadCount).fill(0);
+            for (let i = 0; i < this.workerThreadCount; i++) {
+                const worker = new Worker(workerPath, workerOptions);
+                this.workers.push(worker);
+                const workerIndex = i;
+                worker.on("message", (msg) => {
+                    if (msg.type === "ready") {
+                        worker.postMessage({
+                            type: "config",
+                            maxSockets: Math.max(50, Math.ceil(Math.max(this.maxConcurrentUploads, this.maxConcurrentDownloads) / this.workerThreadCount)),
+                            useODirect: this.useODirect,
+                        });
+                        readyCount++;
+                        if (!settled && readyCount === this.workerThreadCount) {
+                            settled = true;
+                            this.initialized = true;
+                            resolve();
+                        }
+                        return;
+                    }
+                    if (msg.type === "workerLog") {
+                        this.logger?.[msg.level]?.(`[S3TransferManager:worker] ${msg.message}`);
+                        return;
+                    }
+                    if (msg.type === "httpResponse") {
+                        const pending = this.inflightRequests.get(msg.id);
+                        this.inflightRequests.delete(msg.id);
+                        if (pending) {
+                            this.workerInflightCounts[pending.workerIndex]--;
+                            pending.resolve({
+                                response: new HttpResponse({
+                                    statusCode: msg.response.statusCode,
+                                    headers: msg.response.headers,
+                                    body: msg.response.body,
+                                }),
+                            });
+                        }
+                        return;
+                    }
+                    if (msg.type === "httpError") {
+                        const pending = this.inflightRequests.get(msg.id);
+                        this.inflightRequests.delete(msg.id);
+                        if (pending) {
+                            this.workerInflightCounts[pending.workerIndex]--;
+                            const error = Object.assign(new Error(msg.error), {
+                                ...(msg.code && { code: msg.code }),
+                                ...(msg.name && { name: msg.name }),
+                            });
+                            pending.reject(error);
+                        }
+                        return;
+                    }
+                    if (msg.type === "httpDownloadResult") {
+                        const pending = this.inflightDownloads.get(msg.id);
+                        this.inflightDownloads.delete(msg.id);
+                        if (pending) {
+                            this.workerInflightCounts[pending.workerIndex]--;
+                            pending.resolve({
+                                statusCode: msg.statusCode,
+                                headers: msg.headers,
+                                bytesWritten: msg.bytesWritten,
+                                checksum: msg.checksum,
+                            });
+                        }
+                        return;
+                    }
+                    if (msg.type === "httpDownloadError") {
+                        const pending = this.inflightDownloads.get(msg.id);
+                        this.inflightDownloads.delete(msg.id);
+                        if (pending) {
+                            this.workerInflightCounts[pending.workerIndex]--;
+                            pending.reject(buildDownloadError(msg));
+                            return;
+                        }
+                        const pendingTransfer = this.inflightStreamDownloads.get(msg.id);
+                        this.inflightStreamDownloads.delete(msg.id);
+                        if (pendingTransfer) {
+                            this.workerInflightCounts[pendingTransfer.workerIndex]--;
+                            pendingTransfer.reject(buildDownloadError(msg));
+                        }
+                        return;
+                    }
+                    if (msg.type === "httpDownloadHttpError") {
+                        const pending = this.inflightDownloads.get(msg.id);
+                        this.inflightDownloads.delete(msg.id);
+                        if (pending) {
+                            this.workerInflightCounts[pending.workerIndex]--;
+                            pending.resolve({
+                                statusCode: msg.statusCode,
+                                headers: msg.headers,
+                                bytesWritten: 0,
+                                errorBody: msg.body,
+                            });
+                            return;
+                        }
+                        const pendingTransfer = this.inflightStreamDownloads.get(msg.id);
+                        this.inflightStreamDownloads.delete(msg.id);
+                        if (pendingTransfer) {
+                            this.workerInflightCounts[pendingTransfer.workerIndex]--;
+                            pendingTransfer.resolve({
+                                rangeIndex: -1,
+                                buffer: new ArrayBuffer(0),
+                                byteLength: 0,
+                                statusCode: msg.statusCode,
+                                headers: msg.headers,
+                                errorBody: msg.body,
+                            });
+                        }
+                        return;
+                    }
+                    if (msg.type === "httpDownloadStreamResult") {
+                        const pending = this.inflightStreamDownloads.get(msg.id);
+                        this.inflightStreamDownloads.delete(msg.id);
+                        if (pending) {
+                            this.workerInflightCounts[pending.workerIndex]--;
+                            pending.resolve({
+                                rangeIndex: msg.rangeIndex,
+                                buffer: msg.buffer,
+                                byteLength: msg.byteLength,
+                                statusCode: msg.statusCode,
+                                headers: msg.headers,
+                                checksum: msg.checksum,
+                            });
+                        }
+                        return;
+                    }
+                });
+                worker.on("error", (err) => {
+                    if (!settled) {
+                        errorCount++;
+                        if (errorCount + readyCount === this.workerThreadCount) {
+                            settled = true;
+                            this.initPromise = undefined;
+                            reject(new Error(`Failed to initialize worker threads: ${err.message}`));
+                        }
+                    }
+                    for (const [id, pending] of this.inflightRequests) {
+                        if (pending.workerIndex === workerIndex) {
+                            pending.reject(err);
+                            this.inflightRequests.delete(id);
+                        }
+                    }
+                    for (const [id, pending] of this.inflightDownloads) {
+                        if (pending.workerIndex === workerIndex) {
+                            pending.reject(err);
+                            this.inflightDownloads.delete(id);
+                        }
+                    }
+                    for (const [id, pending] of this.inflightStreamDownloads) {
+                        if (pending.workerIndex === workerIndex) {
+                            pending.reject(err);
+                            this.inflightStreamDownloads.delete(id);
+                        }
+                    }
+                    this.workerInflightCounts[workerIndex] = Infinity;
+                });
+            }
+        });
+        return this.initPromise;
+    }
+    pickWorkerIndex() {
+        let minIndex = 0;
+        for (let i = 1; i < this.workerInflightCounts.length; i++) {
+            if (this.workerInflightCounts[i] < this.workerInflightCounts[minIndex]) {
+                minIndex = i;
+            }
+        }
+        return minIndex;
+    }
+    dispatchToWorker(id, message) {
+        const workerIndex = this.pickWorkerIndex();
+        this.workerInflightCounts[workerIndex]++;
+        return new Promise((resolve, reject) => {
+            this.inflightRequests.set(id, { resolve, reject, workerIndex });
+            this.workers[workerIndex].postMessage(message);
+        });
+    }
+    dispatchDownloadToWorker(id, message) {
+        const workerIndex = this.pickWorkerIndex();
+        this.workerInflightCounts[workerIndex]++;
+        return new Promise((resolve, reject) => {
+            this.inflightDownloads.set(id, { resolve, reject, workerIndex });
+            this.workers[workerIndex].postMessage(message);
+        });
+    }
+    extractPartNumber(request) {
+        const partNumber = request.query?.partNumber;
+        if (partNumber) {
+            return Number(Array.isArray(partNumber) ? partNumber[0] : partNumber);
+        }
+        return undefined;
+    }
+    returnBuffer(buffer, workerIndex) {
+        const targetIndex = workerIndex ?? this.nextRequestId % this.workers.length;
+        const message = { type: "returnBuffer", buffer };
+        this.workers[targetIndex]?.postMessage(message, [buffer]);
+    }
+    async handle(request, options) {
+        const downloadDataToFile = options?.downloadDataToFile;
+        if (downloadDataToFile) {
+            await this.ensureInitialized();
+            const id = this.nextRequestId++;
+            const serializedRequest = {
+                method: request.method,
+                protocol: request.protocol,
+                hostname: request.hostname,
+                port: request.port,
+                path: request.path,
+                query: request.query,
+                headers: request.headers,
+            };
+            const message = {
+                type: "httpDownloadToFile",
+                id,
+                request: serializedRequest,
+                filePath: downloadDataToFile.filePath,
+                offset: downloadDataToFile.offset,
+                expectedLength: downloadDataToFile.expectedLength,
+                checksumAlgorithm: downloadDataToFile.checksumAlgorithm,
+            };
+            const result = await this.dispatchDownloadToWorker(id, message);
+            if (result.errorBody !== undefined) {
+                return {
+                    response: new HttpResponse({
+                        statusCode: result.statusCode,
+                        headers: result.headers,
+                        body: Readable.from([Buffer.from(result.errorBody)]),
+                    }),
+                };
+            }
+            if (downloadDataToFile.resultToken) {
+                this.completedDownloads.set(downloadDataToFile.resultToken, {
+                    bytesWritten: result.bytesWritten,
+                    checksum: result.checksum,
+                });
+            }
+            const responseHeaders = {};
+            for (const [key, value] of Object.entries(result.headers)) {
+                if (!key.startsWith("x-amz-checksum-")) {
+                    responseHeaders[key] = value;
+                }
+            }
+            return {
+                response: new HttpResponse({
+                    statusCode: result.statusCode,
+                    headers: responseHeaders,
+                    body: createEmptyReadable(),
+                }),
+            };
+        }
+        const downloadStream = options?.downloadStream;
+        if (downloadStream) {
+            await this.ensureInitialized();
+            const id = this.nextRequestId++;
+            const serializedRequest = {
+                method: request.method,
+                protocol: request.protocol,
+                hostname: request.hostname,
+                port: request.port,
+                path: request.path,
+                query: request.query,
+                headers: request.headers,
+            };
+            const message = {
+                type: "httpDownloadStream",
+                id,
+                request: serializedRequest,
+                expectedSize: downloadStream.expectedSize,
+                rangeIndex: downloadStream.rangeIndex,
+                checksumAlgorithm: downloadStream.checksumAlgorithm,
+            };
+            const workerIndex = this.pickWorkerIndex();
+            this.workerInflightCounts[workerIndex]++;
+            const result = await new Promise((resolve, reject) => {
+                this.inflightStreamDownloads.set(id, { resolve, reject, workerIndex });
+                this.workers[workerIndex].postMessage(message);
+            });
+            if (result.errorBody !== undefined) {
+                return {
+                    response: new HttpResponse({
+                        statusCode: result.statusCode,
+                        headers: result.headers,
+                        body: Readable.from([Buffer.from(result.errorBody)]),
+                    }),
+                };
+            }
+            const responseHeaders = {};
+            for (const [key, value] of Object.entries(result.headers)) {
+                if (!key.startsWith("x-amz-checksum-")) {
+                    responseHeaders[key] = value;
+                }
+            }
+            if (downloadStream.resultToken) {
+                this.completedTransfers.set(downloadStream.resultToken, result);
+            }
+            return {
+                response: new HttpResponse({
+                    statusCode: result.statusCode,
+                    headers: responseHeaders,
+                    body: createEmptyReadable(),
+                }),
+            };
+        }
+        const dataSource = options?.dataSource;
+        const partNumber = this.extractPartNumber(request);
+        if (!dataSource || !partNumber) {
+            return this.fallbackHandler.handle(request, options);
+        }
+        await this.ensureInitialized();
+        const id = this.nextRequestId++;
+        const serializedRequest = {
+            method: request.method,
+            protocol: request.protocol,
+            hostname: request.hostname,
+            port: request.port,
+            path: request.path,
+            query: request.query,
+            headers: request.headers,
+        };
+        if (dataSource.type === "file") {
+            const { filePath, partSize, totalFileSize, startOffset = 0, checksumAlgorithm, checksumHeader } = dataSource;
+            const partOffset = (partNumber - 1) * partSize;
+            const length = Math.min(partSize, totalFileSize - partOffset);
+            const offset = startOffset + partOffset;
+            const message = {
+                type: "httpRequestFromFile",
+                id,
+                request: serializedRequest,
+                filePath,
+                offset,
+                length,
+                checksumAlgorithm,
+                checksumHeader,
+            };
+            return this.dispatchToWorker(id, message);
+        }
+        const { sharedBuffer, startOffset = 0, partSize, totalSize, checksumAlgorithm, checksumHeader } = dataSource;
+        const partOffset = (partNumber - 1) * partSize;
+        const offset = startOffset + partOffset;
+        const length = Math.min(partSize, totalSize - partOffset);
+        const message = {
+            type: "httpRequestFromRAM",
+            id,
+            request: serializedRequest,
+            sharedBuffer,
+            offset,
+            length,
+            checksumAlgorithm,
+            checksumHeader,
+        };
+        return this.dispatchToWorker(id, message);
+    }
+    updateHttpClientConfig(_key, _value) { }
+    httpHandlerConfigs() {
+        return {};
+    }
+    destroy() {
+        for (const worker of this.workers) {
+            worker.postMessage({ type: "done" });
+        }
+        this.workers = [];
+        this.workerInflightCounts = [];
+        this.initialized = false;
+        this.initPromise = undefined;
+        this.fallbackHandler.destroy();
+    }
+}