function canRunCommunicatingJobs = evaluateMJSTask(serializedArgs)
% Evaluates the MJS task stored in the serialized args and submits its
% result to the job manager

% Copyright 2020-2021 The MathWorks, Inc.

    taskExecutionInfo = iGetTaskExecutionInfo(serializedArgs);
    postFcns = iSetup(taskExecutionInfo);
    iEvaluateTask(postFcns);

    % This MATLAB process can only run subsequent communicating jobs if MPI has
    % never been initialized
    canRunCommunicatingJobs = ~feval('_mpiEverInitialized');
    dctSchedulerMessage(2, 'End evaluateMJSTask');
end

%--------------------------------------------------------------------------
% Deserializes the task execution info from its serialized form
%--------------------------------------------------------------------------
function taskExecutionInfo = iGetTaskExecutionInfo(serializedArgs)
    serializedArgs = typecast(serializedArgs, 'uint8'); % Ensure uint8 input
    taskExecutionInfo = parallel.internal.cluster.deserializeTaskExecutionInfo(serializedArgs);
    
    % Convert job and task IDs to their appropriate types
    taskExecutionInfo.jobIDNumAndType = parallel.internal.cluster.MJSClientPropertiesAdaptor.adaptJobIDNumAndType(...
        parallel.internal.cluster.MJSSchedulerPropertiesAdaptor.adaptJobIDNumAndType(taskExecutionInfo.jobIDNumAndType));
    taskExecutionInfo.taskIDAndNum = parallel.internal.cluster.MJSClientPropertiesAdaptor.adaptTaskIDAndNum(...
        parallel.internal.cluster.MJSSchedulerPropertiesAdaptor.adaptTaskIDAndNum(taskExecutionInfo.taskIDAndNum));
    taskExecutionInfo.previousTaskIDAndNum = parallel.internal.cluster.MJSClientPropertiesAdaptor.adaptTaskIDAndNum(...
        parallel.internal.cluster.MJSSchedulerPropertiesAdaptor.adaptTaskIDAndNum(taskExecutionInfo.previousTaskIDAndNum));
    
    % Convert environment variables to map
    entries = taskExecutionInfo.environmentVariables.values;
    taskExecutionInfo.environmentVariables = containers.Map('KeyType', 'char', 'ValueType', 'char');
    for ii = 1:numel(entries)
        entry = entries(ii);
        taskExecutionInfo.environmentVariables(entry.key) = entry.value;
    end
    
    % Captured serialised authentication token bytes
    taskExecutionInfo.authToken = taskExecutionInfo.authToken.hashedCredentials;
    
    % Convert portSelectionInformation to its corresponding MATLAB type
    info = taskExecutionInfo.portSelectionInformation;
    taskExecutionInfo.portSelectionInformation = parallel.internal.cluster.PortSelectionInformation(...
        info.indexOnHost, info.numWorkersOnHost, info.minPortOffset);
    
    % Convert cached work unit properties to maps
    jobData = containers.Map();
    jobData('AdditionalPaths') = taskExecutionInfo.cachedJobData.additionalPaths;
    jobData('ApiTag') = taskExecutionInfo.cachedJobData.apiTag;
    jobData = iCacheOptionalData(jobData, 'AttachedFileData', taskExecutionInfo.cachedJobData.attachedFileData, @uint8, uint8.empty());
    jobData = iCacheOptionalSerialisedData(jobData, 'AttachedFilePaths', taskExecutionInfo.cachedJobData.attachedFilePaths, []);
    jobData('AutoAddClientPath') = taskExecutionInfo.cachedJobData.autoAddClientPath;
    jobData('ExecutionMode') = taskExecutionInfo.cachedJobData.executionMode;
    jobData('EnvironmentVariables') = taskExecutionInfo.cachedJobData.environmentVariables;
    jobData('LeadingTaskNum') = taskExecutionInfo.cachedJobData.leadingTaskNum;
    jobData = iCacheOptionalSerialisedData(jobData, 'MATLABDrivePathMap', taskExecutionInfo.cachedJobData.matlabDrivePathMap, []);
    jobData('MATLABRelease') = taskExecutionInfo.cachedJobData.matlabRelease;
    jobData('NumLabs') = taskExecutionInfo.cachedJobData.numLabs;
    jobData('NumPathsToAppend') = taskExecutionInfo.cachedJobData.numPathsToAppend;
    jobData('ProductList') = taskExecutionInfo.cachedJobData.productKeys;
    taskExecutionInfo = rmfield(taskExecutionInfo, 'cachedJobData');
    
    taskData = containers.Map();
    taskData = iCacheOptionalSerialisedData(taskData, 'Function', taskExecutionInfo.cachedTaskData.function, @parallel.internal.pool.poolWorkerFcn);
    taskData = iCacheOptionalSerialisedData(taskData, 'InputArguments', taskExecutionInfo.cachedTaskData.inputArguments, {});
    taskData('NumOutputArguments') = taskExecutionInfo.cachedTaskData.numOutputArguments;
    taskExecutionInfo = rmfield(taskExecutionInfo, 'cachedTaskData');
    
    clusterData = containers.Map();
    clusterData('AdminUsername') = taskExecutionInfo.cachedClusterData.adminUser;
    clusterData('AllHostAddresses') = taskExecutionInfo.cachedClusterData.allHostAddresses;
    clusterData('AllowClientPasswordCache') = taskExecutionInfo.cachedClusterData.allowClientPasswordCache;
    clusterData('AuthorisationLevel') = parallel.internal.cluster.types.AuthorizationLevel.fromOrdinal(taskExecutionInfo.cachedClusterData.authorizationLevel);
    clusterData('HasSecureCommunication') = taskExecutionInfo.cachedClusterData.hasSecureCommunication;
    clusterData('HostName') = taskExecutionInfo.cachedClusterData.hostname;
    clusterData('IsRunAsUser') = taskExecutionInfo.cachedClusterData.runAsUser;
    clusterData('LookupURL') = taskExecutionInfo.cachedClusterData.lookupURL;
    clusterData('MpiLibraryName') = taskExecutionInfo.cachedClusterData.mpiLibraryName;
    clusterData('Name') = taskExecutionInfo.cachedClusterData.name;
    clusterData('RequiresOnlineLicensing') = taskExecutionInfo.cachedClusterData.requiresOnlineLicensing;
    clusterData('SecurityLevel') = taskExecutionInfo.cachedClusterData.securityLevel;
    clusterData('WorkerProxiesPoolConnections') = taskExecutionInfo.cachedClusterData.workerProxiesPoolConnections;
    
    taskExecutionInfo.cachedData.jobID = taskExecutionInfo.jobIDNumAndType.jobID;
    taskExecutionInfo.cachedData.jobProperties = jobData;
    taskExecutionInfo.cachedData.taskID = taskExecutionInfo.taskIDAndNum.taskID;
    taskExecutionInfo.cachedData.taskProperties = taskData;
    taskExecutionInfo.cachedData.clusterProperties = clusterData;
