PyodideWorkerManager.ts

import path from 'path';
import { Worker } from 'worker_threads';
import { PublishDiagnosticsParams } from 'vscode-languageserver';
import { CloudFormationFileType } from '../../document/Document';
import { CfnLintInitializationSettings, CfnLintSettings } from '../../settings/Settings';
import { LoggerFactory } from '../../telemetry/LoggerFactory';
import { ScopedTelemetry } from '../../telemetry/ScopedTelemetry';
import { Telemetry } from '../../telemetry/TelemetryDecorator';
import { extractErrorMessage } from '../../utils/Errors';
import { retryWithExponentialBackoff } from '../../utils/Retry';
import { WorkerNotInitializedError } from './CfnLintErrors';

interface WorkerTask {
    id: string;
    action: string;
    payload: Record<string, unknown>;
    resolve: (value: unknown) => void;
    reject: (reason: Error) => void;
}

interface WorkerMessage {
    id?: string;
    type?: string;
    result?: unknown;
    error?: string;
    success?: boolean;
    data?: string;
}

export class PyodideWorkerManager {
    private worker: Worker | undefined = undefined;
    private nextTaskId = 1;
    private readonly tasks = new Map<string, WorkerTask>();
    private initialized = false;
    private initializationPromise: Promise<void> | undefined = undefined;

    @Telemetry() private readonly telemetry!: ScopedTelemetry;

    constructor(
        private readonly retryConfig: CfnLintInitializationSettings,
        private cfnLintSettings: CfnLintSettings,
        private readonly log = LoggerFactory.getLogger(PyodideWorkerManager),
    ) {}

    public async initialize(): Promise<void> {
        if (this.initialized) {
            return;
        }

        if (this.initializationPromise) {
            return await this.initializationPromise;
        }

        this.initializationPromise = this.initializeWithRetry();
        return await this.initializationPromise;
    }

    private async initializeWithRetry(): Promise<void> {
        let attemptCount = 0;
        return await retryWithExponentialBackoff(
            async () => {
                if (attemptCount > 0) {
                    this.telemetry.count('worker.restart', 1, { attributes: { attempt: attemptCount.toString() } });
                }
                attemptCount++;
                return await this.initializeWorker();
            },
            {
                maxRetries: this.retryConfig.maxRetries,
                initialDelayMs: this.retryConfig.initialDelayMs,
                maxDelayMs: this.retryConfig.maxDelayMs,
                backoffMultiplier: this.retryConfig.backoffMultiplier,
                jitterFactor: 0.1, // Add 10% jitter to prevent synchronized retry storms
                operationName: 'Pyodide initialization',
                totalTimeoutMs: this.retryConfig.totalTimeoutMs,
            },
            this.log,
        );
    }

    private async initializeWorker(): Promise<void> {
        return await new Promise<void>((resolve, reject) => {
            try {
                // Create worker
                // Use a path relative to the current file
                const workerPath = path.join(__dirname, 'pyodide-worker.js');
                this.log.info(`Loading worker from: ${workerPath}`);
                this.worker = new Worker(workerPath);

                // Add exit event handler to detect crashes
                this.worker.on('exit', (code) => {
                    if (code !== 0) {
                        this.log.error(`Worker exited unexpectedly with code ${code}`);
                        this.telemetry.count('worker.crash', 1, { attributes: { exitCode: code.toString() } });
                        this.initialized = false;
                        this.worker = undefined;

                        // Reject any pending tasks
                        for (const task of this.tasks.values()) {
                            task.reject(new Error(`Worker exited unexpectedly with code ${code}`));
                        }
                        this.tasks.clear();
                    }
                });

                // Set up message handler
                this.worker.on('message', this.handleWorkerMessage.bind(this));

                // Set up error handler
                this.worker.on('error', (error: Error) => {
                    this.log.error(error, 'Worker error');
                    reject(new Error(`Worker error: ${error.message}`));
                });

                // Initialize Pyodide in the worker
                const taskId = this.nextTaskId.toString();
                this.nextTaskId++;
                const pyodideStartTime = performance.now();

                const task: WorkerTask = {
                    id: taskId,
                    action: 'initialize',
                    payload: {},
                    resolve: (result: unknown) => {
                        this.initialized = true;
                        this.telemetry.histogram('pyodide.init.duration', performance.now() - pyodideStartTime, {
                            unit: 'ms',
                        });
                        this.telemetry.count('pyodide.init.success', 1);

                        // Track installation source
                        const initResult = result as { status: string; installSource?: string };
                        if (initResult.installSource === 'pypi') {
                            this.telemetry.count('init.pypi.success', 1);
                        } else if (initResult.installSource === 'wheels') {
                            this.telemetry.count('init.wheels.success', 1);
                        }

                        resolve();
                    },
                    reject: (reason: Error) => {
                        this.worker = undefined;
                        this.telemetry.error('pyodide.init.fault', reason, undefined, { captureErrorType: true });

                        // Try to determine if it was a PyPI or wheels failure
                        const errorMessage = reason.message || '';
                        if (errorMessage.includes('PyPI')) {
                            this.telemetry.error('init.pypi.fault', reason, undefined, {
                                captureErrorType: true,
                            });
                        } else if (errorMessage.includes('wheels') || errorMessage.includes('wheel')) {
                            this.telemetry.error('init.wheels.fault', reason, undefined, {
                                captureErrorType: true,
                            });
                        }

                        reject(reason);
                    },
                };

                this.tasks.set(taskId, task);
                this.worker.postMessage({
                    id: taskId,
                    action: 'initialize',
                    payload: {},
                });
            } catch (error) {
                this.worker = undefined;
                reject(error instanceof Error ? error : new Error(extractErrorMessage(error)));
            }
        });
    }

