Merge 37dca16 into 0adf4dd

azure/datalake/store/core.py

@@ -18,7 +18,6 @@
 import io
 import logging
 import sys
-import time
 import uuid
 import json
 
@@ -30,6 +29,7 @@
 from .utils import ensure_writable, read_block
 from .enums import ExpiryOptionType
 from .retry import ExponentialRetryPolicy
+from .multiprocessor import multi_processor_change_acl
 
 if sys.version_info >= (3, 4):
     import pathlib
@@ -329,11 +329,11 @@ def _acl_call(self, action, path, acl_spec=None, invalidate_cache=False):
 
         return to_return
 
-    def set_acl(self, path, acl_spec):
+    def set_acl(self, path, acl_spec, recursive=False, number_of_sub_process=None):
         """
         Sets the Access Control List (ACL) for a file or folder.
 
-        Note: this is not recursive, and applies only to the file or folder specified.
+        Note: this is by default not recursive, and applies only to the file or folder specified.
 
         Parameters
         ----------
@@ -342,18 +342,21 @@ def set_acl(self, path, acl_spec):
         acl_spec: str
             The ACL specification to set on the path in the format
             '[default:]user|group|other:[entity id or UPN]:r|-w|-x|-,[default:]user|group|other:[entity id or UPN]:r|-w|-x|-,...'
+        recursive: bool
+            Specifies whether to set ACLs recursively or not
         """
+        if recursive:
+            multi_processor_change_acl(adl=self, path=path, method_name="set_acl", acl_spec=acl_spec, number_of_sub_process=number_of_sub_process)
+        else:
+            self._acl_call('SETACL', path, acl_spec, invalidate_cache=True)
 
-        self._acl_call('SETACL', path, acl_spec, invalidate_cache=True)
-
-
-    def modify_acl_entries(self, path, acl_spec):
+    def modify_acl_entries(self, path, acl_spec, recursive=False, number_of_sub_process=None):
         """
         Modifies existing Access Control List (ACL) entries on a file or folder.
         If the entry does not exist it is added, otherwise it is updated based on the spec passed in.
         No entries are removed by this process (unlike set_acl).
 
-        Note: this is not recursive, and applies only to the file or folder specified.
+        Note: this is by default not recursive, and applies only to the file or folder specified.
 
         Parameters
         ----------
@@ -362,18 +365,22 @@ def modify_acl_entries(self, path, acl_spec):
         acl_spec: str
             The ACL specification to use in modifying the ACL at the path in the format
             '[default:]user|group|other:[entity id or UPN]:r|-w|-x|-,[default:]user|group|other:[entity id or UPN]:r|-w|-x|-,...'
+        recursive: bool
+            Specifies whether to modify ACLs recursively or not
         """
-        self._acl_call('MODIFYACLENTRIES', path, acl_spec, invalidate_cache=True)
-
+        if recursive:
+            multi_processor_change_acl(adl=self, path=path, method_name="mod_acl", acl_spec=acl_spec, number_of_sub_process=number_of_sub_process)
+        else:
+            self._acl_call('MODIFYACLENTRIES', path, acl_spec, invalidate_cache=True)
 
-    def remove_acl_entries(self, path, acl_spec):
+    def remove_acl_entries(self, path, acl_spec, recursive=False, number_of_sub_process=None):
         """
         Removes existing, named, Access Control List (ACL) entries on a file or folder.
         If the entry does not exist already it is ignored.
         Default entries cannot be removed this way, please use remove_default_acl for that.
         Unnamed entries cannot be removed in this way, please use remove_acl for that.
 
-        Note: this is not recursive, and applies only to the file or folder specified.
+        Note: this is by default not recursive, and applies only to the file or folder specified.
 
         Parameters
         ----------
@@ -382,8 +389,13 @@ def remove_acl_entries(self, path, acl_spec):
         acl_spec: str
             The ACL specification to remove from the ACL at the path in the format (note that the permission portion is missing)
             '[default:]user|group|other:[entity id or UPN],[default:]user|group|other:[entity id or UPN],...'
+        recursive: bool
+            Specifies whether to remove ACLs recursively or not
         """
-        self._acl_call('REMOVEACLENTRIES', path, acl_spec, invalidate_cache=True)
+        if recursive:
+            multi_processor_change_acl(adl=self, path=path, method_name="rem_acl", acl_spec=acl_spec, number_of_sub_process=number_of_sub_process)
+        else:
+            self._acl_call('REMOVEACLENTRIES', path, acl_spec, invalidate_cache=True)
 
 
     def get_acl_status(self, path):

azure/datalake/store/multiprocessor.py

