% Remote resource set object - gateway to remote SPMD execution. Manages a
% cell array of MPI communicator handles, and a corresponding array of
% process objects and can build the appropriate SPMD block executor


% Copyright 2008-2020 The MathWorks, Inc.

classdef RemoteResourceSet < spmdlang.AbstractResourceSet

    properties (SetAccess = private, Transient)
        % Cell array of remote communicator handles. (Might be character vectors or uint64s).
        RemoteCommCell;
    end
    
    properties (Access = private, Transient)
        % Am I the current "world"
        IsWorld = false;
        
        % Parallel pool
        Pool = parallel.Pool.empty();
    end

    properties (GetAccess = private, Constant)
        COMPOSITE_WARNING_CLASSES = {'distributed', 'Composite'};
    end
    
    methods (Access = public, Static)

        % This function exists mainly to support parallel pool jobs, where the effective
        % world isn't 'world'. 
        function x = getWorldCommsCell(nlabs)
            [wc, labIdFirst] = spmdlang.commForWorld('get');
            
            % Need to pad the cell array so that the labs can index into it by their
            % "World" labindex.
            padCells = cell(1, labIdFirst - 1);
            
            if ischar(wc) % probably 'world'
                x = repmat({wc}, 1, nlabs);
            else
                x = wc;
            end
            x = [padCells, x]; % in parallel pool job, this is e.g.: {[], 6, 6, 6}
        end
        
        % Bootstrap - this static method builds the "world" set from the pool, and
        % requires no communication or remote execution to achieve that.
        function worldSet = buildWorldSet(pool)
            if isempty(pool)
                error(message('parallel:lang:spmd:NoSession'));
            end
            ps = pool.NumWorkers;
            worldSet = spmdlang.RemoteResourceSet(...
                ps, spmdlang.RemoteResourceSet.getWorldCommsCell(ps), pool);
            
            worldSet.IsWorld = true;
        end
    
    end
    
    methods (Access = public, Hidden)
        function obj = RemoteResourceSet(numProcesses, remoteCommCell, pool)
            validateattributes(numProcesses, {'numeric'}, {'integer', 'nonnegative'});
            
            obj = obj@spmdlang.AbstractResourceSet(numProcesses);
            obj.RemoteCommCell = remoteCommCell;
            obj.Pool = pool;
        end
        
        % "Factory method" to build a remote resource set of a specific size from the
        % open pool. It does that by splitting the base set. 
        % "sz" is simply the required size of the subset
        function subSetObj = buildRemoteResourceSet(obj, sz)
            
            if sz == obj.numlabs
                subSetObj = obj;
            elseif sz < 0
                error(message('parallel:lang:spmd:InvalidResourceSelection', sz));
            elseif sz < obj.numlabs

                % Build from the base set, 
                newComms = spmd_feval_fcn(@iBuildCommHandle, { sz }, obj);

                % Pad as necessary to handle mismatch between MPI 'world' labindex and
                % processes. (Happens in parallel pool jobs)
                [~, labIdFirst] = spmdlang.commForWorld('get');
                padCells = cell(1, labIdFirst - 1);
                newComms = {padCells{:}, newComms{:}}; %#ok<CCAT> Dereference the Composite returned from the spmd_feval_fcn
                
                subSetObj = spmdlang.RemoteResourceSet(sz, newComms(:), obj.Pool);
            else
                % Should never get here because there's already a check that will not let us
                % build a resource set greater than the size of the base set
                error(message('parallel:lang:spmd:InvalidResourceSelection', sz));
            end
        end

        
        function tf = isValid(obj)
            pool = obj.Pool;
            tf = ~isempty(pool) && isvalid(pool) && pool.Connected;
        end
        
        function tf = canAccessLabs(obj)
            % Need to test if we are allowed to access the lab and ask it to do an
            % operation - all we need to check is whether the resource set is
            % valid.
            tf = obj.isValid();    
        end
        
        function tf = satisfiesConstraints(obj, minN, maxN)
            obj.errorIfInvalid();
            tf = (obj.numlabs <= maxN && obj.numlabs >= minN);
        end
        
        function val = getFromLab(obj, labidx, key)
        % Talk to the lab, get the value
            obj.errorIfInvalid();
            assistant = obj.getCompositeAssistant(labidx);
            
            % Make sure we can detect transfer of Remote objects
            % Serialize the initialization data and store it locally
            [listener, getListFcn] = ...
                parallel.internal.general.SerializationNotifier.createAccumulatingListener(); %#ok<ASGLU>

            val = remoteRetrieval(obj, assistant, labidx, key);

            serialisationClasses = getListFcn();
            if ~isempty(serialisationClasses)
                if any(ismember(serialisationClasses, obj.COMPOSITE_WARNING_CLASSES))
                    warning(message('parallel:lang:spmd:RemoteObjectRetrievedFromLabs'));
                end
            end
        end
        
        function newKey = setOnLab(obj, labidx, newValue)
        % Talk to the lab, send the value, retrieve the key
            obj.errorIfInvalid();
            assistant = obj.getCompositeAssistant(labidx);

            % Make sure we can detect transfer of Remote objects
            [listener, getListFcn] = ...
                parallel.internal.general.SerializationNotifier.createAccumulatingListener(); %#ok<ASGLU>

            newKey = remoteSet(obj, assistant, labidx, newValue);

            serialisationClasses = getListFcn();
            if ~isempty(serialisationClasses)
                if any(ismember(serialisationClasses, obj.COMPOSITE_WARNING_CLASSES))
                    warning(message('parallel:lang:spmd:RemoteObjectSentToLabs'));
                end
            end
        end
        
        function keyUnreferenced(obj, labidx, key)
        % Push the key onto the stack of things to clear for the particular lab -
        % don't worry if we're not valid (that means the lab has probably
        % gone away)
            if isValid(obj)
                obj.remoteClear(labidx, key);
            end
        end
        
        function forceClear(obj)
        % Queue up all the remote clear commands to execute togther before
        % processing the results
            obj.errorIfSessionClosed();
            assistant = obj.getCompositeAssistant(1:obj.numlabs);
            futures = arrayfun(@(id) obj.doForceClear(assistant, id), 1:obj.numlabs, 'UniformOutput', false);
            breakWaitFcn = @() obj.futurePollFcn();
            for idx = 1:numel(futures)
                future = futures{idx};
                if isempty(future)
                    continue
                end
                future.await(breakWaitFcn);

                try
                    future.getResult();
                catch E
                    warning(message('parallel:lang:spmd:RemoteKeyFlush', idx, getReport(E)));
                end
            end
        end
        
        function blockEx = buildBlockExecutor(obj, bodyFcn, assignOutFcn, getOutFcn, ...
                                               unpackInFcn, initialOuts)
            obj.errorIfInvalid();
            blockEx = spmdlang.RemoteSpmdExecutor(spmdlang.ResourceSetHolder(obj), ...
                                                  bodyFcn, assignOutFcn, getOutFcn, ...
                                                  unpackInFcn, initialOuts, obj.Pool.hGetSession());
        end
        
        function invalidate(obj, labsStillAvailable)
            if ~obj.IsWorld && labsStillAvailable
                % Use SPMD itself to release the communicators
                spmd_feval_fcn(@iFreeMyComm, {}, obj);
            end
            % Finally, call the superclass invalidation
            invalidate@spmdlang.AbstractResourceSet(obj, labsStillAvailable);
        end

        function tf = shouldRetryIfFileNotFound(~)
        % Always retry the computation a file not found error is thrown by the resource set.
            tf = true;
        end
    end
    
    %% Methods to do with Remote retrieval
    methods (Access = private, Hidden)
        function compositeAssistant = getCompositeAssistant(obj, processIndices)
            % Acquire use of the specified labs for remote resource
            % management. The returned value automatically releases the
            % workers on delete.
            try
                compositeAssistant = getCompositeAssistant(obj.Pool.hGetSession(), processIndices);
            catch E
                except = MException(message('parallel:lang:spmd:RemoteSetupError'));
                except = addCause(except, E);
                throw(except);
            end
        end
        
        function errorIfSessionClosed(obj)
            if ~obj.isValid
                error(message('parallel:lang:spmd:PoolSessionHasClosed'));
            end
        end
        
        function tf = futurePollFcn(obj)
            tf = obj.isValid;
            obj.errorIfSessionClosed();
        end
        
        function val = remoteRetrieval(obj, assistant, processIndex, key)
            obj.errorIfSessionClosed();
            future = assistant.retrieveCompositeValue(processIndex,  key);
            val = obj.handleRemoteSetOrRetrievalFuture(future, ...
                                                        'parallel:lang:spmd:WarnGettingComposite', ...
                                                        'parallel:lang:spmd:RemoteRetrievalError');
        end
        
        function newKey = remoteSet(obj, assistant, processIndex, newValue)
            obj.errorIfSessionClosed();
            future = assistant.sendCompositeValue(processIndex,  newValue);           
            newKey = obj.handleRemoteSetOrRetrievalFuture(future, ...
                                                           'parallel:lang:spmd:WarnSettingComposite', ...
                                                           'parallel:lang:spmd:RemoteSet');
        end

        function remoteClear(obj, processIndex, key)
            obj.Pool.hGetSession().getCompositeKeysToClearMap().addKeyForProcess(processIndex, key);
        end
        
        function future = doForceClear(~, assistant, processIndex)
            if assistant.hasKeysToClear(processIndex)
                future = assistant.clearCompositeValues(processIndex);
            else
                future = [];
            end
        end
        
        % Poll the observer for a remoteRetrieval or remoteSet operation, handle
        % warnings, and return the result.
        function result = handleRemoteSetOrRetrievalFuture(obj, future, workerWarningId, errorId)
            breakWaitFcn = @() obj.futurePollFcn();
            future.await(breakWaitFcn);

            % Show the worker-side warnings before deserialization of the result so that
            % anything went wrong on the workers is displayed before anything that
            % might go wrong back at the client.
            warnings = future.getWarnings();
            if ~isempty(warnings)
                warningsList = parallel.internal.display.DisplayHelper.formatWarningsList(warnings);
                parallel.internal.warningNoBackTrace(message(workerWarningId, warningsList));
            end

            % Get the set/retrieval result.
            try
                result = future.getResult();
            catch E
                ME = MException(message(errorId));
                ME = addCause(ME, E);
                throwAsCaller(ME);
            end
        end
    end
    
end


% Used by RemoteResourceSet::invalidate on the labs to free a given communicator.
function iFreeMyComm()
    % Use a special back-door to free communicators - the block cleanup needs to
    % use them for error checking.
    spmdlang.remoteBlockExecution('doFreeComm');
end

% Used by RemoteResourceSet::buildRemoteResourceSet to split the world communicator.
function val = iBuildCommHandle(newSz)
    if newSz > numlabs
        % Should never get here - two layers of previous checks would have to fail
        % to catch this...
        error(message('parallel:lang:spmd:InvalidComms', newSz, numlabs));
    end
    
    % All labs <= labindex get color 1, others get color -1 (so they get COMM_NULL)
    color = (2 * double(labindex <= newSz)) - 1;
    key    = labindex;
    
    val    = mpiCommManip('split', color, key);
end



