Merge c06fe6a into c80bc96

azure/datalake/store/core.py

@@ -18,7 +18,6 @@
 import io
 import logging
 import sys
-import time
 import uuid
 
 
@@ -29,6 +28,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
@@ -328,11 +328,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):
         """
         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
         ----------
@@ -341,18 +341,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)
+        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):
         """
         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
         ----------
@@ -361,18 +364,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)
+        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):
         """
         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
         ----------
@@ -381,8 +388,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)
+        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,145 @@
+from concurrent.futures import ThreadPoolExecutor
+from .utils import CountUpDownLatch
+import threading
+import logging
+import multiprocessing
+import os
+import logging.handlers
+try:
+    from queue import Empty     # Python 3
+except ImportError:
+    from Queue import Empty     # Python 2
+log_sentinel = [None, None]
+
+
+def log_listener_process(queue):
+    while True:
+        try:
+            record = queue.get(timeout=0.1)
+            queue.task_done()
+            if record == log_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=""):
+    log_queue = multiprocessing.JoinableQueue()
+    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)))
+            process_list[-1].start()
+        return process_list
+
+    def walk(walk_path):
+        paths = []
+        all_files = adl.ls(path=walk_path, detail=True)
+        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 = []
+
+        file_path_queue.put(list(paths))  # For leftover paths < bucket_size
+        dir_processed_counter.decrement()         # Processing complete for this directory
+
+    finish_queue_processing_flag = multiprocessing.Event()
+    file_path_queue = multiprocessing.JoinableQueue()
+    cpu_count = multiprocessing.cpu_count()
+    child_processes = launch_processes(2)
+    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 to initialize walk
+    dir_processed_counter.increment()
+    walk(path)                  # Start processing root directory
+
+    if dir_processed_counter.is_zero():  # Done processing all directories. Blocking call.
+        file_path_queue.join()           # Wait for operations to be done
+        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, "Thread pool for worked threads for walk shut down")
+            child.join()
+
+    # Cleanup
+    logger.log(logging.DEBUG, "Sending logger sentinel")
+    log_queue.put(log_sentinel)
+    log_queue.join()
+    log_queue.close()
+    logger.log(logging.DEBUG, "Log queue closed")
+    log_listener.join()
+    logger.log(logging.DEBUG, "Log thread finished")
+    walk_thread_pool.shutdown()
+    logger.log(logging.DEBUG, "Thread pool for worked threads for walk shut down")
+    file_path_queue.close()
+    logger.log(logging.DEBUG, "File path queue closed")
+
+
+def processor(adl, file_path_queue, finish_queue_processing_flag, method_name, acl_spec, log_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}
+        running_thread_count = CountUpDownLatch()
+        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:
+                pass    # Exception is being logged in the relevant acl method. Do nothing here
+            running_thread_count.decrement()
+            logger.log(logging.DEBUG, "Completed running on path:" + str(path))
+
+        while not finish_queue_processing_flag.is_set():
+            try:
+                file_paths = file_path_queue.get(timeout=0.1)
+                file_path_queue.task_done()
+                for file_path in file_paths:
+                    running_thread_count.increment()
+                    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
+
+        if running_thread_count.is_zero():  # Blocking call. Will wait till all threads are finished.
+            pass
+        function_thread_pool.shutdown()
+    except Exception as e:
+        logger.exception("Exception in pid "+str(os.getpid())+"Exception: " + str(e))
+    finally:
+        function_thread_pool.shutdown()
+        logger.log(logging.DEBUG, "Finished processor pid: " + str(os.getpid()))

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

tests/settings.py

@@ -19,6 +19,7 @@
 SUBSCRIPTION_ID = fake_settings.SUBSCRIPTION_ID
 RESOURCE_GROUP_NAME = fake_settings.RESOURCE_GROUP_NAME
 RECORD_MODE = os.environ.get('RECORD_MODE', 'all').lower()
+AZURE_ACL_TEST_APPID = os.environ.get('AZURE_ACL_TEST_APPID')
 '''
 RECORD_MODE = os.environ.get('RECORD_MODE', 'none').lower()
 

tests/test_core.py

@@ -14,6 +14,8 @@
 from azure.datalake.store import utils
 from azure.datalake.store.exceptions import PermissionError, FileNotFoundError
 from tests.testing import azure, second_azure, azure_teardown, my_vcr, posix, tmpfile, working_dir, create_files
+from tests.settings import AZURE_ACL_TEST_APPID
+
 test_dir = working_dir()
 
 a = posix(test_dir / 'a')
@@ -1137,6 +1139,57 @@ def test_chown(azure):
 def test_acl_management(azure):
     pass
 
+
+@my_vcr.use_cassette
+def test_modify_acl_entries(azure):
+    with azure_teardown(azure):
+        acluser = AZURE_ACL_TEST_APPID
+        azure.touch(a)
+
+        permission = "---"
+        azure.modify_acl_entries(a, acl_spec="user:"+acluser+":"+permission)
+        current_acl = azure.get_acl_status(a)
+        aclspec = [s for s in current_acl['entries'] if acluser in s][0]
+        assert aclspec.split(':')[-1] == permission
+
+        permission = "rwx"
+        azure.modify_acl_entries(a, acl_spec="user:" + acluser + ":" + permission)
+        current_acl = azure.get_acl_status(a)
+        aclspec = [s for s in current_acl['entries'] if acluser in s][0]
+        assert aclspec.split(':')[-1] == permission
+
+
+@my_vcr.use_cassette
+def test_remove_acl_entries(azure):
+    with azure_teardown(azure):
+        acluser = AZURE_ACL_TEST_APPID
+        azure.touch(a)
+
+        permission = "rwx"
+        azure.modify_acl_entries(a, acl_spec="user:"+acluser+":"+permission)
+        current_acl = azure.get_acl_status(a)
+        aclspec = [s for s in current_acl['entries'] if acluser in s]
+        assert aclspec != []
+
+        azure.remove_acl_entries(a, acl_spec="user:" + acluser)
+        current_acl = azure.get_acl_status(a)
+        aclspec = [s for s in current_acl['entries'] if acluser in s]
+        assert aclspec == []
+
+@my_vcr.use_cassette
+def test_set_acl(azure):
+    with azure_teardown(azure):
+        acluser = AZURE_ACL_TEST_APPID
+        azure.touch(a)
+        set_acl_base ="user::rwx,group::rwx,other::---,"
+
+        permission = "rwx"
+        azure.set_acl(a, acl_spec=set_acl_base + "user:"+acluser+":"+permission)
+        current_acl = azure.get_acl_status(a)
+        aclspec = [s for s in current_acl['entries'] if acluser in s][0]
+        assert len(current_acl['entries']) == 5
+        assert aclspec.split(':')[-1] == permission
+
 @my_vcr.use_cassette
 def test_set_expiry(azure):
     with azure_teardown(azure):