@@ -0,0 +1,197 @@
+from concurrent.futures import ThreadPoolExecutor
+from .utils import CountUpDownLatch
+import threading
+import logging
+import multiprocessing
+import os
+import logging.handlers
+from .exceptions import  FileNotFoundError
+
+
+try:
+    from queue import Empty     # Python 3
+except ImportError:
+    from Queue import Empty     # Python 2
+end_queue_sentinel = [None, None]
+
+exception = None
+exception_lock = threading.Lock()
+
+
+threading
+def monitor_exception(exception_queue, process_ids):
+    global exception
+    logger = logging.getLogger(__name__)
+    logger.setLevel(logging.DEBUG)
+
+    while True:
+        try:
+            excep = exception_queue.get(timeout=0.1)
+            if excep == end_queue_sentinel:
+                break
+            logger.log(logging.DEBUG, "Setting global exception")
+            exception_lock.acquire()
+            exception = excep
+            exception_lock.release()
+            logger.log(logging.DEBUG, "Closing processes")
+            for p in process_ids:
+                p.terminate()
+            logger.log(logging.DEBUG, "Joining processes")
+            for p in process_ids:
+                p.join()
+            import thread
+            logger.log(logging.DEBUG, "Interrupting main")
+            raise Exception(excep)
+        except Empty:
+            pass
+
+
+def log_listener_process(queue):
+    while True:
+        try:
+            record = queue.get(timeout=0.1)
+            queue.task_done()
+            if record == end_queue_sentinel:  # We send this as a sentinel to tell the listener to quit.
+                break
+            logger = logging.getLogger(record.name)
+            logger.handlers.clear()
+            logger.handle(record)  # No level or filter logic applied - just do it!
+        except Empty:               # Try again
+            pass
+        except Exception as e:
+            import sys, traceback
+            print('Problems in logging')
+            traceback.print_exc(file=sys.stderr)
+
+
+def multi_processor_change_acl(adl, path=None, method_name="", acl_spec="", number_of_sub_process=None):
+    log_queue = multiprocessing.JoinableQueue()
+    exception_queue = multiprocessing.Queue()
+    logger = logging.getLogger(__name__)
+    logger.setLevel(logging.DEBUG)
+    queue_bucket_size = 10
+    worker_thread_num_per_process = 50
+
+    def launch_processes(number_of_processes):
+        process_list = []
+        for i in range(number_of_processes):
+            process_list.append(multiprocessing.Process(target=processor,
+                                    args=(adl, file_path_queue, finish_queue_processing_flag,
+                                          method_name, acl_spec, log_queue, exception_queue)))
+            process_list[-1].start()
+        return process_list
+
+    def walk(walk_path):
+        try:
+            paths = []
+            all_files = adl._ls(path=walk_path)
+
+            for files in all_files:
+                if files['type'] == 'DIRECTORY':
+                    dir_processed_counter.increment()               # A new directory to process
+                    walk_thread_pool.submit(walk, files['name'])
+                paths.append(files['name'])
+                if len(paths) == queue_bucket_size:
+                    file_path_queue.put(list(paths))
+                    paths = []
+            if paths != []:
+                file_path_queue.put(list(paths))  # For leftover paths < bucket_size
+        except FileNotFoundError:
+            pass                    # Continue in case the file was deleted in between
+        except:
+            import traceback
+            logger.exception("Failed to walk for path: " + str(walk_path) + ". Exiting!")
+            exception_queue.put(traceback.format_exc())
+        finally:
+            dir_processed_counter.decrement()           # Processing complete for this directory
+
+    finish_queue_processing_flag = multiprocessing.Event()
+    file_path_queue = multiprocessing.JoinableQueue()
+    if number_of_sub_process == None:
+        number_of_sub_process = max(2, multiprocessing.cpu_count()-1)
+
+    child_processes = launch_processes(number_of_sub_process)
+    exception_monitor_thread = threading.Thread(target=monitor_exception, args=(exception_queue, child_processes))
+    exception_monitor_thread.start()
+    log_listener = threading.Thread(target=log_listener_process, args=(log_queue,))
+    log_listener.start()
+
+    dir_processed_counter = CountUpDownLatch()
+    walk_thread_pool = ThreadPoolExecutor(max_workers=worker_thread_num_per_process)
+
+    file_path_queue.put([path])  # Root directory needs to be passed
+    dir_processed_counter.increment()
+    walk(path)                  # Start processing root directory
+
+    if dir_processed_counter.is_zero(): # Done processing all directories. Blocking call.
+        walk_thread_pool.shutdown()
+        file_path_queue.close()          # No new elements to add
+        file_path_queue.join()           # Wait for operations to be done
+        logger.log(logging.DEBUG, "file path queue closed")
+        finish_queue_processing_flag.set()  # Set flag to break loop of child processes
+        for child in child_processes:  # Wait for all child process to finish
+            logger.log(logging.DEBUG, "Joining process: "+str(child.pid))
+            child.join()
+
+    # Cleanup
+    logger.log(logging.DEBUG, "Sending exception sentinel")
+    exception_queue.put(end_queue_sentinel)
+    exception_monitor_thread.join()
+    logger.log(logging.DEBUG, "Exception monitor thread finished")
+    logger.log(logging.DEBUG, "Sending logger sentinel")
+    log_queue.put(end_queue_sentinel)
+    log_queue.join()
+    log_queue.close()
+    logger.log(logging.DEBUG, "Log queue closed")
+    log_listener.join()
+    logger.log(logging.DEBUG, "Log thread finished")
+
+
+def processor(adl, file_path_queue, finish_queue_processing_flag, method_name, acl_spec, log_queue, exception_queue):
+    logger = logging.getLogger(__name__)
+
+    try:
+        logger.addHandler(logging.handlers.QueueHandler(log_queue))
+        logger.propagate = False                                                        # Prevents double logging
+    except AttributeError:
+        # Python 2 doesn't have Queue Handler. Default to best effort logging.
+        pass
+    logger.setLevel(logging.DEBUG)
+
+    try:
+        worker_thread_num_per_process = 50
+        func_table = {"mod_acl": adl.modify_acl_entries, "set_acl": adl.set_acl, "rem_acl": adl.remove_acl_entries}
+        function_thread_pool = ThreadPoolExecutor(max_workers=worker_thread_num_per_process)
+        adl_function = func_table[method_name]
+        logger.log(logging.DEBUG, "Started processor pid:"+str(os.getpid()))
+
+        def func_wrapper(func, path, spec):
+            try:
+                func(path=path, acl_spec=spec)
+            except FileNotFoundError as e:
+                logger.exception("File "+str(path)+" not found")
+                pass    # Exception is being logged in the relevant acl method. Do nothing here
+            except:
+                # TODO Raise to parent process
+                pass
+
+            logger.log(logging.DEBUG, "Completed running on path:" + str(path))
+
+        while finish_queue_processing_flag.is_set() == False:
+            try:
+                file_paths = file_path_queue.get(timeout=0.1)
+                file_path_queue.task_done()                 # Will not be called if empty
+                for file_path in file_paths:
+                    logger.log(logging.DEBUG, "Starting on path:" + str(file_path))
+                    function_thread_pool.submit(func_wrapper, adl_function, file_path, acl_spec)
+            except Empty:
+                pass
+
+    except Exception as e:
+        import traceback
+        # TODO Raise to parent process
+        logger.exception("Exception in pid "+str(os.getpid())+"Exception: " + str(e))
+        exception_queue.put(traceback.format_exc())
+    finally:
+        function_thread_pool.shutdown()  # Blocking call. Will wait till all threads are done executing.
+        logger.log(logging.DEBUG, "Finished processor pid: " + str(os.getpid()))

