OLD | NEW |
| (Empty) |
1 # -*- coding: utf-8 -*- | |
2 # Copyright 2013 Google Inc. All Rights Reserved. | |
3 # | |
4 # Licensed under the Apache License, Version 2.0 (the "License"); | |
5 # you may not use this file except in compliance with the License. | |
6 # You may obtain a copy of the License at | |
7 # | |
8 # http://www.apache.org/licenses/LICENSE-2.0 | |
9 # | |
10 # Unless required by applicable law or agreed to in writing, software | |
11 # distributed under the License is distributed on an "AS IS" BASIS, | |
12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
13 # See the License for the specific language governing permissions and | |
14 # limitations under the License. | |
15 """Utility classes for the parallelism framework.""" | |
16 | |
17 from __future__ import absolute_import | |
18 | |
19 import multiprocessing | |
20 import threading | |
21 | |
22 | |
23 class BasicIncrementDict(object): | |
24 """Dictionary meant for storing values for which increment is defined. | |
25 | |
26 This handles any values for which the "+" operation is defined (e.g., floats, | |
27 lists, etc.). This class is neither thread- nor process-safe. | |
28 """ | |
29 | |
30 def __init__(self): | |
31 self.dict = {} | |
32 | |
33 def Get(self, key, default_value=None): | |
34 return self.dict.get(key, default_value) | |
35 | |
36 def Put(self, key, value): | |
37 self.dict[key] = value | |
38 | |
39 def Update(self, key, inc, default_value=0): | |
40 """Update the stored value associated with the given key. | |
41 | |
42 Performs the equivalent of | |
43 self.put(key, self.get(key, default_value) + inc). | |
44 | |
45 Args: | |
46 key: lookup key for the value of the first operand of the "+" operation. | |
47 inc: Second operand of the "+" operation. | |
48 default_value: Default value if there is no existing value for the key. | |
49 | |
50 Returns: | |
51 Incremented value. | |
52 """ | |
53 val = self.dict.get(key, default_value) + inc | |
54 self.dict[key] = val | |
55 return val | |
56 | |
57 | |
58 class AtomicIncrementDict(BasicIncrementDict): | |
59 """Dictionary meant for storing values for which increment is defined. | |
60 | |
61 This handles any values for which the "+" operation is defined (e.g., floats, | |
62 lists, etc.) in a thread- and process-safe way that allows for atomic get, | |
63 put, and update. | |
64 """ | |
65 | |
66 def __init__(self, manager): # pylint: disable=super-init-not-called | |
67 self.dict = ThreadAndProcessSafeDict(manager) | |
68 self.lock = multiprocessing.Lock() | |
69 | |
70 def Update(self, key, inc, default_value=0): | |
71 """Atomically update the stored value associated with the given key. | |
72 | |
73 Performs the atomic equivalent of | |
74 self.put(key, self.get(key, default_value) + inc). | |
75 | |
76 Args: | |
77 key: lookup key for the value of the first operand of the "+" operation. | |
78 inc: Second operand of the "+" operation. | |
79 default_value: Default value if there is no existing value for the key. | |
80 | |
81 Returns: | |
82 Incremented value. | |
83 """ | |
84 with self.lock: | |
85 return super(AtomicIncrementDict, self).Update(key, inc, default_value) | |
86 | |
87 | |
88 class ThreadSafeDict(object): | |
89 """Provides a thread-safe dictionary (protected by a lock).""" | |
90 | |
91 def __init__(self): | |
92 """Initializes the thread-safe dict.""" | |
93 self.lock = threading.Lock() | |
94 self.dict = {} | |
95 | |
96 def __getitem__(self, key): | |
97 with self.lock: | |
98 return self.dict[key] | |
99 | |
100 def __setitem__(self, key, value): | |
101 with self.lock: | |
102 self.dict[key] = value | |
103 | |
104 # pylint: disable=invalid-name | |
105 def get(self, key, default_value=None): | |
106 with self.lock: | |
107 return self.dict.get(key, default_value) | |
108 | |
109 def delete(self, key): | |
110 with self.lock: | |
111 del self.dict[key] | |
112 | |
113 | |
114 class ThreadAndProcessSafeDict(ThreadSafeDict): | |
115 """Wraps a multiprocessing.Manager's proxy objects for thread-safety. | |
116 | |
117 The proxy objects returned by a manager are process-safe but not necessarily | |
118 thread-safe, so this class simply wraps their access with a lock for ease of | |
119 use. Since the objects are process-safe, we can use the more efficient | |
120 threading Lock. | |
121 """ | |
122 | |
123 def __init__(self, manager): | |
124 """Initializes the thread and process safe dict. | |
125 | |
126 Args: | |
127 manager: Multiprocessing.manager object. | |
128 """ | |
129 super(ThreadAndProcessSafeDict, self).__init__() | |
130 self.dict = manager.dict() | |
OLD | NEW |