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


% Copyright 2008-2022 The MathWorks, Inc.

classdef AbstractRemoteResourceSet < spmdlang.AbstractResourceSet

    properties (SetAccess = protected, Transient)
        % Cell array of remote communicator handles. (Might be character vectors or uint64s).
        RemoteCommCell;
    end

    properties (GetAccess = protected, SetAccess = immutable, Transient)
        % Am I the current "world"
        IsWorld (1,1) logical = false;

        % Parallel pool
        Pool parallel.Pool {mustBeScalarOrEmpty} = parallel.Pool.empty();
    end

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

    %% Abstract methods related to constructing resource sets.
    methods (Abstract, Access = public, Static)
        % This static method should build the "world" set from the pool.
        worldSet = buildWorldSet(pool)
    end

    methods (Abstract, Access = public, Hidden)
        % "Factory method" to build a remote resource set of a specific size from the
        % open pool.
        % "sz" is simply the required size of the subset
        subSetObj = buildRemoteResourceSet(obj, sz)
    end

    methods (Access = public, Hidden)
        function obj = AbstractRemoteResourceSet(numWorkers, remoteCommCell, pool, isWorld)
            validateattributes(numWorkers, {'numeric'}, {'integer', 'nonnegative'});

            obj = obj@spmdlang.AbstractResourceSet(numWorkers);
            obj.RemoteCommCell = remoteCommCell;
            obj.Pool = pool;
            obj.IsWorld = isWorld;
        end

        function tf = isValid(obj)
            pool = obj.Pool;
            tf = ~isempty(pool) && isvalid(pool) && pool.Connected;
        end

        function tf = canAccessWorkers(obj)
            % Need to test if we are allowed to access the worker 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.NumWorkers <= maxN && obj.NumWorkers >= minN);
        end

        function val = getFromWorker(obj, workerIdx, key)
            % Talk to the worker, get the value
            obj.errorIfInvalid();
            assistant = obj.getCompositeAssistant(workerIdx);

            % 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, workerIdx, key);

            serialisationClasses = getListFcn();
            if ~isempty(serialisationClasses)
                if any(ismember(serialisationClasses, obj.COMPOSITE_WARNING_CLASSES))
                    warning(message('parallel:lang:spmd:RemoteObjectRetrievedFromWorkers'));
                end
            end
        end

        function newKey = setOnWorker(obj, workerIdx, newValue)
            % Talk to the worker, send the value, retrieve the key
            obj.errorIfInvalid();
            assistant = obj.getCompositeAssistant(workerIdx);

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

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

            serialisationClasses = getListFcn();
            if ~isempty(serialisationClasses)
                if any(ismember(serialisationClasses, obj.COMPOSITE_WARNING_CLASSES))
                    warning(message('parallel:lang:spmd:RemoteObjectSentToWorkers'));
                end
            end
        end

        function keyUnreferenced(obj, workerIdx, key)
            % Push the key onto the stack of things to clear for the particular worker -
            % don't worry if we're not valid (that means the worker has probably
            % gone away)
            if isValid(obj)
                obj.remoteClear(workerIdx, 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.NumWorkers);
            futures = arrayfun(@(id) obj.doForceClear(assistant, id), 1:obj.NumWorkers, '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, workersStillAvailable)
            % Simply all the superclass invalidation
            invalidate@spmdlang.AbstractResourceSet(obj, workersStillAvailable);
        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 = protected, Hidden)
        function compositeAssistant = getCompositeAssistant(obj, processIndices)
            % Acquire use of the specified workers for remote resource
            % management. The returned value automatically releases the
            % workers on delete.
            try
                compositeAssistant = obj.Pool.hGetSession().getCompositeAssistant(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