azure/datalake/store/multithread.py

@@ -21,12 +21,14 @@
 import pickle
 import time
 import errno
+import uuid
 
 from io import open
 from .core import AzureDLPath, _fetch_range
 from .exceptions import FileExistsError, FileNotFoundError
 from .transfer import ADLTransferClient
 from .utils import datadir, read_block, tokenize
+from .retry import ExponentialRetryPolicy
 
 logger = logging.getLogger(__name__)
 
@@ -284,9 +286,9 @@ def get_chunk(adlfs, src, dst, offset, size, buffersize, blocksize,
     """
     err = None
     total_bytes_downloaded = 0
-    from azure.datalake.store.retry import ExponentialRetryPolicy
     retry_policy = ExponentialRetryPolicy(max_retries=retries, exponential_retry_interval=delay,
                                           exponential_factor=backoff)
+    filesessionid = str(uuid.uuid4())
     try:
         nbytes = 0
         start = offset
@@ -295,7 +297,8 @@ def get_chunk(adlfs, src, dst, offset, size, buffersize, blocksize,
             fout.seek(start)
             while start < offset+size:
                 with closing(_fetch_range(adlfs.azure, src, start=start,
-                                          end=min(start+blocksize, offset+size), stream=True, retry_policy=retry_policy)) as response:
+                                          end=min(start+blocksize, offset+size), stream=True,
+                                          retry_policy=retry_policy, filesessionid=filesessionid)) as response:
                     chunk = response.content
                     if shutdown_event and shutdown_event.is_set():
                         return total_bytes_downloaded, None

azure/datalake/store/utils.py

@@ -11,6 +11,7 @@
 import os
 import platform
 import sys
+import threading
 
 PY2 = sys.version_info.major == 2
 
@@ -158,3 +159,38 @@ def clamp(n, smallest, largest):
     32
     """
     return max(smallest, min(n, largest))
+
+
+class CountUpDownLatch:
+    """CountUpDownLatch provides a thread safe implementation of Up Down latch
+    """
+    def __init__(self):
+        self.lock = threading.Condition()
+        self.val = 0
+        self.total = 0
+
+    def increment(self):
+        self.lock.acquire()
+        self.val += 1
+        self.total += 1
+        self.lock.release()
+
+    def decrement(self):
+        self.lock.acquire()
+        self.val -= 1
+        if self.val <= 0:
+            self.lock.notifyAll()
+        self.lock.release()
+
+    def total_processed(self):
+        self.lock.acquire()
+        temp = self.total
+        self.lock.release()
+        return temp
+
+    def is_zero(self):
+        self.lock.acquire()
+        while self.val > 0:
+            self.lock.wait()
+        self.lock.release()
+        return True