% remoteparfor - support for parallel_function.
% PCT_JT_CLEANUP (g2863205) remove this

% Copyright 2013-2023 The MathWorks, Inc.
classdef remoteparfor < parallel.internal.parfor.Engine
    properties (GetAccess = private, Constant)
        % Avoid building empty exceptions multiple times.
        EMPTY_EXCEPTION = MException.empty(); 

        % Apply a rate-limiter to displayOutput. This value
        % is the minimum time between calls to draining the
        % underlying output. (At the end of the loop, we always
        % drain all remaining output)
        DISPLAY_RATE_LIMIT_SECONDS = 0.5; 
    end
    properties (GetAccess = private, SetAccess = immutable)
        Pool
        Session
        ParforF % parallel_function 'F' argument.
        MaxNumWorkers (1,1) int32
        InitData % In case we need to restart
    end
    properties (Access = private)
        ParforController
        IntervalCompleteQueue
        NumWorkers (1,1) double

        AllIntervalsAdded (1,1) logical = false
        NormalCompletion  (1,1) logical = false
        AddedFilesToPool  (1,1) logical = false

        % Cell array of intervals that have been submitted but are not yet
        % complete. Indexed by the 'tag' produced by parallel_function - so
        % we rely on those tags being integers starting at 1.
        IncompleteIntervalsCell
        
        % Track the number of intervals that have been submitted to the controller,
        % but not yet collected.
        NumIntervalsInController
        
        % Result of "tic()" at the last point that we actually tried
        % to display command-window output.
        LastDisplayTic (1,1) uint64 = tic();
        
        % Cache for the lifetime of this parfor loop whether or not serialization
        % profiling is enabled. If it is not, we avoid the overhead of
        % checking each time.
        IsSerializationProfilingEnabled (1,1) logical = false
    end

    methods (Static)
        function [OK, pool] = tryRemoteParfor()
        % This is a static method of distcomp.remoteparfor that indicates if it is
        % OK to try constructing a distcomp.remoteparfor object. This will stop
        % obvious errors being thrown.
            [OK, pool] = pctTryCreatePoolIfNecessary();
        end
        
        function oldVal = maxBacklog(newVal)
        % 'Backlog' refers to how many intervals the ParforController will send to the
        % workers without requiring a result.
        %
        % See the discussion in 'numIntervalsFactor' for how this factor
        % interacts with that value.
            persistent MAX_BACKLOG
            if isempty(MAX_BACKLOG)
                MAX_BACKLOG = 2;
            end
            oldVal = MAX_BACKLOG;
            if nargin > 0
                MAX_BACKLOG = newVal;
            end
        end

        function oldVal = initialDispatchFactor(newVal)
        % parallel_function.m uses getInitialDispatchSize() to choose how many intervals
        % to hand off in the first round of dispatch. That method scales the
        % number of workers by this factor to choose the initial dispatch
        % size. The default is to allow parallel_function to dispatch all
        % intervals up-front. This allows distcomp.remoteparfor to receive all
        % the intervals as soon as possible and then send them to the workers.
        %
        % This factor is provided to allow an approximation of the old behaviour
        % to be achieved; however, it is not anticipated that there is any
        % appreciable benefit to changing this value.
            persistent FACTOR
            if isempty(FACTOR)
                FACTOR = Inf;
            end
            oldVal = FACTOR;
            if nargin > 0
                FACTOR = newVal;
            end
        end
        
        function oldVal = smallInitDataBufferSize(newVal)
        % If the ParforController tells us to prefer small buffers when sending the
        % initialization data, this is our value of 'small'.
            persistent BUFFER_SIZE
            if isempty(BUFFER_SIZE)
                % Use 10MB buffers as 'small' buffers.
                BUFFER_SIZE = 10 * 1024 * 1024;
            end
            oldVal = BUFFER_SIZE;
            if nargin > 0
                BUFFER_SIZE = newVal;
            end
        end
    end
    
    methods
        function buildParforController(obj)
        % Called during construction, and also when we need to retry a loop.
            % If anything goes wrong during construction of this interface then we
            % need to clean up the parfor controller, otherwise all subsequent
            % parfor statements will revert to running locally.
            cleanupControllerBeforeAquiringLabs = parallel.internal.general.DisarmableOncleanup(...
                @() obj.Session.releaseCurrentParforController());
                     
            try
                % Make a new parfor controller - this might throw a
                % SessionDestroyedException
                doYield = true; % Allow PLR requests while waiting
                p = parallel.internal.getJavaFutureInterruptibly(...
                    obj.Session.createParforController(), doYield);
            catch err
                if ~isempty(regexp(err.message, 'SessionDestroyedException', 'once'))
                    errorMessageInput = iGetParpoolLinkForError();
                    error(message('parallel:lang:parfor:NoRunningSession', errorMessageInput));
                else
                    rethrow(err);
                end
            end
            
            % Now we have a parforcontroller we can use that to clean up
            % once the acquireWorkers method has been called on it.
            interruptControllerOnError = parallel.internal.general.DisarmableOncleanup(...
                @() p.interrupt());

            serializedInitData = [];
            try
                obj.ParforController = p;
                % We're using a new ParforController - it has no intervals
                obj.NumIntervalsInController = 0;
                % Get the IntervalReturnQueue
                obj.IntervalCompleteQueue = p.getIntervalCompleteQueue;
                % Try to acquire some workers
                obj.NumWorkers = p.acquireWorkers(obj.MaxNumWorkers, ...
                                                  int32(distcomp.remoteparfor.maxBacklog()));
                % Labs acquired, therefore disable the first oncleanup
                cleanupControllerBeforeAquiringLabs.disarm();

                % Serialize the initialization data and store it locally
                [listener, getListFcn] = ...
                    parallel.internal.general.SerializationNotifier.createAccumulatingListener(); %#ok<ASGLU>

                if p.preferSmallInitDataBuffers
                    initDataBufSize = distcomp.remoteparfor.smallInitDataBufferSize();
                else
                    % Here, '0' means unlimited.
                    initDataBufSize = 0;
                end
                serializedInitData = obj.serialize(obj.InitData, initDataBufSize);
                % Error if we're about to broadcast Composites to the labs
                savedClasses = getListFcn();
                if ~isempty(savedClasses) && any(ismember({'distributed', 'Composite'}, savedClasses))
                    error(message('parallel:lang:parfor:IllegalComposite'));
                end
                p.beginLoop(serializedInitData);
                % We've handed the serialized data to the java layer, never
                % free it from here, assume java layer will do that.
                serializedInitData = []; %#ok<NASGU> should never see this value in the 'catch' block...
            catch err
                % If we catch an error then we need to free the data we were going to
                % send to the remote labs.
                if ~isempty( serializedInitData )
                    arrayfun( @free, serializedInitData );
                end
                rethrow(err);
            end
            interruptControllerOnError.disarm();
        end

        function rebuildParforController(obj)
        % rebuildParforController - dispose of old parfor controller, build a new one, and submit any
        % intervals.
            if ~isempty(obj.ParforController)
                obj.ParforController.interruptOnError();
            end
            try
                obj.buildParforController();
            catch E
                % Can't rebuild for some reason - maybe pool has completely
                % crashed.
                if ~obj.Session.isSessionRunning
                    error(message('MATLAB:remoteparfor:AllParforWorkersAborted'));
                else
                    rethrow(E);
                end
            end
            % Re-submit pending intervals. Incomplete intervals are non-empty entries in
            % the IncompleteIntervalsCell array.
            gotIncomplete = ~cellfun(@isempty, obj.IncompleteIntervalsCell);
            tags = find(gotIncomplete);
            dataCell = obj.IncompleteIntervalsCell(gotIncomplete);
            
            % Reset IncompleteIntervalsCell
            obj.IncompleteIntervalsCell = cell(size(obj.IncompleteIntervalsCell));

            dctSchedulerMessage(2, 'Resubmitting intervals with tags: %s\n', ...
                                join(string(tags)));
            for idx = 1:numel(tags)
                OK = obj.addInterval(tags(idx), dataCell{idx}{:});
                assert(OK, ...
                       'Unexpected failure to add interval %d.', idx);
            end
            if obj.AllIntervalsAdded
                obj.ParforController.allIntervalsAdded();
            end
        end
        
        function filesAttached = addMissingFilesToPool(obj, possibleSourceFiles)
        % Before we even try, let's be sure not to do this again.
            obj.AddedFilesToPool = true;
            filesAttached = parallel.internal.pool.attachDependentFilesToPool(...
                obj.Pool, possibleSourceFiles);
        end
        
        function [r, err] = handleIntervalErrorResult(obj, r)
        % handleIntervalErrorResult - called while looping in getCompleteIntervals.
        % Returns empty in r to indicate that this error has been confined.
            
            dctSchedulerMessage(2, 'Got interval error for tag: %d', r.getTag());
            
            rebuildAndRetry = false;
            
            [err, workerAborted] = iIntervalErrorDispatch(r, @(data) deserialize(obj, data));
            if workerAborted
                dctSchedulerMessage(1, 'Worker aborted during parfor, will attempt to retry.');
                warning(message('MATLAB:remoteparfor:ParforWorkerAborted'));
                rebuildAndRetry = true;
            else
                [err, possibleSourceFiles] = ...
                    parallel.internal.parfor.maybeTransformMissingSourceException(...
                        err, obj.ParforF);
                if ~isempty(possibleSourceFiles) && ~obj.AddedFilesToPool
                    % We've added stuff to the pool, build a new controller and send the intervals.
                    rebuildAndRetry = obj.addMissingFilesToPool(possibleSourceFiles);
                end
            end
            if rebuildAndRetry
                dctSchedulerMessage(1, 'Rebuilding parfor controller for retry.');
                obj.rebuildParforController();
                r = [];
                err = distcomp.remoteparfor.EMPTY_EXCEPTION;
            end
        end
    end

    methods (Access = protected)
        function rangeDivided(obj, k)
            if isempty(obj.IncompleteIntervalsCell)
                % Pre-allocate now we know how many intervals will be produced.
                obj.IncompleteIntervalsCell = cell(1, k);
            end
        end
    end
    
    methods (Access = private)
        function data = serialize(obj, data, bufSize)
            if nargin < 3
                bufSize = [];
            end
            data = parallel.internal.pool.serialize(...
                data, bufSize, ...
                obj.IsSerializationProfilingEnabled);
        end
        function data = deserialize(obj, serData)
            data = parallel.internal.pool.deserialize(...
                serData, [], ...
                obj.IsSerializationProfilingEnabled);
        end
    end
    
    methods
        function obj = remoteparfor(partitionMethod, partitionSize, ...
                                    pool, maxLabsToAcquire, initData, parforF, numIterates)
            obj@parallel.internal.parfor.Engine(partitionMethod, partitionSize);

            % Set up the simple immutable data
            iErrorCheckPool(pool);
            obj.Pool                = pool;
            obj.Session             = pool.hGetSession();
            obj.ParforF             = parforF;
            obj.MaxNumWorkers       = min(maxLabsToAcquire, numIterates); 
            obj.InitData            = initData;

            % Check this once for the whole loop.
            [~, obj.IsSerializationProfilingEnabled] = ...
                parallel.internal.pool.serializationProfiling('isenabled');
            
            feval('_pct_parforLogJava');
            
            % Build a java ParforControllerImpl
            obj.buildParforController();
        end

        function W = getNumWorkers(obj)
            W = obj.NumWorkers;
        end
        
        function allIntervalsAdded(obj)
        % We simply need to inform the workers that they can now drop persistent data.
            obj.ParforController.allIntervalsAdded();
            obj.AllIntervalsAdded = true;
        end

        function OK = addInterval(obj, tag, varargin)
            try
                assert(numel(obj.IncompleteIntervalsCell) < tag || ...
                       isempty(obj.IncompleteIntervalsCell{tag}), ...
                       'Unexpectedly got duplicate interval tag: %d', tag);
                obj.IncompleteIntervalsCell{tag} = varargin;
                data = obj.serialize(varargin);
                addedToController = obj.ParforController.addInterval(tag, data);
                if addedToController
                    obj.NumIntervalsInController = 1 + obj.NumIntervalsInController;
                end
                % Always return 'true' here so that parallel_function continues. In the case
                % where the controller has already received an error return
                % from an interval, then getCompleteIntervals will notice that
                % error, and either confine it (in the case of a missing
                % source file), or throw it. If getCompleteIntervals can
                % confine the error, then intervals in IncompleteIntervalsCell
                % will be added to a new controller, so we must always store
                % interval requests there even if the current controller
                % rejects them.
                OK = true;
            catch E
                dctSchedulerMessage(0, 'Problem in addInterval:', E);
                OK = false;
            end
        end

        function complete(obj, errorWasThrown)
        % We get into this function both during normal parfor execution as well as when
        % the user code throws an error.  In both cases, this is our one and only chance
        % to flush partial lines that may be pending in the command window output.

            % Indicate normal completion to stop the workers from being
            % interrupted.
            dctSchedulerMessage(4, 'remoteparfor.complete(%d) called.', errorWasThrown);
                
            if errorWasThrown && ~isempty(obj.ParforController)
                obj.ParforController.interruptOnError();
            end
                
            obj.NormalCompletion = ~errorWasThrown;

            obj.awaitCompletionDisplayingOutput();
        end

        function [tags,  results] = getCompleteIntervals(obj, numIntervals)
            q = obj.IntervalCompleteQueue;
            timeUnitSeconds = iTimeUnitSeconds();
            tags = nan(numIntervals, 1);
            results = cell(numIntervals, 2);
            for i = 1:numIntervals
                r = [];
                err = [];
                while isempty(r)
                    assert(obj.NumIntervalsInController > 0, ...
                           'Internal error in PARFOR - no intervals to retrieve.');
                    r = q.poll(1, timeUnitSeconds);
                    obj.displayOutput();
                    % In each poll tick of the polling wait, yield to any
                    % events allowed to interrupt the parallel language.
                    parallel.internal.pool.yield();
                    if isempty(r)
                        % Only test to see if the session is failing if we didn't get a
                        % results from the queue
                        if ~obj.Session.isSessionRunning
                            errorMessageInput = iGetParpoolLinkForError();
                            error(message('parallel:lang:parfor:SessionShutDown', errorMessageInput));
                        end
                    else
                        obj.NumIntervalsInController = obj.NumIntervalsInController - 1;
                        if r.hasError
                            % Maybe try again
                            [r, err] = obj.handleIntervalErrorResult(r);
                            q = obj.IntervalCompleteQueue;
                        end
                    end
                end
                % Check to see if the interval result has an error
                if ~isempty(err)
                    % Java code has already interrupted the remote execution of the parfor
                    % body and halted the receipt of further IO.  Before throwing an error,
                    % we perform the normal cleanup activities.
                    obj.complete(false);
                    throw(err);
                else
                    tags(i) = r.getTag;
                    data = r.getResult;
                    thisdata = obj.deserialize(data(2));
                    assert(numel(thisdata) == 2, ...
                           'Unexpectedly received the incorrect number of outputs.');
                    results(i,:) = thisdata;
                    obj.IncompleteIntervalsCell{tags(i)} = [];
                end
            end
            % Yield one final time to ensure all events directly generated by
            % code in the PARFOR loop are executed before we give control back
            % to the user.
            parallel.internal.pool.yield();
        end
        function delete(obj)
            % Ensure that any outstanding intervals are cleaned up immediately.
            obj.drainIntervals();

            if obj.NormalCompletion || isempty(obj.ParforController)
                % No need to interrupt nor drain output
                return;
            end
            
            % Interrupt the remote execution and display all the command window output that
            % we receive.

            % Send a Ctrl+C to remote end.
            if ~obj.NormalCompletion
                obj.ParforController.interruptOnError;
            else
                % Do note if the controller is in state 'complete' interrupt does
                % nothing so it is safe to call after the controller is complete.
                obj.ParforController.interrupt;
            end
            obj.awaitCompletionDisplayingOutput();
        end

        function [init, remain] = getDispatchSizes(~, numIntervals, W)
            init = iClamp(distcomp.remoteparfor.initialDispatchFactor() * W, ...
                        1, numIntervals);
            remain = 2 * W;
        end
    end
    
    methods (Access = private)
        
        function drainIntervals(obj)
        % Ensure there are no outstanding intervals
            timeUnitSeconds = iTimeUnitSeconds();
            try
                q = obj.IntervalCompleteQueue;
                r = q.poll(0, timeUnitSeconds);
                while ~isempty(r)
                    if ~r.hasError
                        data = r.getResult;
                        bbhArray = data(2);
                        com.mathworks.toolbox.distcomp.util.ByteBufferHandle.freeBuffers(bbhArray);
                    end
                    r = q.poll(0, timeUnitSeconds);
                end
            catch E
                dctSchedulerMessage(1, 'A problem was encountered freeing return data:', E);
            end
        end
        
        function displayOutput(obj)
        % Rate-limited so that we do not drain output more than once every
        % DISPLAY_RATE_LIMIT_SECONDS seconds.
            if toc(obj.LastDisplayTic) > ...
                    distcomp.remoteparfor.DISPLAY_RATE_LIMIT_SECONDS
                iDisplayStringArray(...
                    obj.ParforController.getDrainableOutput().drainOutput());
                obj.LastDisplayTic = tic();
            end
        end
        function awaitCompletionDisplayingOutput(obj)
        % Make sure we wait for the parfor to complete and continue displaying command
        % window output whilst we are doing so.
            output = obj.ParforController.getDrainableOutput();
            
            if isempty(output)
                % Can get here if no labs were acquired before the
                % ParforController was deleted. If this is the case there
                % is obviously no output to drain.
                return;
            end
            
            millis = java.util.concurrent.TimeUnit.MILLISECONDS;
            while ~obj.ParforController.awaitCompleted(100, millis) ...
                    &&  obj.Session.isSessionRunning
                dctSchedulerMessage(5, 'remoteparfor.awaitCompletionDisplayingOutput() still thinks not completed');
                iDisplayStringArray(output.drainOutput());
            end
            iDisplayStringArray(output.drainAllOutput());
        end
    end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% iDisplayStringArray - display a java.lang.String[] in the command window.
