% RemoteSpmdExecutor - control class to manage remote execution of an SPMD
% block This class manages the M side of dealing with the Java controller,
% and in particular attempts to deal correctly with M-side interrupts.

% Copyright 2008-2020 The MathWorks, Inc.

classdef RemoteSpmdExecutor < spmdlang.AbstractSpmdExecutor

    properties ( Access = private, Hidden, Transient )
        % The java controller
        RemoteSpmdController;
        % An array of processes
        RemoteProcessArray;
        % A cell array of communicators for those processes
        RemoteCommCell;
        % The java session object
        Session;
        % Flag to indicate whether we have sent the block for execution
        SentBlockToLabs = false;
        % Stash how many output arguments to expect (can't access superclass data in destructor)
        SpmdNargout = 0;

        % Have we warned about the interrupt yet?
        WarnedAboutError = false;
        % Time at which we first spotted there was an error - we don't warn for a while
        TimeErrorFirstSeen = [];
    end
    
    properties ( GetAccess = private, Constant )
        % This is the timeout after noticing that a lab has encountered an error
        % (and therefore sending an interrupt) before we warn the users.
        TIMEOUT_BEFORE_WARNING_ABOUT_ERROR = 5; % seconds
    end

    methods ( Access = private, Hidden )
        function errorIfNoSession( rse )
        % Method used to bail out of various situations if we discover that the
        % session has shut down.
            if ~rse.Session.isSessionRunning()
                except = ParallelException.hBuildLikeMException( message('parallel:lang:spmd:SessionNotRunning') );
                if rse.SentBlockToLabs
                    % Tell the RemoteSpmdController that we're out of here. This
                    % is required so that we can extract information about the
                    % block results, even though the block is not necessarily
                    % complete.
                    rse.RemoteSpmdController.dispose();
                    rse.Disposed = true;
                    excepts      = rse.makeBlockExceptionCauses();
                    if ~isempty( excepts )
                        if numel( excepts ) == 1
                            except = ParallelException.hCloneWithNewMessage( excepts{1}, ...
                                'parallel:lang:spmd:SessionNotRunningSingleWorker', excepts{1}.message );
                        else
                            for ii = 1:length( excepts )
                                except = except.addCause( excepts{ii} );
                            end
                        end
                    end
                end
                throw(except);
            end
        end
    end

    methods

        function obj = RemoteSpmdExecutor( resSetH, bodyF, assignOutF, getOutF, unpackInF, initialOuts, session )

        % assert that we've got a remote resource set
            resSet = resSetH.getResourceSet();
            if ~isa( resSet, 'spmdlang.RemoteResourceSet' )
                error(message('parallel:lang:spmd:RemoteInconsistency'));
            end

            % Build the parent class - manages the closures/
            obj = obj@spmdlang.AbstractSpmdExecutor( resSetH, bodyF, assignOutF, ....
                                                     getOutF, unpackInF, initialOuts );

            % Build the java controller which will actually do the work
            obj.RemoteSpmdController = buildController( resSet.RemoteProcessArray, session );
            obj.RemoteProcessArray   = resSet.RemoteProcessArray;
            obj.RemoteCommCell       = resSet.RemoteCommCell;

            % stash the session
            obj.Session              = session;
            
            obj.SpmdNargout          = length( initialOuts );
        end


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

            % Serialize communicator handles and function handles for transmission
            comms = parallel.internal.pool.serialize( obj.RemoteCommCell );
            
            warnStates = spmdlang.AbstractSpmdExecutor.disableRemoteSaveWarnings();
            try
                fcns  = parallel.internal.pool.serialize( ...
                    { obj.BodyFcn, obj.GetOutFcn, obj.UnpackInFcn } );
            catch E
                spmdlang.AbstractSpmdExecutor.restoreRemoteSaveWarnings( warnStates );
                rethrow( E );
            end
            spmdlang.AbstractSpmdExecutor.restoreRemoteSaveWarnings( warnStates );
            
            % Was the serialization OK?
            if any(ismember(getListFcn(), {'distributed', 'Composite'}))
                pOneHiddenCompositeWarning;
            end

            % Actually initiate the remote computation
            obj.RemoteSpmdController.initiateRemoteSpmdBlock( comms, fcns );

            % Remember that we did actually send stuff - see delete()
            obj.SentBlockToLabs = true;
        end

        function tf = isComputationComplete( obj )
            obj.errorIfNoSession();

            % Simply query with a shortish timeout
            import java.util.concurrent.TimeUnit;
            tf = obj.RemoteSpmdController.awaitBlock( 200, TimeUnit.MILLISECONDS );

            % Take the opportunity to drain IO here - "false" because we expect to call
            % this again.
            obj.RemoteSpmdController.drainIO( false );

            obj.maybeWarnIfInterruptedAndWaiting();
        end

        function throwBlockExceptions( obj )
            % Look through what we received from executing the block, see if there are
            % any errors to throw.

            if ~obj.SentBlockToLabs
                % Can't possibly be in error.
                return;
            end

            % Retrieve the exceptions from the controller
            [excepts, exceptLabs] = obj.makeBlockExceptionCauses();

            if ~isempty( excepts )
                labsStr = sprintf( '%d ', exceptLabs );
                labsStr = labsStr(1:end-1);
                
                % Create generic SPMD error, add all excepts as cause and
                % as remote cause, and throw.
                except = pWrapExceptionsIntoParallelException(excepts, labsStr);
                throw( except );
            end
        end

        function dispose( obj )
            
            % If we didn't even send the block to the labs, get out of here
            if ~obj.SentBlockToLabs
                if ~isempty( obj.RemoteSpmdController )
                    obj.RemoteSpmdController.dispose();
                end
                return;
            end

            if obj.Disposed
                return;
            end
            
            import java.util.concurrent.TimeUnit;

            % Ensure we dispose of the controller however we leave this
            % delete() method.
            rsc = obj.RemoteSpmdController;
            x = onCleanup( @() dispose(rsc) );

            % Set disposed to true - we don't need to go through here again.
            obj.Disposed = true;
            
            % First, we must ascertain whether or not we completed correctly.
            if ~obj.isComputationComplete()
                obj.errorIfNoSession();
                % The client must have hit CTRL-C or similar - send the interrupt to the labs
                obj.RemoteSpmdController.interrupt();
                
                % carry on waiting - for the block + interrupts to complete. If the user
                % gets bored, they might keep hitting CTRL-C, but the Java
                % layer will continue.
                while ~obj.isComputationComplete()
                    % Wait for the interrupt to take its course
                    obj.errorIfNoSession();
                end
            end

            % Normal completion, attempt to retrieve output arguments - long timeout, as
            % we're not doing anything really useful.
            while ~obj.RemoteSpmdController.awaitCleanup( 1000, TimeUnit.MILLISECONDS )
                % Wait a while
                obj.errorIfNoSession();
            end

            % Retrieve the deserialized parcels
            [deserParcels, E] = obj.extractReturnParcels();

            obj.RemoteSpmdController.drainIO( true );

            % Let the controller know that we're done with it.
            %obj.RemoteSpmdController.dispose();
            obj.callAssignOuts( deserParcels );

            % Now that we've overwritten the old values, flush any remotes - this makes
            % a remote call to remoteClear - but only on those labs that need it.
            obj.RemoteSpmdController.initiateKeyFlush();
            while ~obj.RemoteSpmdController.awaitKeyFlush( 1000, TimeUnit.MILLISECONDS )
                % Wait a while
                obj.errorIfNoSession();
            end
            obj.handleKeyFlushErrors();
            
            % Throw any errors relating to calling the cleanup method.
            if ~isempty( E )
                throw( E );
            end
        end
    end

    methods ( Access = private )

        function [excepts, exceptLabs] = makeBlockExceptionCauses( obj )
        % Look through what we received from executing the block, see if there are
        % any errors to throw.

        % Retrieve the exceptions from the controller
            nlabs      = length( obj.RemoteProcessArray );
            excepts    = cell( 1, nlabs );
            gotExcept  = false( 1, nlabs );

            for lab = 1:nlabs
                % get the FevalResult corresponding to remoteBlockBody
                result = obj.RemoteSpmdController.getBlockResult( obj.RemoteProcessArray( lab ) );
                if ~isempty(result) && result.isError()
                    ex = result.getException();
                    if ~isempty( ex )
                        [isMvmException, thisExcept] = ...
                            parallel.internal.general.transformIfMvmException( ex );
                        if isMvmException
                            % 'ex' is a Java exception, log it.
                            parallel.internal.general.logJavaException(...
                                0, 'Failed to get result from SPMD block.', ex);
                        else
                            thisExcept     = distcompdeserialize( ex );
                            thisExcept     = iBuildParallelException( thisExcept );
                        end
                        excepts{lab}   = obj.maybeTransformMissingSourceException( thisExcept );
                        gotExcept(lab) = true;
                    else
                        % Assumption is that this is an interrupted lab, don't report that.
                    end
                end
            end
            excepts = excepts( gotExcept );
            exceptLabs = find( gotExcept );
        end
        
        function maybeWarnIfInterruptedAndWaiting( obj )
        % If an error has occurred on the labs, warn the users if we've been waiting
        % a long time for the labs to come back.
            
            if obj.WarnedAboutError
                % Seen an error, already warned about it, get out of here
                return;
            end
            
            if isempty( obj.TimeErrorFirstSeen )
                % Haven't seen an error, let's check
                proc = obj.RemoteSpmdController.getProcessCausingInterrupt();
                if ~isempty( proc )
                    obj.TimeErrorFirstSeen = tic;
                end
                % Can't possibly be time to warn, so get out of here
                return;
            end
            
            % Have seen an error - we must be in possession of a TimeErrorFirstSeen - was
            % it long enough ago for us to warn?
            if ( toc( obj.TimeErrorFirstSeen ) >= ...
                 spmdlang.RemoteSpmdExecutor.TIMEOUT_BEFORE_WARNING_ABOUT_ERROR ) 
                
                % Ensure we don't come back through
                obj.WarnedAboutError = true;
                
                % Get the information from the SpmdControllerImpl about the exception
                proc = obj.RemoteSpmdController.getProcessCausingInterrupt();
                msg = obj.RemoteSpmdController.getMessageCausingInterrupt();
                
                if ~isempty( msg ) && ~isempty( proc ) ...
                        && msg.isError() && ~isempty( msg.getException() )
                    % Got a real error
                    except = distcompdeserialize( msg.getException() );
                    except  = iBuildParallelException( except );
                    
                    warning(message('parallel:lang:spmd:WaitingForBlockInterruption', double( proc.getLabIndex() ), except.getReport()));
                end
            end
        end

        function E = maybeTransformMissingSourceException( obj, origE )
        % If the lab threw a SourceCodeNotAvailableOnLab exception, transform that to
        % describe the source code that wasn't found.
            if strcmp( origE.identifier, 'parallel:lang:spmd:SourceCodeNotAvailableOnLab' )
                E = ParallelException.hCloneWithNewMessage( origE, ...
                        'parallel:lang:spmd:SourceCodeNotAvailable', obj.getBodyFunctionFile() );
                E = E.addCause(origE.remotecause{1});
            elseif strcmp( origE.identifier, 'MATLAB:UndefinedFunction' ) && ...
                ~isempty(origE.remotecause{1}.arguments)
                undefinedFunction = origE.remotecause{1}.arguments{1};
                E = ParallelException.hCloneWithNewMessage( origE, ...
                        'parallel:lang:spmd:UndefinedFunctionOnWorker', undefinedFunction, undefinedFunction );
                E = E.addCause(origE.remotecause{1});    
            else
                % No transformation
                E = origE;
            end
        end
        
        function handleKeyFlushErrors( obj )
        % dig through to see if anything bad happened during the remote call to remoteClear
            nlabs = length( obj.RemoteProcessArray );
            for lab = 1:nlabs
                % FevalResult corresponding to remoteClear
                res = obj.RemoteSpmdController.getKeyFlushResult( ...
                    obj.RemoteProcessArray( lab ) );
                if isempty( res )
                    % Nothing happened
                else
                    actResult = res.getResult();
                    if ~logical( actResult(1) )
                        except = distcompdeserialize( res(2) );
                        warning(message('parallel:lang:spmd:RemoteKeyFlush', lab, getReport( except )));
                    end
                end
            end
        end
        
        function [deserParcels, E] = extractReturnParcels( obj )
        % Process the results of the cleanup - return the parcels, and an error that
        % we think ought to be thrown.
            
            nlabs = length( obj.RemoteProcessArray );
            E     = [];
            
            % Pre-allocate the return:
            deserParcels = cell( 1, obj.SpmdNargout );
            for argout = 1:obj.SpmdNargout
                deserParcels{argout} = cell( 1, nlabs );
            end
            
            for lab = 1:nlabs
                % Retrieve FevalResult corresponding to remoteBlockCleanup
                res = obj.RemoteSpmdController.getCleanupResult( ...
                    obj.RemoteProcessArray( lab ) );

                if isempty( res )
                    resultOk = false;
                else
                    tmp = res.getResult;
                    if isempty(tmp)
                        res = res.getException();
                        resultOk = false;
                    else
                        res = tmp;
                        resultOk = logical(res(1));
                    end
                end

                if resultOk
                    thisLabsReturn = distcompdeserialize( res(2) );
                    if ~iscell( thisLabsReturn ) || isempty( thisLabsReturn )
                        % If thisLabsReturn is completely empty - this means that the lab didn't
                        % execute any of the lines needed to produce
                        % output. If it's not even a cell, then the
                        % remoteBlockBody>iStashOutputs didn't get called.
                    else
                        for argout = 1:obj.SpmdNargout
                            deserParcels{argout}{lab} = thisLabsReturn{argout};
                        end
                    end
                else
                    E = obj.handleBadCleanupResult( E, res, lab );
                end
            end
        end

        function E = handleBadCleanupResult( ~, E, res, lab )
        % Called during extractReturnParcels() to handle an FevalResult object
        % returned from a lab which indicates problems
            if isempty( E )
                E = ParallelException.hBuildLikeMException(...
                    message('parallel:lang:spmd:RemoteCleanupErrored'));
            end
            if isempty( res )
                % This is really serious - there was no remoteBlockCleanup return from the
                % lab, which may mean that the parallel pool is stuck.
                E = iAddRemoteCause( E, MException(message('parallel:lang:spmd:NoRemoteCleanup', lab)) );
            else
                try
                    if isa( res, 'java.lang.Exception' )
                        parallel.internal.general.logJavaException(...
                            0, 'Failed to get result from SPMD cleanup.', res);
                        [isMvmException, parException] = ...
                            parallel.internal.general.transformIfMvmException( res );
                        if isMvmException
                            E = iAddRemoteCause( E, parException );
                        end
                    else
                        if ~logical( res(1) )
                            E = iAddRemoteCause( E, distcompdeserialize( res(2) ) );
                        end
                    end
                catch EE
                    E = addCause( E, EE );
                end
            end
        end

    end
end

function E = iAddRemoteCause(E, cause)
E = addCause(E, cause);
E = hAddRemoteCause(E, cause);
end

function ctrl = buildController( procArray, session )
    try
        ctrl = parallel.internal.getJavaFutureInterruptibly(...
            session.createSpmdController());
        % If we successfully built the controller, must ensure
        % that it gets disposed if acquireLabs fails.
        tmpCleanupCtrl = parallel.internal.general.DisarmableOncleanup(...
            @() ctrl.dispose());
        ctrl.acquireLabs( procArray );
    catch E
        %% if this is a java exception that someone has provided a non-generic ID for
        if isJavaException( E ) &&  ~ isequal( E.identifier , 'Matlab:Java:GenericException' )
            rethrow(E);
        else
            except = MException(message('parallel:lang:spmd:SpmdCreationError'));
            except = addCause( except, E );
        end

        throw( except );
    end
    tmpCleanupCtrl.disarm();
end

function E = iBuildParallelException( origErr )
    clientStackToIgnore = 'spmd_feval';
    E = ParallelException.hBuildRemoteParallelException( ...
        origErr, clientStackToIgnore );
end