end

function map = iCacheOptionalData(map, propertyName, optionalData, prepareDataFcn, emptyValue)
    if optionalData.isPresent
        if isempty(optionalData.data)
            map(propertyName) = emptyValue;
        else
            map(propertyName) = prepareDataFcn(optionalData.data);
        end
    end
end

function map = iCacheOptionalSerialisedData(map, propertyName, optionalData, emptyValue)
    createDeferredDeserialisationFcn = @(data) parallel.internal.cluster.DeferredEvaluation(@() distcompdeserialize(data));
    map = iCacheOptionalData(map, propertyName, optionalData, createDeferredDeserialisationFcn, emptyValue);
end

%--------------------------------------------------------------------------
% Prepares worker, job, and task for evaluation and returns a collection of
% post-task cleanup functions to run
%--------------------------------------------------------------------------
function postFcns = iSetup(taskExecutionInfo)
    postFcns = {};
    
    %% Perform user initialization (should be done first!)
    if taskExecutionInfo.performUserInit
        preMJSTask(taskExecutionInfo.username);
        postFcns = [{{@postMJSTask}}, postFcns];
    end
    
    %% Perform worker setup
    [justStarted, resetWorkerFcn] = iSetupWorker(taskExecutionInfo);
    dctSchedulerMessage(2, 'Begin evaluateMJSTask');

    %% Perform job setup
    if justStarted
        taskExecutionInfo.performJobInit = true; % Always perform job initialization after a worker start
    end
    
    if taskExecutionInfo.performJobInit
        dctSchedulerMessage(4, 'Performing job initialization');
        resetWorkerFcn();
        
        % Clean out files from all previous jobs by deleting the file
        % dependencies directory
        if exist(taskExecutionInfo.fileDependencyDir, 'dir')
            try
                dctSchedulerMessage(4, 'Removing old dependency directory %s', taskExecutionInfo.fileDependencyDir);
                rmdir(taskExecutionInfo.fileDependencyDir, 's');
            catch E
                % Don't worry too much if we can't do this, it is actually
                % quite likely to happen on Windows if MATLAB hasn't
                % restarted. At some point, we will be able to delete this
                % directory.
                dctSchedulerMessage(6, 'Removing dependency directory failed: %s', E.message);
            end
        end

        % Make sure the worker directory exists
        if ~exist(taskExecutionInfo.workerDir, 'dir')
            dctSchedulerMessage(4, 'Making new worker directory %s', taskExecutionInfo.workerDir);
            mkdir(taskExecutionInfo.workerDir);
        end
    end

    % Ensure that any changes made by mpiprofile to general profiler
    % settings have been reverted
    try
        dctSchedulerMessage(4, 'Resetting mpiprofile');
        mpiprofile('reset');
    catch E
        dctSchedulerMessage(1, 'Warning - failed to reset profiler state. Error thrown was:\n%s\n', E.message);
    end

    % We wish to ensure that the current working directory of MATLAB is not
    % a UNC path on Windows and that it is likely to be writable by this
    % process, so change directory to the local work directory
    dctSchedulerMessage(4, 'Changing directory to %s', taskExecutionInfo.workerDir);
    cd(taskExecutionInfo.workerDir);

    % Create the appropriate type of desktop client
    webLicenseInfo = taskExecutionInfo.webLicenseInfo;
    usingMHLM = ~isempty(webLicenseInfo.webUserCredentials);
    if usingMHLM
        desktopClient = parallel.internal.mwa.MDCSDesktopClient.buildWithFixedToken(webLicenseInfo.webUserCredentials, webLicenseInfo.webUserName);
    else
        desktopClient = parallel.internal.mwa.MDCSDesktopClient();
    end
    parallel.internal.mwa.currentDesktopClient(desktopClient);

    % Set the number of computational threads
    dctSchedulerMessage(4, 'Setting number of computational threads to %d.', taskExecutionInfo.numThreads);
    try
        feature('numthreads', taskExecutionInfo.numThreads);
    catch E
        dctSchedulerMessage(1, sprintf('Warning - failed to set number of computational threads to %d. Error thrown was:', taskExecutionInfo.numThreads), E);
    end

    %% Create the job manager
    % Build a new MJS cluster object. Since each cluster object stores its
    % own cache of jobs and tasks, there is no need to clear previous jobs.
    dctSchedulerMessage(4, 'Creating job manager');
    jobManager = parallel.cluster.MJS('Name', taskExecutionInfo.jobManagerName, ...
        'Host', taskExecutionInfo.connectUrl, ...
        'Username', taskExecutionInfo.username, ...
        'Credentials', taskExecutionInfo.authToken, ...
        'CachedData', taskExecutionInfo.cachedData);
    if usingMHLM
        entitlement = parallel.internal.mwa.Entitlement(webLicenseInfo.licenseID, webLicenseInfo.licenseNumber, '');
        jobManager.hSetLicenseEntitlementNoCheck(entitlement);
    end

    %% Create the worker
    dctSchedulerMessage(4, 'Creating worker');
    workerProperties = taskExecutionInfo.workerProperties;
    worker = parallel.cluster.MJSWorker(jobManager, ...
        workerProperties.hostName, ...
        workerProperties.computerMLType, ...
        workerProperties.allHostAddresses, ...
        workerProperties.name);
    
    %% Create the job and task
    dctSchedulerMessage(4, 'Creating job and task');
    [job, task] = jobManager.hGetJobAndTaskFromStructs(taskExecutionInfo.jobIDNumAndType, taskExecutionInfo.taskIDAndNum);
    task.hSetAttempt(taskExecutionInfo.taskAttempt);
    [~, previousTask] = jobManager.hGetJobAndTaskFromStructs(taskExecutionInfo.jobIDNumAndType, taskExecutionInfo.previousTaskIDAndNum);
    previousTask.hSetAttempt(taskExecutionInfo.taskAttempt);
    task.hSetPreviousTask(previousTask);

    %% Create the runprop data
    runProp = distcomp.runprop();
    runProp.AppendPathDependencies = taskExecutionInfo.performJobInit; % Add PathDependencies if job init is required
    runProp.AppendFileDependencies = taskExecutionInfo.performJobInit; % Add FileDependencies if job init is required
    runProp.IsFirstTask = taskExecutionInfo.performJobInit; % Call jobStartup() if job init is required
    runProp.MATLABDriveEnabled = strcmpi(getenv('MDCS_MATLAB_DRIVE_ENABLED_ON_WORKER'), 'true');
    dctSchedulerMessage(4, 'MATLAB Drive Enabled %d', runProp.MATLABDriveEnabled);
    runProp.PortSelectionInformation = taskExecutionInfo.portSelectionInformation;
    runProp.EnvironmentVariablesMap = taskExecutionInfo.environmentVariables;
    runProp.DependencyDirectory = fullfile(taskExecutionInfo.fileDependencyDir, num2str(taskExecutionInfo.jobIDNumAndType.num));
    runProp.TaskEvaluator = parallel.internal.evaluator.BasicEvaluator();

    % Make sure to reuse the AttachedFilesPathMap if we're not performing
    % job initialization. See g1565643.
    if ~taskExecutionInfo.performJobInit
        oldRunProp = distcomp.getdistcompobjectroot().CurrentRunprop;
        runProp.AttachedFilesPathMap = oldRunProp.AttachedFilesPathMap;
    end

    %% Create the handlers
    handlers = struct('errorFcn', @iError);

    %% Set all properties on the root object
    dctSchedulerMessage(4, 'Setting current task info');
    root = distcomp.getdistcompobjectroot();
    root.setCurrentTaskInfo(jobManager, worker, job, task, runProp, handlers);
    
    %% Create the post task functions
    % Define any post processing tasks that we want to evaluate - these
    % will be run when the finish function is called

    % Clean up dependency directory only if the job is a concurrent type
    if isa(job, 'parallel.ConcurrentJob') || isa(job, 'parallel.CommunicatingJob')
        postFcns = [{{@removeDirectory, runProp.DependencyDirectory}}, postFcns];
    end
    
    % Always clean up the per-host directory
    postFcns = [{{@removeDirectory, runProp.PerHostDirectory}}, postFcns];

    % Always clean up parallel pool state if applicable.
    postFcns = [{{@iCleanupPoolState}}, postFcns];