function iDisplayStringArray(msgs)
    % Tiny performance improvement for the common case where 'msgs' is empty - avoid
    % conversion of empty array etc.
    if ~isempty(msgs)
        cellfun(@(msg) fprintf(1, '%s', char(msg)), cell(msgs));
    end
end

function errorMessageInput = iGetParpoolLinkForError()
if feature('hotlinks')
    errorMessageInput = '<a href="matlab: helpPopup parpool">parpool</a>';
else
    errorMessageInput = 'parpool';
end
end

function val = iClamp(inVal, minBound, maxBound)
val = min(max(inVal, minBound), maxBound);
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Return an appropriate error from a failing interval result.
function [err, workerAborted] = iIntervalErrorDispatch(r, deserFcn)
    intervalError = r.getError;
    workerAborted = r.isWorkerAbortedError;
    if workerAborted
        err = MException(message('parallel:lang:pool:WorkerAborted'));
    elseif ischar(intervalError)
        err = MException(message('parallel:lang:pool:UnexpectedParforFailure', ...
                                 intervalError));
    else
        [isMvmErr, mvmErr] = ...
            parallel.internal.general.transformIfMvmException(intervalError);
        if isMvmErr
            origErr = mvmErr;
            % intervalError must have been a Java Exception - log it.
            parallel.internal.general.logJavaException(...
                0, 'Failed to get result from PARFOR.', intervalError);
        else
            origErr = deserFcn(intervalError);
        end
        clientStackToIgnore = 'parallel_function';
        % Create a new exception that stitches together the client and worker stack.
        err = ParallelException.hBuildRemoteParallelException(...
            origErr, clientStackToIgnore);
    end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function iErrorCheckPool(pool)
    if isempty(pool) || ~pool.Connected
        errorMessageInput = iGetParpoolLinkForError();
        error(message('parallel:lang:parfor:NoSession', errorMessageInput));
    end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function t = iTimeUnitSeconds()
    persistent TIME_UNIT_SECONDS
    if isempty(TIME_UNIT_SECONDS)
        TIME_UNIT_SECONDS = java.util.concurrent.TimeUnit.SECONDS;
    end
    t = TIME_UNIT_SECONDS;
end