    private handleWorkerMessage(message: WorkerMessage): void {
        // Handle stdout/stderr messages
        if (message.type === 'stdout') {
            this.log.info({ message }, 'Pyodide stdout');
            return;
        }

        if (message.type === 'stderr') {
            this.log.error({ message }, 'Pyodide stderr');
            return;
        }

        // Handle task responses
        const id = message.id;
        if (!id) {
            return; // Ignore messages without an ID
        }

        const task = this.tasks.get(id);
        if (!task) {
            this.log.error(`Received response for unknown task: ${id}`);
            return;
        }

        this.tasks.delete(id);

        if (message.success) {
            task.resolve(message.result);
        } else {
            task.reject(new Error(message.error));
        }
    }

    public async lintTemplate(
        content: string,
        uri: string,
        fileType: CloudFormationFileType,
    ): Promise<PublishDiagnosticsParams[]> {
        return await this.executeTask<PublishDiagnosticsParams[]>('lint', {
            content,
            uri,
            fileType,
            settings: this.cfnLintSettings,
        });
    }

    public async getCfnLintVersion(): Promise<string> {
        return await this.executeTask<string>('getVersion', {});
    }

    public async lintFile(
        path: string,
        uri: string,
        fileType: CloudFormationFileType,
    ): Promise<PublishDiagnosticsParams[]> {
        return await this.executeTask<PublishDiagnosticsParams[]>('lintFile', {
            path,
            uri,
            fileType,
            settings: this.cfnLintSettings,
        });
    }

    public updateSettings(settings: CfnLintSettings): void {
        this.cfnLintSettings = settings;
    }

    public async mountFolder(fsDir: string, mountDir: string): Promise<void> {
        await this.executeTask<{ mounted: boolean; mountDir: string }>('mountFolder', { fsDir, mountDir });
    }

    private async executeTask<T>(action: string, payload: Record<string, unknown>): Promise<T> {
        if (!this.initialized) {
            await this.initialize();
        }

        if (!this.worker) {
            throw new WorkerNotInitializedError();
        }

        // Track queue depth
        this.telemetry.countUpDown('worker.queue.depth', this.tasks.size, { unit: '1' });

        const startTime = performance.now();

        return await new Promise<T>((resolve, reject) => {
            const taskId = this.nextTaskId.toString();
            this.nextTaskId++;

            const task: WorkerTask = {
                id: taskId,
                action,
                payload,
                resolve: (result: unknown) => {
                    this.telemetry.histogram('worker.response.time', performance.now() - startTime, { unit: 'ms' });
                    resolve(result as T);
                },
                reject: (error: Error) => {
                    this.telemetry.histogram('worker.response.time', performance.now() - startTime, { unit: 'ms' });
                    reject(error);
                },
            };

            this.tasks.set(taskId, task);
            if (this.worker) {
                this.worker.postMessage({ id: taskId, action, payload });
            } else {
                reject(new WorkerNotInitializedError());
            }
        });
    }

    public async shutdown(): Promise<void> {
        if (this.worker) {
            // Reject all pending tasks
            for (const task of this.tasks.values()) {
                task.reject(new Error('Worker shutdown'));
            }
            this.tasks.clear();

            // Terminate worker and wait for completion
            try {
                await this.worker.terminate();
            } catch (error) {
                this.log.error(error, 'Error terminating worker');
            }
            this.worker = undefined;
            this.initialized = false;
            this.initializationPromise = undefined;
        }
    }
}