/**
 * Copyright Microsoft Corporation. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
import colors from 'colors/safe';
import { ManualPromise } from '@isomorphic/manualPromise';
import { eventsHelper } from '@utils/eventsHelper';
import { addSuggestedRebaseline } from './rebase';
import { WorkerHost } from './workerHost';
import { ipc } from '../common';
import { addLocationAndSnippetToError } from '../reporters/internalReporter';
import { serializeError } from '../util';
export class Dispatcher {
    // Worker slot is claimed when it has jobDispatcher assigned.
    _workerSlots = [];
    _queue = [];
    _workerLimitPerProjectId = new Map();
    _queuedOrRunningHashCount = new Map();
    _finished = new ManualPromise();
    _isStopped = true;
    _testRun;
    _extraEnvByProjectId = new Map();
    _producedEnvByProjectId = new Map();
    constructor(testRun) {
        this._testRun = testRun;
        for (const project of testRun.config.projects) {
            if (project.workers)
                this._workerLimitPerProjectId.set(project.id, project.workers);
        }
    }
    _findFirstJobToRun() {
        // Always pick the first job that can be run while respecting the project worker limit.
        for (let index = 0; index < this._queue.length; index++) {
            const job = this._queue[index];
            const projectIdWorkerLimit = this._workerLimitPerProjectId.get(job.projectId);
            if (!projectIdWorkerLimit)
                return index;
            const runningWorkersWithSameProjectId = this._workerSlots.filter(w => w.jobDispatcher?.job.projectId === job.projectId).length;
            if (runningWorkersWithSameProjectId < projectIdWorkerLimit)
                return index;
        }
        return -1;
    }
    _scheduleJob() {
        // NOTE: keep this method synchronous for easier reasoning.
        // 0. No more running jobs after stop.
        if (this._isStopped)
            return;
        // 1. Find a job to run.
        const jobIndex = this._findFirstJobToRun();
        if (jobIndex === -1)
            return;
        const job = this._queue[jobIndex];
        // 2. Find a worker with the same hash, or just some free worker.
        let workerIndex = this._workerSlots.findIndex(w => !w.jobDispatcher && w.worker && w.worker.hash() === job.workerHash && !w.worker.didSendStop());
        if (workerIndex === -1)
            workerIndex = this._workerSlots.findIndex(w => !w.jobDispatcher);
        if (workerIndex === -1) {
            // No workers available, bail out.
            return;
        }
        // 3. Claim both the job and the worker slot.
        this._queue.splice(jobIndex, 1);
        const jobDispatcher = new JobDispatcher(job, this._testRun, () => this.stop().catch(() => { }));
        this._workerSlots[workerIndex].jobDispatcher = jobDispatcher;
        // 4. Run the job. This is the only async operation.
        void this._runJobInWorker(workerIndex, jobDispatcher).then(() => {
            // 5. Release the worker slot.
            this._workerSlots[workerIndex].jobDispatcher = undefined;
            // 6. Check whether we are done or should schedule another job.
            this._checkFinished();
            this._scheduleJob();
        });
    }
    async _runJobInWorker(index, jobDispatcher) {
        const job = jobDispatcher.job;
        // 0. Perhaps skip the whole job?
        if (jobDispatcher.skipWholeJob())
            return;
        let worker = this._workerSlots[index].worker;
        // 1. Restart the worker if it has the wrong hash or is being stopped already.
        if (worker && (worker.hash() !== job.workerHash || worker.didSendStop())) {
            await worker.stop();
            worker = undefined;
            if (this._isStopped) // Check stopped signal after async hop.
                return;
        }
        // 2. Start the worker if it is down.
        let startError;
        if (!worker) {
            worker = this._createWorker(job, index, ipc.serializeConfig(this._testRun.config, true));
            this._workerSlots[index].worker = worker;
            worker.on('exit', () => this._workerSlots[index].worker = undefined);
            startError = await worker.start();
            if (this._isStopped) // Check stopped signal after async hop.
                return;
        }
        // 3. Finally, run some tests in the worker! Or fail all of them because of startup error...
        if (startError)
            jobDispatcher.onExit(startError);
        else
            jobDispatcher.runInWorker(worker);
        const result = await jobDispatcher.jobResult;
        this._updateCounterForWorkerHash(job.workerHash, -1);
        // 4. When worker encounters error, we stop it and create a new one.
        //    We also do not keep the worker alive if it cannot serve any more jobs.
        if (result.didFail)
            void worker.stop(true /* didFail */);
        else if (this._isWorkerRedundant(worker))
            void worker.stop();
        // 5. Possibly queue a new job with leftover tests and/or retries.
        if (!this._isStopped && result.newJob) {
            this._queue.unshift(result.newJob);
            this._updateCounterForWorkerHash(result.newJob.workerHash, +1);
        }
    }
    _checkFinished() {
        if (this._finished.isDone())
            return;
        // Check that we have no more work to do.
        if (this._queue.length && !this._isStopped)
            return;
        // Make sure all workers have finished the current job.
        if (this._workerSlots.some(w => !!w.jobDispatcher))
            return;
        this._finished.resolve();
    }
    _isWorkerRedundant(worker) {
        let workersWithSameHash = 0;
        for (const slot of this._workerSlots) {
            if (slot.worker && !slot.worker.didSendStop() && slot.worker.hash() === worker.hash())
                workersWithSameHash++;
        }
        return workersWithSameHash > this._queuedOrRunningHashCount.get(worker.hash());
    }
    _updateCounterForWorkerHash(hash, delta) {
        this._queuedOrRunningHashCount.set(hash, delta + (this._queuedOrRunningHashCount.get(hash) || 0));
    }
    async run(testGroups, extraEnvByProjectId) {
        this._extraEnvByProjectId = extraEnvByProjectId;
        this._queue = testGroups;
        for (const group of testGroups)
            this._updateCounterForWorkerHash(group.workerHash, +1);
        this._isStopped = false;
        this._workerSlots = [];
        // 0. Stop right away if we have reached max failures.
        if (this._testRun.hasReachedMaxFailures())
            void this.stop();
        // 1. Allocate workers.
        for (let i = 0; i < this._testRun.config.config.workers; i++)
            this._workerSlots.push({});
        // 2. Schedule enough jobs.
        for (let i = 0; i < this._workerSlots.length; i++)
            this._scheduleJob();
        this._checkFinished();
        // 3. More jobs are scheduled when the worker becomes free.
        // 4. Wait for all jobs to finish.
        await this._finished;
    }
    _createWorker(testGroup, parallelIndex, loaderData) {
        const project = this._testRun.config.projects.find(p => p.id === testGroup.projectId);
        const pauseAtEnd = this._testRun.topLevelProjects.includes(project) && !!this._testRun.options.pauseAtEnd;
        const worker = new WorkerHost(testGroup, {
            parallelIndex,
            config: loaderData,
            extraEnv: this._extraEnvByProjectId.get(testGroup.projectId) || {},
            outputDir: project.project.outputDir,
            pauseOnError: !!this._testRun.options.pauseOnError,
            pauseAtEnd,
        });
        const handleOutput = (params) => {
            const chunk = chunkFromParams(params);
            if (worker.didFail()) {
                // Note: we keep reading stdio from workers that are currently stopping after failure,
                // to debug teardown issues. However, we avoid spoiling the test result from
                // the next retry.
                return { chunk };
            }
            const currentlyRunning = this._workerSlots[parallelIndex].jobDispatcher?.currentlyRunning();
            if (!currentlyRunning)
                return { chunk };
            return { chunk, test: currentlyRunning.test, result: currentlyRunning.result };
        };
        worker.on('stdOut', (params) => {
            const { chunk, test, result } = handleOutput(params);
            result?.stdout.push(chunk);
            this._testRun.reporter.onStdOut?.(chunk, test, result);
        });
        worker.on('stdErr', (params) => {
            const { chunk, test, result } = handleOutput(params);
            result?.stderr.push(chunk);
            this._testRun.reporter.onStdErr?.(chunk, test, result);
        });
        worker.on('teardownErrors', (params) => {
            this._testRun.hasWorkerErrors = true;
            const workerInfo = {
                config: this._testRun.config.config,
                project: project.project,
                workerIndex: worker.workerIndex,
                parallelIndex: worker.parallelIndex,
            };
            for (const error of params.fatalErrors)
                this._testRun.reporter.onError?.(error, workerInfo);
        });
        worker.on('exit', () => {
            const producedEnv = this._producedEnvByProjectId.get(testGroup.projectId) || {};
            this._producedEnvByProjectId.set(testGroup.projectId, { ...producedEnv, ...worker.producedEnv() });
        });
        return worker;
    }
    producedEnvByProjectId() {
        return this._producedEnvByProjectId;
    }
    async stop() {
        if (this._isStopped)
            return;
        this._isStopped = true;
        await Promise.all(this._workerSlots.map(({ worker }) => worker?.stop()));
        this._checkFinished();
    }
}
class JobDispatcher {
    jobResult = new ManualPromise();
    job;
    _testRun;
    _stopCallback;
    _listeners = [];
    _failedTests = new Set();
    _failedWithNonRetriableError = new Set();
    _remainingByTestId = new Map();
    _dataByTestId = new Map();
    _parallelIndex = 0;
    _workerIndex = 0;
    _currentlyRunning;
    constructor(job, testRun, stopCallback) {
        this.job = job;
        this._testRun = testRun;
        this._stopCallback = stopCallback;
        this._remainingByTestId = new Map(this.job.tests.map(e => [e.id, e]));
    }
    _onTestBegin(params) {
        const test = this._remainingByTestId.get(params.testId);
        if (!test) {
            // TODO: this should never be the case, report an internal error?
            return;
        }
        const result = test._appendTestResult();
        this._dataByTestId.set(test.id, { test, result, steps: new Map() });
        result.parallelIndex = this._parallelIndex;
        result.workerIndex = this._workerIndex;
        result.startTime = new Date(params.startWallTime);
        this._testRun.reporter.onTestBegin?.(test, result);
        this._currentlyRunning = { test, result };
    }
    _onTestEnd(params) {
        if (this._testRun.hasReachedMaxFailures()) {
            // Do not show more than one error to avoid confusion, but report
            // as interrupted to indicate that we did actually start the test.
            params.status = 'interrupted';
            params.errors = [];
        }
        const data = this._dataByTestId.get(params.testId);
        if (!data) {
            // TODO: this should never be the case, report an internal error?
            return;
        }
        this._dataByTestId.delete(params.testId);
        this._remainingByTestId.delete(params.testId);
        const { result, test } = data;
        result.duration = params.duration;
        result.errors = params.errors;
        result.error = result.errors[0];
        result.status = params.status;
        result.annotations = params.annotations;
        test.annotations = [...params.annotations]; // last test result wins
        test.expectedStatus = params.expectedStatus;
        test.timeout = params.timeout;
        const isFailure = result.status !== 'skipped' && result.status !== test.expectedStatus;
        if (isFailure)
            this._failedTests.add(test);
        if (params.hasNonRetriableError)
            this._addNonretriableTestAndSerialModeParents(test);
        this._reportTestEnd(test, result);
        this._currentlyRunning = undefined;
    }
    _addNonretriableTestAndSerialModeParents(test) {
        this._failedWithNonRetriableError.add(test);
        for (let parent = test.parent; parent; parent = parent.parent) {
            if (parent._parallelMode === 'serial')
                this._failedWithNonRetriableError.add(parent);
        }
    }
    _onStepBegin(params) {
        const data = this._dataByTestId.get(params.testId);
        if (!data) {
            // The test has finished, but steps are still coming. Just ignore them.
            return;
        }
        const { result, steps, test } = data;
        const parentStep = params.parentStepId ? steps.get(params.parentStepId) : undefined;
        const step = {
            title: params.title,
            titlePath: () => {
                const parentPath = parentStep?.titlePath() || [];
                return [...parentPath, params.title];
            },
            parent: parentStep,
            category: params.category,
            startTime: new Date(params.wallTime),
            duration: -1,
            steps: [],
            attachments: [],
            annotations: [],
            location: params.location,
        };
        steps.set(params.stepId, step);
        (parentStep || result).steps.push(step);
        this._testRun.reporter.onStepBegin?.(test, result, step);
    }
    _onStepEnd(params) {
        const data = this._dataByTestId.get(params.testId);
        if (!data) {
            // The test has finished, but steps are still coming. Just ignore them.
            return;
        }
        const { result, steps, test } = data;
        const step = steps.get(params.stepId);
        if (!step) {
            this._testRun.reporter.onStdErr?.('Internal error: step end without step begin: ' + params.stepId, test, result);
            return;
        }
        step.duration = params.wallTime - step.startTime.getTime();
        if (params.error)
            step.error = params.error;
        if (params.suggestedRebaseline)
            addSuggestedRebaseline(step.location, params.suggestedRebaseline);
        step.annotations = params.annotations;
        steps.delete(params.stepId);
        this._testRun.reporter.onStepEnd?.(test, result, step);
    }
    _onAttach(params) {
        const data = this._dataByTestId.get(params.testId);
        if (!data) {
            // The test has finished, but attachments are still coming. Just ignore them.
            return;
        }
        const attachment = {
            name: params.name,
            path: params.path,
            contentType: params.contentType,
            body: params.body !== undefined ? Buffer.from(params.body, 'base64') : undefined
        };
        data.result.attachments.push(attachment);
        if (params.stepId) {
            const step = data.steps.get(params.stepId);
            if (step)
                step.attachments.push(attachment);
            else
                this._testRun.reporter.onStdErr?.('Internal error: step id not found: ' + params.stepId);
        }
    }
    _failTestWithErrors(test, errors) {
        const runData = this._dataByTestId.get(test.id);
        // There might be a single test that has started but has not finished yet.
        let result;
        if (runData) {
            result = runData.result;
        }
        else {
            result = test._appendTestResult();
            this._testRun.reporter.onTestBegin?.(test, result);
        }
        result.errors = [...errors];
        result.error = result.errors[0];
        result.status = errors.length ? 'failed' : 'skipped';
        this._reportTestEnd(test, result);
        this._failedTests.add(test);
    }
    _massSkipTestsFromRemaining(testIds, errors) {
        for (const test of this._remainingByTestId.values()) {
            if (!testIds.has(test.id))
                continue;
            if (!this._testRun.hasReachedMaxFailures()) {
                this._failTestWithErrors(test, errors);
                errors = []; // Only report errors for the first test.
            }
            this._remainingByTestId.delete(test.id);
        }
        if (errors.length) {
            // We had fatal errors after all tests have passed - most likely in some teardown.
            // Let's just fail the test run.
            this._testRun.hasWorkerErrors = true;
            for (const error of errors)
                this._testRun.reporter.onError?.(error);
        }
    }
    _onDone(params) {
        // We won't file remaining if:
        // - there are no remaining
        // - we are here not because something failed
        // - no unrecoverable worker error
        if (!this._remainingByTestId.size && !this._failedTests.size && !params.fatalErrors.length && !params.skipTestsDueToSetupFailure.length && !params.fatalUnknownTestIds && !params.unexpectedExitError && !params.stoppedDueToUnhandledErrorInTestFail) {
            this._finished({ didFail: false });
            return;
        }
        for (const testId of params.fatalUnknownTestIds || []) {
            const test = this._remainingByTestId.get(testId);
            if (test) {
                this._remainingByTestId.delete(testId);
                this._failTestWithErrors(test, [{ message: `Test not found in the worker process. Make sure test title does not change.` }]);
            }
        }
        if (params.fatalErrors.length) {
            // In case of fatal errors, report first remaining test as failing with these errors,
            // and all others as skipped.
            this._massSkipTestsFromRemaining(new Set(this._remainingByTestId.keys()), params.fatalErrors);
        }
        // Handle tests that should be skipped because of the setup failure.
        this._massSkipTestsFromRemaining(new Set(params.skipTestsDueToSetupFailure), []);
        if (params.unexpectedExitError) {
            // When worker exits during a test, we blame the test itself.
            //
            // The most common situation when worker exits while not running a test is:
            //   worker failed to require the test file (at the start) because of an exception in one of imports.
            // In this case, "skip" all remaining tests, to avoid running into the same exception over and over.
            if (this._currentlyRunning)
                this._massSkipTestsFromRemaining(new Set([this._currentlyRunning.test.id]), [params.unexpectedExitError]);
            else
                this._massSkipTestsFromRemaining(new Set(this._remainingByTestId.keys()), [params.unexpectedExitError]);
        }
        const retryCandidates = new Set();
        const serialSuitesWithFailures = new Set();
        for (const failedTest of this._failedTests) {
            if (this._failedWithNonRetriableError.has(failedTest))
                continue;
            retryCandidates.add(failedTest);
            let outermostSerialSuite;
            for (let parent = failedTest.parent; parent; parent = parent.parent) {
                if (parent._parallelMode === 'serial')
                    outermostSerialSuite = parent;
            }
            if (outermostSerialSuite && !this._failedWithNonRetriableError.has(outermostSerialSuite))
                serialSuitesWithFailures.add(outermostSerialSuite);
        }
        // If we have failed tests that belong to a serial suite,
        // we should skip all future tests from the same serial suite.
        const testsBelongingToSomeSerialSuiteWithFailures = [...this._remainingByTestId.values()].filter(test => {
            let parent = test.parent;
            while (parent && !serialSuitesWithFailures.has(parent))
                parent = parent.parent;
            return !!parent;
        });
        this._massSkipTestsFromRemaining(new Set(testsBelongingToSomeSerialSuiteWithFailures.map(test => test.id)), []);
        for (const serialSuite of serialSuitesWithFailures) {
            // Add all tests from failed serial suites for possible retry.
            // These will only be retried together, because they have the same
            // "retries" setting and the same number of previous runs.
            serialSuite.allTests().forEach(test => retryCandidates.add(test));
        }
        const remaining = [...this._remainingByTestId.values()];
        for (const test of retryCandidates) {
            if (test.results.length < test.retries + 1)
                remaining.push(test);
        }
        // This job is over, we will schedule another one.
        const newJob = remaining.length ? { ...this.job, tests: remaining } : undefined;
        this._finished({ didFail: true, newJob });
    }
    onExit(data) {
        const unexpectedExitError = data.unexpectedly ? {
            message: `Error: worker process exited unexpectedly (code=${data.code}, signal=${data.signal})`
        } : undefined;
        this._onDone({ skipTestsDueToSetupFailure: [], fatalErrors: [], unexpectedExitError });
    }
    _finished(result) {
        eventsHelper.removeEventListeners(this._listeners);
        this.jobResult.resolve(result);
    }
    runInWorker(worker) {
        this._parallelIndex = worker.parallelIndex;
        this._workerIndex = worker.workerIndex;
        const runPayload = {
            file: this.job.requireFile,
            entries: this.job.tests.map(test => {
                return { testId: test.id, retry: test.results.length };
            }),
        };
        worker.runTestGroup(runPayload);
        this._listeners = [
            eventsHelper.addEventListener(worker, 'testBegin', this._onTestBegin.bind(this)),
            eventsHelper.addEventListener(worker, 'testEnd', this._onTestEnd.bind(this)),
            eventsHelper.addEventListener(worker, 'stepBegin', this._onStepBegin.bind(this)),
            eventsHelper.addEventListener(worker, 'stepEnd', this._onStepEnd.bind(this)),
            eventsHelper.addEventListener(worker, 'attach', this._onAttach.bind(this)),
            eventsHelper.addEventListener(worker, 'testPaused', this._onTestPaused.bind(this, worker)),
            eventsHelper.addEventListener(worker, 'done', this._onDone.bind(this)),
            eventsHelper.addEventListener(worker, 'exit', this.onExit.bind(this)),
        ];
    }
    _onTestPaused(worker, params) {
        const data = this._dataByTestId.get(params.testId);
        if (!data)
            return;
        const { result, test } = data;
        const sendMessage = async (message) => {
            try {
                if (this.jobResult.isDone())
                    throw new Error('Test has already stopped');
                const response = await worker.sendCustomMessage({ testId: test.id, request: message.request });
                if (response.error)
                    addLocationAndSnippetToError(this._testRun.config.config, response.error);
                return response;
            }
            catch (e) {
                const error = serializeError(e);
                addLocationAndSnippetToError(this._testRun.config.config, error);
                return { response: undefined, error };
            }
        };
        result.status = params.status;
        result.errors = params.errors;
        result.error = result.errors[0];
        void this._testRun.reporter.onTestPaused?.(test, result).then(() => {
            worker.sendResume({});
        });
        this._testRun.onTestPaused({ ...params, sendMessage });
    }
    skipWholeJob() {
        // If all the tests in a group are skipped, we report them immediately
        // without sending anything to a worker. This avoids creating unnecessary worker processes.
        //
        // However, if there is at least one non-skipped test in a group, we'll send
        // the whole group to the worker process and report tests in the natural order,
        // with skipped tests mixed in-between non-skipped. This makes
        // for a better reporter experience.
        const allTestsSkipped = this.job.tests.every(test => test.expectedStatus === 'skipped');
        if (allTestsSkipped && !this._testRun.hasReachedMaxFailures()) {
            for (const test of this.job.tests) {
                const result = test._appendTestResult();
                this._testRun.reporter.onTestBegin?.(test, result);
                result.status = 'skipped';
                // This must mirror _onTestEnd() above
                result.annotations = [...test.annotations];
                this._reportTestEnd(test, result);
            }
            return true;
        }
        return false;
    }
    currentlyRunning() {
        return this._currentlyRunning;
    }
    _reportTestEnd(test, result) {
        this._testRun.reporter.onTestEnd?.(test, result);
        const hadMaxFailures = this._testRun.hasReachedMaxFailures();
        // Test is considered failing after the last retry.
        if (test.outcome() === 'unexpected' && test.results.length > test.retries)
            ++this._testRun.failedTestCount;
        if (!hadMaxFailures && this._testRun.hasReachedMaxFailures()) {
            this._stopCallback();
            this._testRun.reporter.onError?.({ message: colors.red(`Testing stopped early after ${this._testRun.config.config.maxFailures} maximum allowed failures.`) });
        }
    }
}
function chunkFromParams(params) {
    if (typeof params.text === 'string')
        return params.text;
    return Buffer.from(params.buffer, 'base64');
}
//# sourceMappingURL=dispatcher.js.map