%LeadWorkers Implement lead-worker-on-host behaviour

% Copyright 2014-2020 The MathWorks, Inc.

classdef (Sealed, Hidden) LeadWorkers < handle

    properties (GetAccess = private, SetAccess = immutable)
        Pool
    end
    properties (Access = private)
        PoolSize    = -1
        LeadTIDs    = []
        TIDOfLeader = []
    end

    methods (Access = private)
        function updateLeadInfo(obj)
        % updateLeadInfo Synchronously update lead-worker information.

            if obj.Pool.NumWorkers ~= obj.PoolSize
                % Pool has changed size, recalculate everything.

                partitionFcn = ...
                    parallel.internal.pool.LeadWorkers.getSetWorkerPartitionFunction();

                % Invoke the partition
                [allTIDs, tidsOfLeadWorkers, tidOfMyLeadWorker] = ...
                    partitionFcn(obj.Pool);

                % Inform all workers of the new leader status.
                fetchOutputs(...
                    parfevalOnAll(...
                        obj.Pool, ...
                        @iGetSetLeadStatus, 1, tidsOfLeadWorkers, allTIDs, tidOfMyLeadWorker));

                % Update object properties.
                obj.LeadTIDs    = tidsOfLeadWorkers;
                obj.TIDOfLeader = tidOfMyLeadWorker;
                obj.PoolSize    = obj.Pool.NumWorkers;
            end
        end
        
        function obj = LeadWorkers(pool)
        %LeadWorkers Build LeadWorkers for this pool.
            obj.Pool = pool;
            obj.updateLeadInfo();
        end
    end

    methods (Static)
        function obj = getForPool(pool)
        %LeadWorkers.getForPool Get or create LeadWorkers for specified pool.

            persistent LAST_OBJ
            narginchk(1,1);
            if ~isempty(LAST_OBJ) && ...
                pool == LAST_OBJ.Pool
                % Nothing to do
            else
                assert(isa(pool, 'parallel.Pool') && ...
                       isscalar(pool) && ...
                       isvalid(pool) && ...
                       pool.Connected);
                LAST_OBJ = parallel.internal.pool.LeadWorkers(pool);
            end
            obj = LAST_OBJ;
        end
    end
    methods (Hidden, Static)
        function fcn = getSetWorkerPartitionFunction(newFcn)
        %LeadWorkers.getSetWorkerPartitionFunction
        % Allow QE to define how to partition workers.
            persistent FCN
            defaultFcn = @iDefaultPartition;
            if isempty(FCN)
                FCN = defaultFcn;
            end
            fcn = FCN;
            if nargin > 0
                if isempty(newFcn)
                    % reset to the default
                    FCN = defaultFcn;
                else
                    FCN = newFcn;
                end
            end
        end
    end

    methods
        function f = fevalOnLeadWorkers(obj, fcn, numOut, varargin)
        % Implementation of parallel.pool.parfevalOnEachHost
            obj.updateLeadInfo();
            % This is a little ugly, but unfortunately
            % inevitable until this code becomes more integrated
            % with the main parfeval infrastructure.
            if isempty(getCurrentTask())
                % Interactive case. WorkerID matches current task ID
                workerIDOffset = 0;
            else
                % Batch case. WorkerID is task ID - 1.
                workerIDOffset = 1;
            end
            f = parallel.internal.queue.FilteringFuture.createAndSubmit(...
                obj.Pool, obj.LeadTIDs, workerIDOffset, ...
                fcn, numOut, varargin{:});
            f = parallel.Future.hBuild(f);
            wait(f);
        end

        function f = fevalOnAll(obj, varargin)
        % Wraps parfevalOnAll, ensuring the workers have up-to-date
        % information about lead workers.
            obj.updateLeadInfo();
            f = parfevalOnAll(obj.Pool, varargin{:});
        end
    end
    methods (Static)
        function tf = isLeadWorker()
        % isLeadWorker TRUE if this worker is a leader.
            assert(parallel.internal.pool.isPoolWorker());
            tf = iGetSetLeadStatus();
        end
        function idx = getLeadOnThisHost()
        % Return the task ID of the lead worker on this host.
            assert(parallel.internal.pool.isPoolWorker());
            [~, idx] = iGetSetLeadStatus();
        end
    end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [allTIDs, tidsOfLeadWorkers, tidOfMyLeadWorker] = iDefaultPartition(pool)
    % Get task IDs and corresponding hosts
    [tids, hosts] = fetchOutputs(...
        parfevalOnAll(pool, @iGetWorkerInfo, 2), ...
        'UniformOutput', false);

    % Decide who is the lead worker using UNIQUE and then for each task, work out
    % which task is the local leader.
    [~, correspondingIndices, ic] = unique(hosts);
    tidsOfLeadWorkers = [tids{correspondingIndices}];
    tidOfMyLeadWorker = reshape(tidsOfLeadWorkers(ic), 1, []);
    allTIDs = [tids{:}];
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [tf, leadTID] = iGetSetLeadStatus(leadTIDs, allTIDs, tidOfMyLeadWorker)
    mlock
    persistent IS_LEAD MY_LEADER
    assert(nargin > 0 || ~isempty(IS_LEAD));
    if nargin > 0
        myTID = iGetTID();
        IS_LEAD = ismember(myTID, leadTIDs);

        locInAllTIDs = find(allTIDs == myTID, 1, 'first');
        assert(~isempty(locInAllTIDs));
        MY_LEADER = tidOfMyLeadWorker(locInAllTIDs);
    end
    tf = IS_LEAD;
    leadTID = MY_LEADER;
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function tid = iGetTID()
    t = getCurrentTask();
    assert(~isempty(t));
    tid = t.ID;
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function host = iGetHostname()
    config = pctconfig();
    host   = config.hostname;
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [taskID, host] = iGetWorkerInfo()
    taskID = iGetTID();
    host   = iGetHostname();
end


