function [OK, out] = remoteBlockExecutionPlain(action, varargin)
%remoteBlockExecutionPlain controls the SPMD block execution on the workers

% Copyright 2020-2023 The MathWorks, Inc.

    persistent EXECUTION_HANDLES BLOCK_COMM_HANDLE PREVIOUS_COMM_HANDLE FREE_COMM_ON_CLEANUP PRELUDE_EXCEPTION
    % EXECUTION_HANDLES    - stores the function handles sent across from the client
    % BLOCK_COMM_HANDLE    - the communicator handle used in the context of this SPMD block
    % PREVIOUS_COMM_HANDLE - the communicator handle in use outside this SPMD context
    % FREE_COMM_ON_CLEANUP - boolean flag to indicate whether or not we should free the block's communicator
    % PRELUDE_EXCEPTION    - MException thrown during prelude phase, if any.

    mlock;

    out = [];

    logger = parallel.internal.logging.LogGuard('spmd action="%s"', action);

    try
        switch action
          case 'prelude'
            % Grab the world spmdIndex before we go any further.
            try
                % Catch problems from pctPreRemoteEvaluation, will use as PRELUDE_EXCEPTION, but
                % only after we've extracted the communicators.
                pctPreRemoteEvaluation('mwmpi');
                preEvalException = [];
            catch preEvalException
                logger.logError(preEvalException);
            end
            PREVIOUS_COMM_HANDLE = mpiCommManip('select', 'world');
            worldSpmdIndex = spmdIndex;
            mpiCommManip('select', PREVIOUS_COMM_HANDLE);

            % Set up defaults in case iPrelude errors out
            FREE_COMM_ON_CLEANUP = false;
            BLOCK_COMM_HANDLE = [];
            EXECUTION_HANDLES = [];
            PRELUDE_EXCEPTION = [];

            try
                [EXECUTION_HANDLES, BLOCK_COMM_HANDLE, PRELUDE_EXCEPTION] = iPrelude(worldSpmdIndex, varargin{:});
            catch E
                % Unexpected error in prelude. We have no way to continue, so abort.
                logger.logError(E, 'Unexpected error during SPMD prelude, aborting.');
                mpigateway('abort');
            end
            if ~isempty(preEvalException)
                PRELUDE_EXCEPTION = preEvalException;
            end

          case 'interruptibleExecution'
            % In this case, we expect interruption, and so do not rely on modifying any
            % of the persistent data.

            % Use onCleanup to ensure we always call setidle.
            ensureSetIdle = onCleanup(@() iSetIdle(BLOCK_COMM_HANDLE));

            if isempty(PRELUDE_EXCEPTION)
                bodyFcn = EXECUTION_HANDLES{1};
                bodyFcn();
            else
                % Don't even try to execute the block
                throw(PRELUDE_EXCEPTION);
            end
          case 'postlude'
            % Always call setidle here - we may not have hit it because interruptibleExecution might
            % never have run.
            iSetIdle(BLOCK_COMM_HANDLE);

            if isempty(PRELUDE_EXCEPTION)
                getOutFcn = EXECUTION_HANDLES{2};
                out = iPackOutputs(getOutFcn());
            else
                % No block execution - bad setup
                out = [];
            end
            % Done with EXECUTION_HANDLES - clear it now because it can cause
            % data and Composites to stay in memory (g2420794).
            EXECUTION_HANDLES = [];

            % Always perform deadlock detection
            iDeadlockDetection(BLOCK_COMM_HANDLE, PREVIOUS_COMM_HANDLE, FREE_COMM_ON_CLEANUP);

            % Need to flush all in-flight data queue messages before exiting the
            % spmd block to ensure consistent order of operations.
            workerSession = parallel.internal.pool.workerSession();
            if ~isempty(workerSession)
                cppWorkerSession = workerSession.Session;
                cppWorkerSession.ensureAllDataQueueMessagesDelivered();
            end
          case 'doFreeComm'
            FREE_COMM_ON_CLEANUP = true;
        end
        OK = true;
    catch E
        logger.logError(E, 'Error during execution.');
        E = parallel.internal.pool.processUserCodeException(E);
        OK = false;
        % This is where we are actually evaluating the user code, so the
        % error that we send back needs to include the stack info
        % so we can stitch together a new exception later.
        stackFilesToIgnore = {'LocalSpmdExecutor', 'remoteBlockExecutionPlain'};
        out = ParallelException.hBuildFromRemoteException(E, stackFilesToIgnore);
    end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% call "mpigateway('setidle')", while guaranteeing to do that only on the block
