%JavaBackedSession Wrapper around a Java session object.

% Copyright 2020-2021 The MathWorks, Inc.

classdef JavaBackedSession < parallel.internal.pool.ISession
    properties (Constant, GetAccess = private)
        EnforceInitMatlabRef = iInitMatlabRef();
    end
    
    properties (Access = private)
        Session % The underlying Java object.
        ConnectionManager % only valid during the setup phase
        ConnectionStrategy % only valid during the setup phase
        WriterOnCleanup % Used by workers to cleanup MatlabOutputWriters
        PathNotificationOnCleanup % Used by clients to turn off path/clear notification
    end

    % AbstractDataQueue needs to access detailed communication workings of the
    % underlying Java session object.
    methods (Access = ?parallel.internal.dataqueue.RemoteJavaQueueFactory)
        function s = getInternalSession(obj)
            s = obj.Session;
        end
    end

    methods (Access = {?parallel.internal.pool.InteractivePmodeClient, ...
                       ?parallel.internal.pool.InteractiveWorker, ...
                       ?parallel.internal.pool.ISession})
        function l = pmodeGetLabs(obj)
            l = obj.Session.getLabs();
        end

        function l = pmodeGetTransfer(obj)
            l = obj.Session.getTransfer();
        end
    end

    methods (Access = {?parallel.internal.pool.AbstractClusterPool, ...
                       ?parallel.internal.pool.ISession})
        function uuid = getSessionUUID(obj)
            uuid = char(obj.Session.getSessionUUID());
        end

        function tf = canInitializeSpmd(obj)
            tf = obj.Session.canInitializeSpmd();
        end

        function tf = initializeSpmd(obj)
            tf = obj.Session.initializeSpmd();
        end
        
        function fq = getFevalQueue(obj, pool)
            fq = parallel.FevalQueue.getQueueForSession(pool, ...
                @(id)  parallel.internal.queue.JavaFevalQueue(obj.Session, id));
        end
    end

    % Worker-only methods
    methods (Access = {?parallel.internal.pool.AbstractWorker, ...
                       ?parallel.internal.pool.ISession})
        function setDependencyDir(obj, dependencyDir)
            fda = obj.getFileDependenciesAssistant();
            fda.setDependencyDir(dependencyDir);
        end
    end
    
    % These methods are *not* exposed through the ISession interface as they are
    % too tied to the implementation of SessionLogger
    methods (Access = ?parallel.internal.profiling.SessionLogger)
        function addProfilingListener(obj, l)
            obj.Session.addProfilingListener(l)
        end

        function removeProfilingListener(obj, l)
            obj.Session.removeProfilingListener(l)
        end
    end
    % These methods are *not* exposed through the ISession interface as they are
    % too tied to the current implementation of PoolProfiler
    methods (Access = ?parallel.internal.profiling.PoolProfiler)
        function didBecomeIdle = waitUntilPoolIdle(obj, timeoutMillis)
            rm = obj.Session.getResourceManager();
            didBecomeIdle = rm.blockUntilWorkersIdle(timeoutMillis);
        end
    end

    methods (Access = {?parallel.pool.TicBytesResult, ...
                       ?parallel.internal.pool.ISession})
        function info = getCurrentBytesTransferredToInstances(obj)
            info = obj.Session.getCurrentBytesTransferredToInstances();
        end
    end

    methods (Access = {?parallel.internal.pool.AbstractClient, ...
                       ?parallel.internal.pool.ISession})
        function tf = destroyClientSession(obj)
            % Run the Java garbage collector to clean up leaked file
            % descriptors on Linux -- See G904795 for details
            cleanupJavaGC = onCleanup(@()java.lang.System.gc());
            tf = obj.Session.destroyClientSession();
        end
        function startSendPathAndClearNotificationToLabs(obj)
            obj.Session.startSendPathAndClearNotificationToLabs()
            dctPathAndClearNotificationGateway('on');
            % Ensure that no matter how the pool is shut down, we
            % turn off path/clear notificaton (g2335187).
            obj.PathNotificationOnCleanup = ...
                onCleanup(@() dctPathAndClearNotificationGateway('off'));
        end
        function tf = waitForSessionToStart(obj, timeInSeconds)
            tf = obj.Session.waitForSessionToStart(timeInSeconds * 1000, ...
                                                   java.util.concurrent.TimeUnit.MILLISECONDS);
        end
        function waitForConnections(obj, spmdInitialized, sessionInfo, checkFcn)
        % Here, the client has submitted the job, and we must build the session and
        % wait for it to establish communication with the workers.
            import com.mathworks.toolbox.distcomp.pmode.matlabpool.sessions.ClientSessionBuilder
            import java.util.concurrent.atomic.AtomicInteger
            import parallel.internal.pool.JavaBackedSessionNotifier

            connectionCounter = AtomicInteger(0);

            javaInfo = JavaBackedSessionNotifier.sessionInfoToJava(sessionInfo);
            sessionFuture = ClientSessionBuilder.initiate(...
                obj.ConnectionManager, spmdInitialized, javaInfo, ...
                connectionCounter, obj.ConnectionStrategy);
            stopSessionFuture = onCleanup(@() sessionFuture.cancel(true));
            % We've handed the ConnectionManager to the java session object, so we no
            % longer need to ensure that it is closed here.
            obj.ConnectionManager = [];
            obj.ConnectionStrategy = [];
            obj.Session = obj.waitForSessionCreation(sessionFuture, connectionCounter, ...
                                                     checkFcn);
        end
    end

    methods (Access = private)
        function obj = JavaBackedSession(varargin)
            if nargin == 1
                obj.Session = varargin{1};
            else
                obj.ConnectionManager = varargin{1};
                obj.ConnectionStrategy = varargin{2};
            end
        end
    end
    
    methods
        function notifySerializationStatus(obj, isSer, isStart, id, bytes)
            obj.Session.notifySerializationStatus(isSer, isStart, id, bytes)
        end

        function evalOnClient(obj, cmd)
            obj.Session.getClient().evalConsoleOutput(cmd);
        end

        function releaseCurrentParforController(obj)
            obj.Session.releaseCurrentParforController();
        end

        function fda = getFileDependenciesAssistant(obj)
            fda = obj.Session.getFileDependenciesAssistant();
        end

        function tf = isSessionRunning(obj)
            tf = obj.Session.isSessionRunning();
        end

        function setRestartOnClusterChange(obj, tf)
            obj.Session.setRestartOnClusterChange(tf)
        end
        
        function info = getClientSessionInfo(obj)
            javaInfo = obj.Session.getClientSessionInfo();
            info = parallel.internal.pool.JavaBackedSessionNotifier.sessionInfoFromJava(javaInfo);
        end
        
        function rehashed = rehashIfNecessary(~)
           rehashed = com.mathworks.toolbox.distcomp.pmode.RehashNotificationCommand.resetRehashRequestCount();
           if rehashed
               rehash();
           end
        end
        
        function itemsToClear = drainItemsToClear(~, javaObj)
            % Java layer aggregates clear requests, each stored as a
            % serialized PathInfo. NB: For performance reasons, deserialize
            % items into a cell array and then convert to an array; it is
            % much faster when there are lots of commands!
            clearCommands = javaObj.drainClearCommands();
            itemsToClear = arrayfun(@(cmd) parallel.internal.apishared.PathInfo.getDeserializedPathInfo(cmd.getClearData()), clearCommands, 'UniformOutput', false);
            itemsToClear = [itemsToClear{:}]; % Convert to array
        end

        function setShutdownIdleTimeout(obj, timeInSeconds)
            obj.Session.setShutdownIdleTimeout(timeInSeconds)
        end


        function n = getPoolSize(obj)
            n = double(obj.Session.getPoolSize());
        end

        function ctrl = createSpmdController(obj, labIndices)
            jctrl = parallel.internal.getJavaFutureInterruptibly(...
                obj.Session.createSpmdController());
            ctrl = parallel.internal.pool.JavaRemoteSpmdController.create(jctrl, labIndices);
        end

        function futureController = createParforController(obj, ~, ~)
            futureController = obj.Session.createParforController();
        end

        function assistant = getCompositeAssistant(obj, labIndices)
            assistant = parallel.internal.pool.JavaRemoteCompositeAssistant.create(...
                    obj.Session, labIndices);
        end
        
        function map = getCompositeKeysToClearMap(obj)
            map = obj.Session.getCompositeKeysToClearMap();
        end

        function factory = createRemoteQueueFactory(obj, queueUuid)
            factory = parallel.internal.dataqueue.RemoteJavaQueueFactory(queueUuid, obj.Session);
        end
        
        function shutdown(obj)
            obj.Session.shutdown();
        end
        
        function delete(obj)
        % Run the Java garbage collector to clean up leaked file
        % descriptors on Linux -- See G904795 for details
            java.lang.System.gc();
            if ~isempty(obj.ConnectionManager)
                obj.ConnectionManager.close();
            end
        end
    end

    methods (Access = private)
        function session = waitForSessionCreation(~, sessionFuture, connectionCounter, ...
                                                  checkFcn)
        % Block until the session has been created - which completes only when all the
        % connections are available.
            gotSession = false;
            session = [];
            previouslyConnectedTo = 0;
            while ~gotSession
                % This throws an appropriate error in the case where things go wrong.
                [gotSession, session] = parallel.internal.getJavaFutureResult(...
                    sessionFuture, 1, java.util.concurrent.TimeUnit.SECONDS);
                if gotSession
                    return
                end
                % If we get here, we have no session. Let's check to see how things are getting
                % on using the injected checkFcn - this might throw an error if things are bad.
                checkFcn();
                currentlyConnectedTo = double(connectionCounter.get());
                if currentlyConnectedTo > previouslyConnectedTo
                    dctSchedulerMessage(2, 'Currently connected to: %d', currentlyConnectedTo);
                    previouslyConnectedTo = currentlyConnectedTo;
                end
            end
        end
        function connInfo = getConnectionInfo(obj)
        % We need the name of this machine to connect back to from the parallel
        % job.  Get from pctconfig as it allows the user to override the default
        % hostname.
            [~, undoc] = pctconfig();

            % Build the connection info object that we are going to send to the remote
            % machines. This contains an InetSocketAddress object and a UUID that
            % identifies this connection attempt. The undoc method below allows us to
            % simulate the situation where we get the wrong IP for the local machine
            % and still expect the remote machine to undertake a lookup to see if they
            % can resolve the correct address. This simulation only makes sense if
            % we are sending a ServerSocketConnectInfo to the workers
            connMgr = obj.ConnectionManager;
            connInfo = connMgr.getInfoToConnect();
            if ~isempty( undoc.hostinetaddress ) && ...
                    connInfo.getClass().getSimpleName().equals('ServerSocketConnectInfo')
                connInfo = connMgr.getInfoToConnect().createMangledConnectionInfo( undoc.hostinetaddress );
            end
        end
    end
    
    methods (Static)
        function buildThreadSession(sessionInfo, pool)
        % Used by doLocalThreadsParpool. This doesn't actually need to return
        % an ISession instance as it is never needed.
            import parallel.internal.pool.JavaBackedSessionNotifier
            import com.mathworks.toolbox.distcomp.pmode.SessionFactory
            
            javaSessionInfo = JavaBackedSessionNotifier.sessionInfoToJava(sessionInfo);
            session = SessionFactory.createThreadSession(javaSessionInfo, pool.NumWorkers);
            [~] = addlistener(pool, 'ObjectBeingDestroyed', @(~, ~) session.notifyShutdown());
        end
        
        
        function [obj, workerBuild] = buildInteractiveClient(cluster)
        % This build method is used for parpool and pmode.
            
            import parallel.internal.pool.JavaBackedSession
            try
                [connMgr, connStrategy] = cluster.pCreateConnectionManager();
            catch err
                rethrow(err);
            end
            % Put the connMgr into the JavaBackedSession right now so that
            % we can be sure it gets cleaned up if necessary.
            obj = JavaBackedSession(connMgr, connStrategy);

            connInfo = obj.getConnectionInfo();
            workerBuild = {@(workerIdx, numWorkers, broadcastFcn) ...
                JavaBackedSession.buildWorker(...
                    workerIdx, numWorkers, broadcastFcn, connInfo, connStrategy)};
        end
        
        function [obj, workerBuild] = buildBatchClient()
        % This build method is used for the client-worker of a batch pool.
            
            import com.mathworks.toolbox.distcomp.pmode.matlabpool.sessions.ConnectionTopology
            import parallel.internal.apishared.ConnMgrBuilder
            import parallel.internal.pool.JavaBackedSession
            
            connMgr = ConnMgrBuilder.buildPoolClientConnMgr();
            % Put the connMgr into the JavaBackedSession right now so that
            % we can be sure it gets cleaned up if necessary.
            obj = JavaBackedSession(connMgr, []);
            connInfo = obj.getConnectionInfo();
            connStrategy = ConnectionTopology.ALL_WORKERS_CONNECT_TO_CLIENT;

            workerBuild = @(workerIdx, numWorkers, broadcastFcn) ...
                JavaBackedSession.buildWorker(...
                    workerIdx, numWorkers, broadcastFcn, connInfo, connStrategy);
        end

        function obj = buildWorker(workerIdx, numWorkers, broadcastFunctor, connInfo, connStrategy)
        % This build method is used for interactive workers and batch pool workers.
            import com.mathworks.toolbox.distcomp.mjs.jobmanager.PortSelectionInformation
            import com.mathworks.toolbox.distcomp.pmode.poolmessaging.ConnectionManager
            import com.mathworks.toolbox.distcomp.pmode.matlabpool.sessions.WorkerSessionBuilder
            import com.mathworks.toolbox.distcomp.pmode.matlabpool.sessions.ConnectionTopology
            import com.mathworks.toolbox.distcomp.pmode.poolmessaging.ProcessInstance
            import parallel.internal.apishared.ConnMgrBuilder

            dctSchedulerMessage(4, 'Creating session with strategy: %s', ...
                                char(connStrategy));

            [cfg, undoc]      = pctconfig();
            hostname          = cfg.hostname;
            internalhostname  = undoc.internalhostname;
            numContending     = numWorkers;
            connectingProcess = ProcessInstance.getClientInstance();
            
            root = distcomp.getdistcompobjectroot();
            portSelectionInformation = root.CurrentRunprop.PortSelectionInformation;
            
            % Convert PortSelectionInformation to a Java class, if necessary
            if isa(portSelectionInformation, 'parallel.internal.cluster.PortSelectionInformation')
                portSelectionInformation = PortSelectionInformation(...
                    portSelectionInformation.IndexOnHost, ...
                    portSelectionInformation.NumWorkersOnHost, ...
                    portSelectionInformation.MinPortOffset);
            end

            if connStrategy == ConnectionTopology.WORKERS_PROXIED_BY_LAB_ONE
                if workerIdx == 1
                    sessionFuture = iBuildProxyingWorkerSession(workerIdx, numWorkers, broadcastFunctor, ...
                                                                hostname, connInfo, internalhostname, ...
                                                                connectingProcess, portSelectionInformation);
                else
                    sessionFuture = iBuildProxiedWorkerSession(workerIdx, numWorkers, broadcastFunctor, ...
                                                               numContending, hostname, portSelectionInformation);
                end
            else
                % "Standard" unproxied worker session.
                connMgr = ConnectionManager.buildLabConnManager(workerIdx, ...
                                                                numWorkers, ...
                                                                numContending, ...
                                                                hostname, ...
                                                                connInfo, ...
                                                                connectingProcess, ...
                                                                portSelectionInformation);

                sessionFuture = WorkerSessionBuilder.initiatePlain(connMgr, workerIdx);
            end

            try
                obj = parallel.internal.pool.JavaBackedSession(...
                    parallel.internal.getJavaFutureInterruptibly(sessionFuture));

                % Set up the streams for asynchronous output redirection,
                % only used by/needed for pmode
                com.mathworks.toolbox.distcomp.pmode.MatlabOutputWriters.getInstance.setup();
                obj.WriterOnCleanup = onCleanup(@() com.mathworks.toolbox.distcomp.pmode.MatlabOutputWriters.getInstance.teardown());
            catch err
                dctSchedulerMessage(1, ['Failed to create session object with ', ...
                                    'the message:\n''%s'''], err.message);
                rethrow(err);
            end
        end
    end
    
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function success = iInitMatlabRef()
    com.mathworks.toolbox.distcomp.util.MatlabRefStore.initMatlabRef();
    success = true;
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function sessionFuture = iBuildProxyingWorkerSession(workerIdx, numWorkers, broadcastFunctor, ...
                                                     hostname, connInfo, internalhostname, ...
                                                     connectingProcess, portSelectionInformation)
    import com.mathworks.toolbox.distcomp.pmode.poolmessaging.ConnectionManager
    import com.mathworks.toolbox.distcomp.pmode.matlabpool.sessions.WorkerSessionBuilder
    import parallel.internal.apishared.ConnMgrBuilder 

    % numContending is required by the broker to understand how many requests it
    % might expect. In this case, we know that number to be 1.
    numContending = 1;
    connMgr = ConnectionManager.buildLabConnManager(workerIdx, ...
                                                    numWorkers, ...
                                                    numContending, ...
                                                    hostname, ...
                                                    connInfo, ...
                                                    connectingProcess, ...
                                                    portSelectionInformation);

    % We are acting as the proxy - build an acceptor to receive connections.
    dctSchedulerMessage(4, 'About to buildLabOneAcceptor');
    acceptor = ConnMgrBuilder.buildLabOneAcceptor(...
        connMgr.getLocalInstance());
    dctSchedulerMessage(5, 'Built lab one acceptor: %s.', ...
                        char(acceptor.toString()));
    dctSchedulerMessage(5, 'Lab one connect info: %s.', ...
                        char(acceptor.getInfoToConnect().toString()));

    % Tell the other workers how to get connected
    infoToConnectToLab1 = acceptor.getInfoToConnect();
    
    % If the 'internalhostname' is set, override infoToConnectToLab1 so that workers
    % connect on that name.
    if ~isempty(internalhostname)
        dctSchedulerMessage(2, 'Overriding host of connect info from %s to %s.', ...
                            char(infoToConnectToLab1.getSocketAddress()), ...
                            internalhostname);
        try
            internalInetAddr = java.net.InetAddress.getByName(internalhostname);
            infoToConnectToLab1 = infoToConnectToLab1.createMangledConnectionInfo(...
                internalInetAddr);
        catch E
            dctSchedulerMessage(1, 'Failed to override connect info:', E);
        end
    end
    
    lab1Instance = connMgr.getLocalInstance();
    broadcastFunctor({infoToConnectToLab1, lab1Instance});

    sessionFuture = WorkerSessionBuilder.initiateProxyingWorker(...
        connMgr, workerIdx, lab1Instance, acceptor, ...
        mpiInitialized);
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function sessionFuture = iBuildProxiedWorkerSession(workerIdx, numWorkers, broadcastFunctor, ...
                                                    numContending, hostname, portSelectionInformation)
    import com.mathworks.toolbox.distcomp.pmode.poolmessaging.ConnectionManager
    import com.mathworks.toolbox.distcomp.pmode.matlabpool.sessions.WorkerSessionBuilder
    import com.mathworks.toolbox.distcomp.pmode.poolmessaging.ProcessInstance

    % Get connection information.
    infoAndInstance = broadcastFunctor();
    infoToConnectToLab1 = infoAndInstance{1};
    lab1Instance = infoAndInstance{2};
    lab1Process = ProcessInstance.getLabInstance(1);
    connMgr = ConnectionManager.buildLabConnManager(workerIdx, ...
                                                    numWorkers, ...
                                                    numContending, ...
                                                    hostname, ...
                                                    infoToConnectToLab1, ...
                                                    lab1Process, ...
                                                    portSelectionInformation);

    sessionFuture = WorkerSessionBuilder.initiateProxiedWorker(...
        connMgr, workerIdx, lab1Instance, infoToConnectToLab1);
end

%#ok<*JAPIMATHWORKS> transition class.