end

%--------------------------------------------------------------------------
% Initializes the worker and returns true if it has just started, otherwise
% returns a function that resets the worker to its initial state.
%
% NB: If the worker has just started, the returned function is a no-op.
%--------------------------------------------------------------------------
function [justStarted, resetWorkerFcn] = iSetupWorker(taskExecutionInfo)
    mlock(); % Prevents persistent variables from being cleared
    persistent resetFcn
    justStarted = isempty(resetFcn);
    
    % Always set up logging
    iSetupLogging(justStarted, taskExecutionInfo);
    
    % Quick exit if the worker hasn't just started
    if ~justStarted
        resetWorkerFcn = resetFcn;
        return
    end
    
    dctSchedulerMessage(2, 'Start initworker');

    % MATLAB has been launched with -nodesktop, yet the diary output is
    % likely to be viewed on a desktop client. We set the hotlinks feature
    % on here for backwards compatibility as tasks used to be evaluated
    % with evalc, which always turns on hotlinks.
    feature('hotlinks', 'on');

    % Turn off beep for workers
    beep('off');

    % Make sure this MATLAB cannot core dump, and will exit if anything
    % goes wrong
    dct_psfcns('ensureProcessExitsOnFault');
    
    % PCT_Java_Transition_TODO (g2233336): Handle SSL certificate
    % For now we install any java certificate to allow java pool
    % communication if required
    encodedCertificate = getenv('WORKER_CERTIFICATE');
    if ~isempty( encodedCertificate )
        dctSchedulerMessage(2, 'Adding certificate: %s', encodedCertificate);
        certStore = com.mathworks.toolbox.distcomp.mjs.security.ClientCertificateStore.getInstance();
        certStore.setCertificateEntry('SharedSecret', encodedCertificate);
    end

    % Store the original path and warning state for resetting the worker
    originalPath = path();
    originalWarningState = warning('query', 'all');
    resetFcn = @() iResetWorker(originalPath, originalWarningState);
    
    dctSchedulerMessage(2, 'End initworker');
    
    resetWorkerFcn = @distcomp.nop; % No reset is necessary, worker just started
