function [OK, out] = remoteBlockExecution( action, varargin )
%remoteBlockExecution - this function controls the SPMD block execution on the labs

% Copyright 2008-2017 The MathWorks, Inc.
    
    persistent EXECUTION_HANDLES BLOCK_COMM_HANDLE PREVIOUS_COMM_HANDLE FREE_COMM_ON_CLEANUP PRELUDE_EXCEPTION
    % EXECUTION_HANDLES    - stores the deserialized 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 = [];

    iJavaLog( 6, 'Entering remoteBlockExecution with action "%s".', action );
    logOnExit = onCleanup( @() iJavaLog( 6, 'Leaving remoteBlockExecution with action "%s".', action ) );
    
    try
        switch action
          case 'prelude'
            % Grab the world labindex before we go any further.
            pctPreRemoteEvaluation( 'mwmpi' );
            PREVIOUS_COMM_HANDLE = mpiCommManip( 'select', 'world' );
            worldLabIndex = labindex;
            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( worldLabIndex, varargin{:} );
            catch E
                % Unexpected error in prelude. We have no way to continue, so abort.
                dctSchedulerMessage( 0, 'Unexpected error during SPMD prelude, aborting: %s', E.identifier )
                mpigateway( 'abort' );
            end
                                    
          case 'interruptibleExecution'
            % In this case, we expect interruption, and so do not rely on modifying any
            % of the persistent data.
            if isempty( PRELUDE_EXCEPTION )
                bodyFcn = EXECUTION_HANDLES{1};

                feval( '_workspace_transparency', 1 );
                bodyFcn();
                feval( '_workspace_transparency', 0 );

                % Always call setidle here - if we've got this far, we know it's safe, and
                % it just might speed an error-exit.
                iSetIdle(BLOCK_COMM_HANDLE);
            else
                % Always call setidle here - if we've got this far, we know it's safe, and
                % it just might speed an error-exit.
                iSetIdle(BLOCK_COMM_HANDLE);
                % 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 of error or
            % interruption during the body.
            iSetIdle(BLOCK_COMM_HANDLE);

            if isempty( PRELUDE_EXCEPTION )
                getOutFcn = EXECUTION_HANDLES{2};
                out = iPackOutputs( getOutFcn() );
            else
                % No block execution - bad setup
                out = distcompserialize( [] );
            end
            % Always perform deadlock detection
            iDeadlockDetection( BLOCK_COMM_HANDLE, PREVIOUS_COMM_HANDLE, FREE_COMM_ON_CLEANUP );
          case 'doFreeComm'
            FREE_COMM_ON_CLEANUP = true;
        end
        OK = true;
    catch E
        OK = false;
        % This is where we are actually evaluating the user code, so the 
        % error that we serialize back needs to include the stack info
        % so we can stitch together a new exception later.
        stackFilesToIgnore = {'LocalSpmdExecutor', 'remoteBlockExecution'};
        errToSerialize = ParallelException.hBuildFromRemoteException( E, stackFilesToIgnore );
        out = distcompserialize( errToSerialize );
    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( worldLabIndex, serCommHandleCell, ...
                                                  serializedFcnHandles, arrayHolderMap )

    % Clear any keys that are no longer in use. We need to find out our 'worker
    % index' rather than our world labindex to handle the differences between
    % an interactive and a batch mode pool.
    [~, sessionObj]   = parallel.internal.pool.PoolArrayManager.getCurrent();
    workerKey = java.lang.Integer( sessionObj.WorkerIndex );
    if arrayHolderMap.containsKey( workerKey )
        pClearIfNecessary( arrayHolderMap.get( workerKey ) );
    end

    % Unpack the data we've been given
    [commHCell, fcnHandles, preludeExcept] = ...
        iDeserializeInputs( serCommHandleCell, serializedFcnHandles );
    
    % Set up the communicators
    commHandle = iSelectCommunicator( worldLabIndex, 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:SourceCodeNotAvailableOnLab', labindex));
        else
            rethrow(E);
        end
    end
    
end

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

    commCellValid = iscell( commHCell ) && length( commHCell ) >= worldLabIndex;
    
    commSet = false;
    if commCellValid
        myComm = commHCell{worldLabIndex};
        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


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% iDeserializeInputs - unpack everything sent over, and return along with a
% flag indicating whether any Remote type objects were sent.
function [commHCell, fcnHandles, deserExc] = iDeserializeInputs( serCommHandleCell, serializedFcnHandles )
    state = [];
    % Since we're returning an exception, give these default invalid values.
    commHCell = []; fcnHandles = [];
    try
        commHCell = parallel.internal.pool.deserialize( serCommHandleCell );

        spmdlang.AbstractRemote.saveLoadCount( 'clear' );

        fcnHandles = ...
            parallel.internal.pool.deserializeFunction( serializedFcnHandles );

        E = [];
    catch E
        % E will be rethrown later.
    end
    
    if ~isempty( state )
        warning( state );
    end

    % Return the error if we've got one
    deserExc = E;
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% iPackOutputs - pack up and return the body outputs
function out = 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
    out = distcompserialize( outCellParcel );
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% iDeadlockDetection - perform deadlock detection postlude stage
function iDeadlockDetection( blockComm, prevComm, freeBlockComm )
    try
        mpiCommManip( 'select', blockComm );
        mpigateway( 'setidle' );
        mpigateway( 'setrunning' );
    catch E %#ok
            % Ignore
    end
    
    % Unconditionally reset parfor_depth, as we do for PMODE (in
    % distcomp.pInterPPromptFcn)
    parfor_depth( 0 );
    
    mpiCommManip( 'select', prevComm );
    
    if freeBlockComm
        mpiCommManip( 'free', blockComm );
    end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% iJavaLog - log to the pmode logger. Note we do not use dctSchedulerMessage
% here as that prints output to the command window.
function iJavaLog( level, fmt, varargin )
try
    loglvl = com.mathworks.toolbox.parallel.pctutil.logging.DistcompLevel.getLevelFromValue( level );
    LOGGER = com.mathworks.toolbox.distcomp.pmode.PackageInfo.LOGGER;
    LOGGER.log( loglvl, sprintf( fmt, varargin{:} ) );
catch E %#ok<NASGU>
    % ignore any logging problems
end
end