% communicators.
function iSetIdle(blockCommHandle)
currentCommHandle = mpiCommManip('queryState');
x = onCleanup(@() mpiCommManip('select', currentCommHandle));
mpiCommManip('select', blockCommHandle);
mpigateway('setidle');
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% iPrelude: guarantees to set commHandle - may set fcnHandles if preludeExcept
% is not empty. If something is wrong with commHandle, will abort.
function [fcnHandles, commHandle, preludeExcept] = iPrelude(worldSpmdIndex, commHCell, fcnHandles, keysToClear, preludeExcept)
    if nargin < 5
        % No deserialisation exception from wrapping layer
        preludeExcept = [];
    end

    % Clear any keys that are no longer in use.
    pClearIfNecessary(keysToClear);

    % Set up the communicators
    commHandle = iSelectCommunicator(worldSpmdIndex, commHCell);

    if ~isempty(preludeExcept)
        return;
    end
    
    % Unpack the inputs for later
    try
        unpackInFcn = fcnHandles{3};
        unpackInFcn(@spmdlang.AbstractSpmdExecutor.unpack);
    catch E
        if strcmp(E.identifier, 'MATLAB:UndefinedFunction')
            preludeExcept = MException(message('parallel:lang:spmd:SourceCodeNotAvailableOnWorker', spmdIndex));
        else
            rethrow(E);
        end
    end

end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% iSelectCommunicator - deal with setting the communicator
function myComm = iSelectCommunicator(worldSpmdIndex, commHCell)

    commCellValid = iscell(commHCell) && length(commHCell) >= worldSpmdIndex;

    commSet = false;
    if commCellValid
        myComm = commHCell{worldSpmdIndex};
        if mpiCommManip('isValid', myComm)
            mpiCommManip('select', myComm);
            commSet = true;
        end
    end
    if ~commSet
        % This ought to be a fatal error, as we can't recover (OTOH, I've never hit
        % this...)
        dctSchedulerMessage(0, 'Unexpectedly failed to get MPI communicators, aborting');
        mpigateway('abort');
    end
end


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% iPackOutputs - pack up and return the body outputs
function outCellParcel = iPackOutputs(outCell)
    outCellParcel = cell(1, length(outCell));
    for ii=1:length(outCell)
        try
            if isempty(outCell{ii})
                % Do nothing - cell already empty
            else
                data = outCell{ii}{1};
                key = spmdlang.ValueStore.store(data);
                [fcn, data] = getRemoteFromSPMD(data);
                outCellParcel{ii} = spmdlang.ReturnableParcel(key, fcn, data);
            end
        catch E
            warning(message('parallel:lang:spmd:SpmdOutputs', getReport(E)));
        end
    end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% iDeadlockDetection - perform deadlock detection postlude stage
function iDeadlockDetection(blockComm, prevComm, freeBlockComm)
% NB: onCleanup inside 'interruptibleExecution' ensures we are in setidle mode here.
    try
        mpiCommManip('select', blockComm);
        mpigateway('setrunning');
    catch E %#ok
            % Ignore
    end

    % Unconditionally reset parfor_depth
    parfor_depth(0);

    mpiCommManip('select', prevComm);

    if freeBlockComm
        mpiCommManip('free', blockComm);
    end
end
