%FolderManagement Object supporting directory and file transfers for parallel.pool

% Copyright 2014-2016 The MathWorks, Inc.

classdef (Sealed, Hidden) FolderManagement < handle

    properties (GetAccess = private, SetAccess = immutable)
        Workers % Instance of LeadWorkers
        Pool    % The pool we operate on.
    end
    methods (Access = private)
        function obj = FolderManagement(pool)
            obj.Pool    = pool;
            obj.Workers = parallel.internal.pool.LeadWorkers.getForPool(pool);

            [tids, dirs] = fetchOutputs(obj.Workers.fevalOnAll(...
                @iGetTidAndPerHostRoot, 2), 'UniformOutput', false);

            % Use fetchOutputs here simply to block and assert no errors. Set
            % 'UniformOutput' to false to ensure no problems with per-host
            % directories of differing lengths.
            fetchOutputs(obj.Workers.fevalOnAll(...
                @iGetSetActivePerHostDir, 1, [tids{:}], dirs), ...
                         'UniformOutput', false);
        end
    end

    methods (Access = private, Static)
        function d = getPerHostOrWorkerDir(baseDirFcn, subdir)
            assert(parallel.internal.pool.isPoolWorker(), ...
                   message('parallel:pool:DirectoryOnlyOnWorkers'));
            d = fullfile(baseDirFcn(), subdir);
            assert(exist(d, 'dir') == 7, ...
                   message('parallel:pool:NoSuchDirectory', subdir));
        end
    end
    
    methods (Static)
        function obj = getForPool(pool)
        % Static access to the FolderManagement for specified pool.
            assert(~parallel.internal.pool.isPoolWorker(), ...
                   message('parallel:pool:DirectoryManagementOnClient'));
            assert(isscalar(pool) && ...
                   isa(pool, 'parallel.Pool') && ...
                   isvalid(pool) && ...
                   pool.Connected, ...
                   message('parallel:pool:ValidPoolRequired'));
            persistent LAST_OBJ
            narginchk(1, 1);
            if ~isempty(LAST_OBJ) && ...
                    pool == LAST_OBJ.Pool
                % Nothing to do.
            else
                LAST_OBJ = parallel.internal.pool.FolderManagement(pool);
            end
            obj = LAST_OBJ;
        end
        
        function d = getPerWorkerFolder(subdir)
        % Worker-side access to the per-worker directory.
            d = parallel.internal.pool.FolderManagement.getPerHostOrWorkerDir(...
                @iPerWorkerRoot, subdir);
        end
        function d = getPerHostFolder(subdir)
        % Worker-side access to the per-host directory.
            d = parallel.internal.pool.FolderManagement.getPerHostOrWorkerDir(...
                @iGetSetActivePerHostDir, subdir);
        end
        
        function OK = writeZipData(data, destDir)
        % Write zip data to destination folder
            tempZip = [tempname, '.zip'];
            x = onCleanup(@() delete(tempZip));
            fh = fopen(tempZip, 'wb');
            fwrite(fh, data, 'int8');
            fclose(fh);
            try
                unzip(tempZip, destDir);
            catch E
                throw(addCause(MException(message('parallel:pool:FileTransferError')), E));
            end
            OK = true;
        end
    end
    methods
        function f = createFolder(obj, name)
            f = obj.Workers.fevalOnAll(@iCreateFolder, 1, name);
        end
        function f = deleteFolder(obj, name)
            f = obj.Workers.fevalOnAll(@iDeleteFolder, 1, name);
        end
        function f = listFolders(obj)
            f = obj.Workers.fevalOnAll(@iListFolders, 1);
        end
        function f = replicateFolders(obj, name)
            f = obj.Workers.fevalOnLeadWorkers(@iZipPerHostLocation, 2, name);
            [tid, loc] = fetchOutputs(f, 'UniformOutput', false);
            map = containers.Map(tid, loc);
            f = obj.Workers.fevalOnAll(@iUnzipFromPerHostLocation, 1, map, name);
        end
        function f = copyToPerHost(obj, subdir, filesOrFolders)
            zippedData = iZipFilesAndFolders(filesOrFolders);
            f = obj.Workers.fevalOnLeadWorkers(@iUnzipPerHost, 1, subdir, zippedData);
        end
        function f = copyToPerWorker(obj, subdir, filesOrFolders)
            zippedData = iZipFilesAndFolders(filesOrFolders);
            f = obj.Workers.fevalOnAll(@iUnzipPerWorker, 1, subdir, zippedData);
        end
        function f = getZippedFolderDataFromPerHost(obj, subdir)
            f = obj.Workers.fevalOnLeadWorkers(@iZipDataPerHostLocation, 2, subdir);
        end
        function f = getZippedFolderDataFromPerWorker(obj, subdir)
            f = obj.Workers.fevalOnAll(@iZipDataPerWorkerLocation, 2, subdir);
        end
    end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function data = iZipFilesAndFolders(filesOrFolders)
    data = parallel.internal.apishared.AttachedFiles.zipFiles(filesOrFolders);
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function OK = iUnzipPerHost(subdir, data)
    OK = parallel.internal.pool.FolderManagement.writeZipData(data, fullfile(iMyPerHostRoot(), subdir));
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function OK = iUnzipPerWorker(subdir, data)
    OK = parallel.internal.pool.FolderManagement.writeZipData(data, fullfile(iPerWorkerRoot(), subdir));
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function iEnsureDirOrThrow(d)
    if exist(d, 'dir') ~= 7
        [success, message, messageId] = mkdir(d);
        if ~success
            error(messageId, '%s', message);
        end
    end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [tid, zipFile] = iZipPerHostLocation(name)
    t = getCurrentTask();
    assert(~isempty(t));
    tid = t.ID;
    dirToZip = fullfile(iMyPerHostRoot(), name);
    assert(exist(dirToZip, 'dir') == 7);
    zipFile  = [dirToZip, '.zip'];
    relRoot  = iMyPerHostRoot();
    try
        zip(zipFile, dirToZip, relRoot);
    catch E
        throw(addCause(MException(message('parallel:pool:FileTransferError')), E));
    end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [tid, data] = iZipDataPerWorkerLocation(subdir)
    t = getCurrentTask();
    assert(~isempty(t));
    tid = t.ID;
    dirToZip = fullfile(iPerWorkerRoot(), subdir);
    assert(exist(dirToZip, 'dir') == 7);
    try
        data = parallel.internal.apishared.AttachedFiles.zipFiles({dirToZip});
    catch E
        throw(addCause(MException(message('parallel:pool:FileTransferError')), E));
    end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [tid, data] = iZipDataPerHostLocation(subdir)
    t = getCurrentTask();
    assert(~isempty(t));
    tid = t.ID;
    dirToZip = fullfile(iMyPerHostRoot(), subdir);
    assert(exist(dirToZip, 'dir') == 7);
    try
        data = parallel.internal.apishared.AttachedFiles.zipFiles({dirToZip});
    catch E
        throw(addCause(MException(message('parallel:pool:FileTransferError')), E));
    end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function OK = iUnzipFromPerHostLocation(map, name)
    leadTID = parallel.internal.pool.LeadWorkers.getLeadOnThisHost();
    zipFile = map(leadTID);
    assert(exist(zipFile, 'file') == 2);
    try
        unzip(zipFile, iPerWorkerRoot());
    catch E
        throw(addCause(MException(message('parallel:pool:FileTransferError')), E));
    end
    expectedDirectory = fullfile(iPerWorkerRoot(), name);
    OK = exist(expectedDirectory, 'dir') == 7;
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function OK = iRemoveDir(d)
    if exist(d, 'dir') == 7
        [success, message, messageId] = rmdir(d, 's');
        if ~success
            error(messageId, '%s', message);
        end
    end
    OK = exist(d, 'dir') ~= 7;
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [tid, phr] = iGetTidAndPerHostRoot()
    t = getCurrentTask();
    tid = t.ID;
    phr = iMyPerHostRoot();
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function rd = iMyPerHostRoot()
    root    = distcomp.getdistcompobjectroot;
    runprop = root.CurrentRunprop;
    assert(~isempty(runprop));
    rd      = getPerHostFolder(runprop);
    iEnsureDirOrThrow(rd);
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function rd = iPerWorkerRoot()
    root    = distcomp.getdistcompobjectroot;
    runprop = root.CurrentRunprop;
    assert(~isempty(runprop));
    rd      = getPerWorkerFolder(runprop);
    iEnsureDirOrThrow(rd);
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function rd = iGetSetActivePerHostDir(allTids, roots)
    mlock % protect worker-side persistent data
    persistent ACTIVE_ROOT
    if nargin > 0
        activeTid = parallel.internal.pool.LeadWorkers.getLeadOnThisHost();
        idx = find(allTids == activeTid, 1, 'first');
        assert(~isempty(idx));
        ACTIVE_ROOT = roots{idx};
    end
    rd = ACTIVE_ROOT;
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function OK = iCreateFolder(name)
    if parallel.internal.pool.LeadWorkers.isLeadWorker()
        iEnsureDirOrThrow(fullfile(iMyPerHostRoot(), name));
    end
    iEnsureDirOrThrow(fullfile(iPerWorkerRoot(), name));
    OK = true;
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function OK = iDeleteFolder(name)
    if parallel.internal.pool.LeadWorkers.isLeadWorker()
        OK1 = iRemoveDir(fullfile(iMyPerHostRoot(), name));
    else
        OK1 = true;
    end
    OK2 = iRemoveDir(fullfile(iPerWorkerRoot(), name));
    OK = OK1 && OK2;
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function d = iGetSubDirs(r)
    listing = dir(r);
    directories = listing([listing.isdir]);
    d = setdiff({directories.name}, {'.', '..'});
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% On a lead worker, we choose the intersection of the subdirectories of
% per-worker and per-host directories. Under normal circumstances, these should
% be identical.
function l = iListFolders()
    l = iGetSubDirs(iPerWorkerRoot());
    if parallel.internal.pool.LeadWorkers.isLeadWorker()
        ph = iGetSubDirs(iMyPerHostRoot());
        l = intersect(ph, l);
    end
end
