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 17from __future__ import absolute_import 18from __future__ import division 19from __future__ import print_function 20 21import threading 22 23 24class 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 __slots__ = ["_ready", "_num_groups", "_group_member_counts"] 55 56 def __init__(self, num_groups=2): 57 """Initialize a group lock. 58 59 Args: 60 num_groups: The number of groups that will be accessing the resource under 61 consideration. Should be a positive number. 62 63 Returns: 64 A group lock that can then be used to synchronize code. 65 66 Raises: 67 ValueError: If num_groups is less than 1. 68 """ 69 if num_groups < 1: 70 raise ValueError( 71 "Argument `num_groups` must be a positive integer. " 72 f"Received: num_groups={num_groups}") 73 self._ready = threading.Condition(threading.Lock()) 74 self._num_groups = num_groups 75 self._group_member_counts = [0] * self._num_groups 76 77 def group(self, group_id): 78 """Enter a context where the lock is with group `group_id`. 79 80 Args: 81 group_id: The group for which to acquire and release the lock. 82 83 Returns: 84 A context manager which will acquire the lock for `group_id`. 85 """ 86 self._validate_group_id(group_id) 87 return self._Context(self, group_id) 88 89 def acquire(self, group_id): 90 """Acquire the group lock for a specific group `group_id`.""" 91 self._validate_group_id(group_id) 92 93 self._ready.acquire() 94 while self._another_group_active(group_id): 95 self._ready.wait() 96 self._group_member_counts[group_id] += 1 97 self._ready.release() 98 99 def release(self, group_id): 100 """Release the group lock for a specific group `group_id`.""" 101 self._validate_group_id(group_id) 102 103 self._ready.acquire() 104 self._group_member_counts[group_id] -= 1 105 if self._group_member_counts[group_id] == 0: 106 self._ready.notifyAll() 107 self._ready.release() 108 109 def _another_group_active(self, group_id): 110 return any( 111 c > 0 for g, c in enumerate(self._group_member_counts) if g != group_id) 112 113 def _validate_group_id(self, group_id): 114 if group_id < 0 or group_id >= self._num_groups: 115 raise ValueError( 116 "Argument `group_id` should verify `0 <= group_id < num_groups` " 117 f"(with `num_groups={self._num_groups}`). " 118 f"Received: group_id={group_id}") 119 120 class _Context(object): 121 """Context manager helper for `GroupLock`.""" 122 123 __slots__ = ["_lock", "_group_id"] 124 125 def __init__(self, lock, group_id): 126 self._lock = lock 127 self._group_id = group_id 128 129 def __enter__(self): 130 self._lock.acquire(self._group_id) 131 132 def __exit__(self, type_arg, value_arg, traceback_arg): 133 del type_arg, value_arg, traceback_arg 134 self._lock.release(self._group_id) 135