end

%--------------------------------------------------------------------------
% Sets up logging for this MATLAB worker
%--------------------------------------------------------------------------
function iSetupLogging(justStarted, taskExecutionInfo)
    if justStarted
        logFilePath = getenv('WORKER_LOG_FILE_PATH');
        maxLogFileSizeBytes = str2double(getenv('WORKER_MAX_LOG_FILE_SIZE_BYTES'));
        maxNumLogFiles = str2double(getenv('WORKER_MAX_NUM_LOG_FILES'));
        parallel.internal.logging.initLogging(logFilePath, maxLogFileSizeBytes, maxNumLogFiles);
    end
    
    % Always set the log level/spec and scheduler message handler in case
    % they have been changed by the previous task
    if isempty(taskExecutionInfo.logSpec)
        parallel.internal.logging.setLogLevel(taskExecutionInfo.logLevel);
    else
        parallel.internal.logging.setLogSpec(taskExecutionInfo.logSpec); % Overrides log level
    end
    setSchedulerMessageHandler(@(msg, levelNum) parallel.internal.logging.log(levelNum, msg));
end

%--------------------------------------------------------------------------
% Helper function to restore path and warning state
%--------------------------------------------------------------------------
function iResetWorker(originalPath, originalWarningState)
    dctSchedulerMessage(4, 'Restoring original path and warning state');
    path(originalPath);
    warning(originalWarningState);
