%FileManager Handle attached files for a parallel pool.

% Copyright 2013-2021 The MathWorks, Inc.

classdef FileManager < handle
    properties (SetAccess = immutable)
        Session % The Java session object
        Client % The MATLAB session object
    end
        
    properties (SetAccess = immutable, GetAccess = private)
        FileDependenciesAssistant
    end
       
    properties (SetAccess = private, GetAccess = public)
        AttachedFiles = {}     % Files that have been explicitly added by users
        AutoAttachedFiles = {} % Files that have been automatically added for users
        TransferCallbacks = {}
        TransferTimer = [];
    end
    
    % Public for testing only
    properties (Hidden)
        % All files that have been successfully attached to the pool,
        % including both user attached files and auto attached files
        AllAttachedFiles
    end
    
    properties (Access = private)
        LastUpdateTime
    end
    
    methods
        function obj = FileManager(client)
            obj.Client = client;
            obj.Session = client.Session;
            obj.FileDependenciesAssistant = obj.Session.getFileDependenciesAssistant();
            obj.LastUpdateTime = iGetTime();
            
            % Initialize all attached files to be all attached files on the
            % job; these have been "added" for us already
            obj.AllAttachedFiles = obj.Client.ParallelJob.hGetAllAttachedFiles();
        end
        
        % Called when a user edits a file on the client that has been
        % attached to an open pool
        function maybeUpdateAttachedFiles(obj, files, asyncWait)
            iErrorIfSessionNotRunning(obj.Session);
            if nargin < 3
                asyncWait = false;
            end
            if ~iscell(files)
                files = {files};
            end
            [~, filesToUpdate] = obj.getFilesToAddAndUpdate(files);
            if isempty(filesToUpdate)
                return
            end
            
            % Only update those files that have already been added
            changedFilesMap = obj.findSpecificChangedFiles(filesToUpdate);
            obj.doUpdateAttachedFiles(changedFilesMap, asyncWait);
        end
        
        % Called when pool workers cannot find files on their path and we
        % want to try to automatically attach them
        function addOrUpdateAutoAttachedFiles(obj, files)
            iErrorIfSessionNotRunning(obj.Session);
            if ~iscell(files)
                files = {files};
            end

            % Ensure we don't specify the same file multiple times
            files = unique(files);
            [filesToAdd, filesToUpdate] = obj.getFilesToAddAndUpdate(files);
            
            if ~isempty(filesToAdd)
                % Ensure we have a column cellstr of file names
                filesToAdd = filesToAdd(:);
                areAutoAttachedFiles = true;
                obj.addAttachedFiles(filesToAdd, areAutoAttachedFiles);
            end
            
            if ~isempty(filesToUpdate)
                % Only update those files that have already been added
                changedFilesMap = obj.findSpecificChangedFiles(filesToUpdate);
                obj.doUpdateAttachedFiles(changedFilesMap);
            end
        end
        
        % Called when a user explicitly adds files to a pool with
        % pool.addAttachedFiles() (or internally from
        % addOrUpdateAutoAttachedFiles())
        function successfullyAttachedFiles = addAttachedFiles(obj, files, areAutoAttachedFiles)
            if nargin < 3
                areAutoAttachedFiles = false;
            end
            
            iErrorIfSessionNotRunning(obj.Session);
            
            % Canonicalize the list of files (removes empty files, i.e. '')
            files = distcomp.pCanonicalizeFileDependenciesList(files);

            % Remove duplicate file names
            files = unique(files, 'stable');
            
            % Check each file to add correctly references a file or directory
            for ii = 1:numel(files)
                if ~exist(files{ii}, 'file')
                    error(message('parallel:lang:pool:InvalidDependency', files{ii}));
                end
            end
            
            % Filter out files that have already been added
            files = obj.getFilesToAddAndUpdate(files);
            numFiles = numel(files);
            
            % Quick exit if there are no files to add
            if numFiles == 0
                successfullyAttachedFiles = cell(0, 1);
                return
            end
            
            % Work out where the remote files are supposed to go.
            job = obj.Client.ParallelJob;
            % Note that even though pathMap is a handle, we must explicitly set it again
            % on the job so that it gets serialized correctly.
            [relativeWorkerPaths, pathMap] = parallel.internal.apishared.AttachedFiles.generateTempDirectoryStructure(files, job.AttachedFilePaths);
            job.AttachedFilePaths = pathMap;
            
            % Create a cell array of completion observers
            observers = cell(numFiles, 1);
            
            % Loop backwards because we expect the first file in the list
            % to be at the top of the path, not the end
            for ii = numFiles:-1:1
                file = files{ii};
                relativeWorkerPath = relativeWorkerPaths{ii};
                
                if isfolder(file)
                    % Send directories by zipping them up and sending the zip. Note
                    % that the cleanup object will take responsibility for removing the zip file
                    % once the message has been received by all
                    [tempFilesDir, tempDirCleanupObj] = parallel.internal.apishared.AttachedFiles.copyAllFilesToTempDir({file}, {relativeWorkerPath}); %#ok<ASGLU>
                    [zipName, cleanupObject] = parallel.internal.apishared.AttachedFiles.createZipFile(relativeWorkerPath, tempFilesDir); %#ok<ASGLU>
                    
                    % clientFolders and workerFolders map client to worker folders
                    [clientFolders, workerFolders] = iRetrieveMapFoldersUnderTree(pathMap, file);
                    observer = obj.FileDependenciesAssistant.addNewDirectoryDependency(zipName, file, relativeWorkerPath, clientFolders, workerFolders);
                else
                    % Send files as is.
                    observer = obj.FileDependenciesAssistant.addNewFileDependency(file, relativeWorkerPath);
                end
                % Make sure we retain the observer for each requested transfer
                observers{ii} = observer;
            end
            observers = [observers{:}];
            [transferErrors, erroredFiles] = iWaitForTransfersToComplete(observers);
            
            files = files(:); % Ensure we have a column cellstr of file names
            successfullyAttachedFiles = setdiff(files, erroredFiles, 'stable');
            
            % Add successfully attached files to the list of all attached files
            obj.AllAttachedFiles = [obj.AllAttachedFiles; successfullyAttachedFiles];
            
            % Keep track of the successfully added files in the appropriate list
            if areAutoAttachedFiles
                obj.AutoAttachedFiles = [obj.AutoAttachedFiles; successfullyAttachedFiles];
            else
                obj.AttachedFiles = [obj.AttachedFiles; successfullyAttachedFiles];
            end
            
            if ~isempty(transferErrors)
                errorToThrow = MException(message('parallel:lang:pool:FileTransferError'));
                for ii = 1:numel(transferErrors)
                    errorToThrow = errorToThrow.addCause(MException(message('parallel:lang:pool:UnableToCompleteTransfer', erroredFiles{ii}, transferErrors{ii})));
                end
                throw(errorToThrow);
            end
        end
        
        % Called when a user explicitly updates a pool's attached files
        % with pool.updateAttachedFiles()
        function updateAttachedFiles(obj)
            % Do nothing if no session is running, simply return
            session = obj.Session;
            if ~iIsSessionRunning(session)
                return
            end          
            % Find files that have changed since we last looked
            
            mapFiles = obj.findChangedFiles();
            obj.doUpdateAttachedFiles(mapFiles);
        end
        
        % This method should ONLY be called on the timer - NOTE needs to be
        % public so that it can be called as a timer callback
        function executeTransferCallbacks(obj)
            fcns = obj.TransferCallbacks;
            % Iterate backwards so we can remove correctly without messing
            % up the elements earlier in the list
            for i = numel(fcns):-1:1
                try
                    % Execute this callback
                    complete = fcns{i}();
                    if complete                        
                        obj.TransferCallbacks(i) = [];
                    end
                catch me %#ok<NASGU>
                    obj.TransferCallbacks(i) = [];
                end
            end
            % Nothing left to do? Stop & delete the timer running.
            if isempty(obj.TransferCallbacks)
                stop(obj.TransferTimer);
                delete(obj.TransferTimer);
                obj.TransferTimer = [];
            end
        end
    end
    
    methods (Access = private)
        % Returns whether the given (canonicalized) file or files have been
        % attached to the pool
        function tf = isAttachedFile(obj, canonicalizedFileOrFiles)
            tf = startsWith(canonicalizedFileOrFiles, obj.AllAttachedFiles);
        end
        
        % Searches through all attached files and returns a map of the
        % attached files that have changed since the last update time to
        % their corresponding (relative) paths on the workers
        function changedFilesMap = findChangedFiles(obj)
            changedFilesMap = obj.findSpecificChangedFiles(obj.AllAttachedFiles);
        end
        
        % Searches through the given attached files and returns a map of
        % those that have changed since the last update time to their
        % corresponding (relative) paths on the workers
        function changedFilesMap = findSpecificChangedFiles(obj, attachedFiles)
            changedFilesMap = cell(0, 2);
            
            pathMap = [];
            for ii = 1:numel(attachedFiles)
                attachedFile = attachedFiles{ii};
                changedFiles = obj.doFindChangedFiles(attachedFile);
                if ~isempty(changedFiles)
                    % Lazily retrieve path map from the job
                    if isempty(pathMap)
                        pathMap = obj.Client.ParallelJob.AttachedFilePaths;
                    end
                    containingAttachedFile = obj.getContainingAttachedFile(attachedFile);
                    workerPath = pathMap.get(containingAttachedFile);
                    relativeWorkerPaths = strrep(changedFiles, containingAttachedFile, workerPath);
                    changedFilesMap = [changedFilesMap; [changedFiles', relativeWorkerPaths']]; %#ok<AGROW>
                end
            end
        end
        
        % Recursively searches the given attached file for all files that
        % have changed since the last update time and returns them
        function changedFiles = doFindChangedFiles(obj, attachedFile)
            % For a file, just check whether it has changed
            if isfile(attachedFile)
                if obj.hasChanged(attachedFile)
                    changedFiles = {attachedFile};
                else
                    changedFiles = {};
                end
                return
            end
            
            % For a folder, recursively find all changed files in the
            % folder or any subfolder
            changedFiles = {};
            folders = {attachedFile};
            while ~isempty(folders)
                newFolders = {};
                for ii = 1:numel(folders)
                    [files, subFolders] = iGetFilesAndFolders(folders{ii});
                    
                    % Add changed files to the list
                    hasChanged = false(numel(files));
                    for jj = 1:numel(files)
                        hasChanged(jj) = obj.hasChanged(files{jj});
                    end
                    changedFiles = [changedFiles, files(hasChanged)]; %#ok<AGROW>
                    
                    % Add new folders to check to the list
                    newFolders = [newFolders, subFolders]; %#ok<AGROW>
                end
                
                % Update 'folders' for the next loop iteration
                folders = newFolders;
            end
        end
        
        % Returns whether the given file has changed since the last update
        function tf = hasChanged(obj, file)
            tf = parallel.internal.filesystem.getModifyTime(file) + 1 >= obj.LastUpdateTime ...
                || parallel.internal.filesystem.getChangeTime(file) + 1 >= obj.LastUpdateTime;
        end
        
        % Returns either the given file itself or the attached folder that
        % contains the given file
        function attachedFile = getContainingAttachedFile(obj, file)
            for ii = 1:numel(obj.AllAttachedFiles)
                if startsWith(file, obj.AllAttachedFiles{ii})
                    attachedFile = obj.AllAttachedFiles{ii};
                    return
                end
            end
            assert(false, "Unable to find containing attached file for: " + file);
        end
        
        function addTransferCallback(obj, callback)
            % Add this transfer watcher to the list of callbacks
            obj.TransferCallbacks{end+1} = callback;
            % Make sure that the TransferCallbackTimer is running
            if isempty(obj.TransferTimer) || ~isvalid(obj.TransferTimer)
                % Create a timer that will eventually delete itself. We only do
                % this if there is no existing timer running.
                obj.TransferTimer = timer('TimerFcn', @(~,~) executeTransferCallbacks(obj), ...
                    'Period', 1, 'ExecutionMode', 'fixedRate', 'ObjectVisibility', 'on', ...
                    'BusyMode', 'drop');
                start(obj.TransferTimer);
            end
        end
        
        function doUpdateAttachedFiles(obj, mapFiles, asyncWait)
            if nargin < 3
                asyncWait = false;
            end
            
            if isempty(mapFiles)
                % If there is nothing to change then we can return early
                return
            end
            lastUpdateTime = iGetTime();
            % Ask the FileDependenciesAssistant to send this list of changed files -
            % getting back a list of observers to those sends
            observer = obj.FileDependenciesAssistant.updateFileDependencies(mapFiles(:,1), mapFiles(:,2));
            if asyncWait
                obj.asyncTransferCompletion(observer, lastUpdateTime);
            else
                obj.syncTransferCompletion(observer, lastUpdateTime);
            end
        end
        
        % Canonicalizes the given files and returns those that should be
        % added to the pool (i.e. haven't been attached to the pool) and
        % those that should be updated (i.e. have been attached to the pool)
        function [filesToAdd, filesToUpdate] = getFilesToAddAndUpdate(obj, files)
            files = cellfun(@parallel.internal.apishared.FilenameUtils.getCanonicalPath, files, 'UniformOutput', false);
            isAttachedFile = obj.isAttachedFile(files);
            filesToAdd = files(~isAttachedFile);
            filesToUpdate = files(isAttachedFile);
        end
        
        function syncTransferCompletion(obj, observers, lastUpdateTime)
            [transferErrors, erroredFilenames] = iWaitForTransfersToComplete(observers);
        
            % If all were successful, update the last check time
            obj.LastUpdateTime = lastUpdateTime;
            iWarnOnTransferError(transferErrors, erroredFilenames);
        end
        
        % Arranges for a timer to perform actions once transfer is complete
        function asyncTransferCompletion(obj, observers, lastUpdateTime)
            % Add the nested function below to be called periodically on
            % the timer
            obj.addTransferCallback(@nAsyncTransferTimerFcn);
            
            function transferComplete = nAsyncTransferTimerFcn
                transferComplete = all(arrayfun(@isComplete, observers));
                if ~transferComplete
                    return
                end
                
                [transferErrors, erroredFilenames] = iGetTransferErrors(observers);
                
                % If all was successful then update the time that we last checked at.
                obj.LastUpdateTime = lastUpdateTime;
                iWarnOnTransferError(transferErrors, erroredFilenames);
            end
        end
    end
end

% Get all the subfolders stored in the map which are subfolders of
% rootFolder. The first clientFolders cell contains the base folder
% originally added by the user (i.e., not automatically added). For
% instance, if the user adds 'a/b' but 'a/b/c', 'a/b/d', 'a/b/d/e' are on
% the path, clientFolders{1} always contains 'a/b'
function [clientFolders, workerFolders] = iRetrieveMapFoldersUnderTree(pathMap, rootFolder)
    % Make sure pathMap is accessed by absolute paths
    import parallel.internal.apishared.FilenameUtils

    keys = pathMap.keys();
    values = pathMap.values();

    canonicalRootPathForwardSlashes = strrep(FilenameUtils.getCanonicalPath(rootFolder), '\', '/');
    indexes = contains(keys, canonicalRootPathForwardSlashes);

    % Build local and remote list of folders from the map keys and values,
    % respectively
    clientFolders = keys(indexes);
    workerFolders = values(indexes);
end

function [transferErrors, erroredFiles] = iWaitForTransfersToComplete(observers)
    % Wait for all transfers to complete
    for ii = 1:numel(observers)
        while ~observers(ii).waitForCompletion(1000) % Wait for 1 second intervals to allow for interrupts
        end
    end
    
    [transferErrors, erroredFiles] = iGetTransferErrors(observers);
end

function [transferErrors, erroredFilenames] = iGetTransferErrors(observers)
    numObservers = numel(observers);
    transferErrors = cell(numObservers, 1);
    erroredFilenames = cell(numObservers, 1);

    for ii = 1:numObservers
        % Get the transfer errors from this file send to see if we
        % should indicate something went wrong
        transferErrors{ii} = iGetTransferError(observers(ii));
        if ~isempty(transferErrors{ii})
            % Which file were we trying to send across.
            erroredFilenames{ii} = char(observers(ii).getFile());
        end
    end

    empties = cellfun(@isempty, transferErrors);
    transferErrors = transferErrors(~empties);
    erroredFilenames = erroredFilenames(~empties);
end

function err = iGetTransferError(fileTransferObserver)
    % Get the transfer errors for this file send to see if we should
    % indicate something went wrong.
    transferErrors = cellstr(string(fileTransferObserver.getErrorMessages()));
    if isempty(transferErrors)
        err = '';
    else
        % Concatenate the transfer error messages for all failed destinations.
        % Each transfer error message should contain the destination and some
        % explanation as to what went wrong.
        err = sprintf('\n%s', transferErrors{:});
    end
end

function iWarnOnTransferError(transferErrors, erroredFilenames)
    if ~isempty(transferErrors)
        for ii = 1:numel(transferErrors)
            warning(message('parallel:lang:pool:UnableToCompleteUpdate', ...
                erroredFilenames{ii}, transferErrors{ii}));
        end
    end
end

function iErrorIfSessionNotRunning(session)
    if ~iIsSessionRunning(session)
        error(message('parallel:lang:pool:PoolNotRunning'));
    end
end

function tf = iIsSessionRunning(session)
    tf = ~isempty(session) && isvalid(session) && session.isSessionRunning();
end

function time = iGetTime()
    time = posixtime(datetime('now', 'TimeZone', 'local'));
end

% Returns the full path of all files and folders (excluding '.' and '..')
% in the given folder
function [files, folders] = iGetFilesAndFolders(folder)
    dirContents = dir(folder);
    dirContents(ismember({dirContents.name}, {'.', '..'})) = []; % Exclude '.' and '..'
    allFiles = cellstr(folder + string(filesep) + {dirContents.name});
    isdir = [dirContents.isdir];
    files = allFiles(~isdir);
    folders = allFiles(isdir);
end
