1 # Copyright 2018 The TensorFlow Authors. All Rights Reserved. 2 # 3 # Licensed under the Apache License, Version 2.0 (the "License"); 4 # you may not use this file except in compliance with the License. 5 # You may obtain a copy of the License at 6 # 7 # http://www.apache.org/licenses/LICENSE-2.0 8 # 9 # Unless required by applicable law or agreed to in writing, software 10 # distributed under the License is distributed on an "AS IS" BASIS, 11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 # See the License for the specific language governing permissions and 13 # limitations under the License. 14 # ============================================================================== 15 """Locking related utils.""" 16 17 from __future__ import absolute_import 18 from __future__ import division 19 from __future__ import print_function 20 21 import threading 22 23 24 class GroupLock(object): 25 """A lock to allow many members of a group to access a resource exclusively. 26 27 This lock provides a way to allow access to a resource by multiple threads 28 belonging to a logical group at the same time, while restricting access to 29 threads from all other groups. You can think of this as an extension of a 30 reader-writer lock, where you allow multiple writers at the same time. We 31 made it generic to support multiple groups instead of just two - readers and 32 writers. 33 34 Simple usage example with two groups accessing the same resource: 35 36 ```python 37 lock = GroupLock(num_groups=2) 38 39 # In a member of group 0: 40 with lock.group(0): 41 # do stuff, access the resource 42 # ... 43 44 # In a member of group 1: 45 with lock.group(1): 46 # do stuff, access the resource 47 # ... 48 ``` 49 50 Using as a context manager with `.group(group_id)` is the easiest way. You 51 can also use the `acquire` and `release` method directly. 52 """ 53 54 def __init__(self, num_groups=2): 55 """Initialize a group lock. 56 57 Args: 58 num_groups: The number of groups that will be accessing the resource under 59 consideration. Should be a positive number. 60 61 Returns: 62 A group lock that can then be used to synchronize code. 63 64 Raises: 65 ValueError: If num_groups is less than 1. 66 """ 67 if num_groups < 1: 68 raise ValueError("num_groups must be a positive integer, got {}".format( 69 num_groups)) 70 self._ready = threading.Condition(threading.Lock()) 71 self._num_groups = num_groups 72 self._group_member_counts = [0] * self._num_groups 73 74 def group(self, group_id): 75 """Enter a context where the lock is with group `group_id`. 76 77 Args: 78 group_id: The group for which to acquire and release the lock. 79 80 Returns: 81 A context manager which will acquire the lock for `group_id`. 82 """ 83 self._validate_group_id(group_id) 84 return self._Context(self, group_id) 85 86 def acquire(self, group_id): 87 """Acquire the group lock for a specific group `group_id`.""" 88 self._validate_group_id(group_id) 89 90 self._ready.acquire() 91 while self._another_group_active(group_id): 92 self._ready.wait() 93 self._group_member_counts[group_id] += 1 94 self._ready.release() 95 96 def release(self, group_id): 97 """Release the group lock for a specific group `group_id`.""" 98 self._validate_group_id(group_id) 99 100 self._ready.acquire() 101 self._group_member_counts[group_id] -= 1 102 if self._group_member_counts[group_id] == 0: 103 self._ready.notifyAll() 104 self._ready.release() 105 106 def _another_group_active(self, group_id): 107 return any( 108 c > 0 for g, c in enumerate(self._group_member_counts) if g != group_id) 109 110 def _validate_group_id(self, group_id): 111 if group_id < 0 or group_id >= self._num_groups: 112 raise ValueError( 113 "group_id={} should be between 0 and num_groups={}".format( 114 group_id, self._num_groups)) 115 116 class _Context(object): 117 """Context manager helper for `GroupLock`.""" 118 119 def __init__(self, lock, group_id): 120 self._lock = lock 121 self._group_id = group_id 122 123 def __enter__(self): 124 self._lock.acquire(self._group_id) 125 126 def __exit__(self, type_arg, value_arg, traceback_arg): 127 del type_arg, value_arg, traceback_arg 128 self._lock.release(self._group_id) 129