end

%--------------------------------------------------------------------------
% Helper function for erroring out of MATLAB
%--------------------------------------------------------------------------
function iError(err, msg)
    % The task errored and the integrity of the running MATLAB can no
    % longer be assumed. Log the error and notify the worker by rethrowing
    % the error. The worker will terminate this process.
    dctSchedulerMessage(0, msg, err);
    rethrow(err);
end

%--------------------------------------------------------------------------
% Helper function for cleaning pool state
%--------------------------------------------------------------------------
function iCleanupPoolState()
    try
        delete(gcp('nocreate'));
    catch E
        dctSchedulerMessage(2, 'Error cleaning up pool state', E);
    end
end

%--------------------------------------------------------------------------
% Evaluates the task and executes post-task functions after task evaluation
%--------------------------------------------------------------------------
function iEvaluateTask(postFcns)
    dctSchedulerMessage(4, 'About to evaluate task');
    try
        dctEvaluateTask(postFcns, @iFinishTask);
    catch E
        if isa(E, 'distcomp.ExitException')
            % Extract and rethrow the more useful underlying error
            rethrow(E.CauseException);
        else
            % Otherwise just rethrow the original error
            rethrow(E);
        end
    end
    dctSchedulerMessage(4, 'Finished evaluating task');
end

%--------------------------------------------------------------------------
% Submits the task result and performs post-task work
%--------------------------------------------------------------------------
function iFinishTask(job, task, out)
    % NB: Normally, we call pPostJobEvaluate() before submitting a task's
    % result so that MPI-based jobs can wait for all tasks to reach the end
    % before allowing any task to submit its result (see comment below).
    % However, to ensure that MPI-based jobs finish quickly if one of their
    % tasks errors, we submit the result of such a task *first*. Doing so
    % makes the task 'finished' and hence uncancellable, though, so to
    % ensure that the task doesn't get stuck in pPostJobEvaluate() we
    % forcibly exit MATLAB.
    taskErrored = ~isempty(out.errOutput);
    if taskErrored && mpiInitialized()
        dctSchedulerMessage(4, 'About to submit task result of a parallel task that errored in user code');
        task.hSubmitResult(out.output, out.errOutput, out.warnOutput, out.textOutput);
        
        dctSchedulerMessage(0, 'About to exit MATLAB because a parallel task errored');
        exit('force');
    end

    % NB: It is important that we perform the post job evaluation *before*
    % submitting the task result because MPI-based jobs (i.e. all
    % communicating jobs and some concurrent jobs) rely on the blocking
    % behaviour in pPostJobEvaluate() to ensure that all tasks remain
    % running and hence cancellable until all tasks reach the end.
    % Otherwise, if some but not all tasks error, the tasks that do *not*
    % error will submit their results, making them no longer cancellable,
    % and then block in the mpiCleanup() in pPostJobEvaluate() (possibly
    % forever) waiting for the erroring tasks to reach the same point,
    % which they may not if MATLAB crashed for those erroring tasks.
    dctSchedulerMessage(4, 'About to pPostJobEvaluate');
    job.pPostJobEvaluate();

    dctSchedulerMessage(4, 'About to submit task result');
    task.hSubmitResult(out.output, out.errOutput, out.warnOutput, out.textOutput);

    pctSetmcrappkeys('', '');
    % Reset the CWD-on-path state if necessary.
    parallel.internal.general.revertCwdOnPath();
end
