18d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi# -*- coding: utf-8 -*- 28d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi# Copyright 2013 Google Inc. All Rights Reserved. 38d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi# 48d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi# Licensed under the Apache License, Version 2.0 (the "License"); 58d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi# you may not use this file except in compliance with the License. 68d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi# You may obtain a copy of the License at 78d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi# 88d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi# http://www.apache.org/licenses/LICENSE-2.0 98d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi# 108d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi# Unless required by applicable law or agreed to in writing, software 118d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi# distributed under the License is distributed on an "AS IS" BASIS, 128d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 138d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi# See the License for the specific language governing permissions and 148d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi# limitations under the License. 158d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi"""Utility classes for the parallelism framework.""" 168d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi 178d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoifrom __future__ import absolute_import 188d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi 198d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoiimport threading 208d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi 218d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi 22cef7893435aa41160dd1255c43cb8498279738ccChris Craikclass AtomicDict(object): 23cef7893435aa41160dd1255c43cb8498279738ccChris Craik """Thread-safe (and optionally process-safe) dictionary protected by a lock. 248d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi 25cef7893435aa41160dd1255c43cb8498279738ccChris Craik If a multiprocessing.Manager is supplied on init, the dictionary is 26cef7893435aa41160dd1255c43cb8498279738ccChris Craik both process and thread safe. Otherwise, it is only thread-safe. 278d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi """ 288d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi 29cef7893435aa41160dd1255c43cb8498279738ccChris Craik def __init__(self, manager=None): 30cef7893435aa41160dd1255c43cb8498279738ccChris Craik """Initializes the dict. 318d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi 328d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi Args: 33cef7893435aa41160dd1255c43cb8498279738ccChris Craik manager: multiprocessing.Manager instance (required for process safety). 348d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi """ 35cef7893435aa41160dd1255c43cb8498279738ccChris Craik if manager: 36cef7893435aa41160dd1255c43cb8498279738ccChris Craik self.lock = manager.Lock() 37cef7893435aa41160dd1255c43cb8498279738ccChris Craik self.dict = manager.dict() 38cef7893435aa41160dd1255c43cb8498279738ccChris Craik else: 39cef7893435aa41160dd1255c43cb8498279738ccChris Craik self.lock = threading.Lock() 40cef7893435aa41160dd1255c43cb8498279738ccChris Craik self.dict = {} 418d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi 428d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi def __getitem__(self, key): 438d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi with self.lock: 448d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi return self.dict[key] 458d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi 468d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi def __setitem__(self, key, value): 478d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi with self.lock: 488d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi self.dict[key] = value 498d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi 508d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi # pylint: disable=invalid-name 518d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi def get(self, key, default_value=None): 528d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi with self.lock: 538d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi return self.dict.get(key, default_value) 548d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi 558d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi def delete(self, key): 568d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi with self.lock: 578d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi del self.dict[key] 588d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi 59cef7893435aa41160dd1255c43cb8498279738ccChris Craik def Increment(self, key, inc, default_value=0): 60cef7893435aa41160dd1255c43cb8498279738ccChris Craik """Atomically updates the stored value associated with the given key. 618d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi 62cef7893435aa41160dd1255c43cb8498279738ccChris Craik Performs the atomic equivalent of 63cef7893435aa41160dd1255c43cb8498279738ccChris Craik dict[key] = dict.get(key, default_value) + inc. 648d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi 658d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi Args: 66cef7893435aa41160dd1255c43cb8498279738ccChris Craik key: lookup key for the value of the first operand of the "+" operation. 67cef7893435aa41160dd1255c43cb8498279738ccChris Craik inc: Second operand of the "+" operation. 68cef7893435aa41160dd1255c43cb8498279738ccChris Craik default_value: Default value if there is no existing value for the key. 69cef7893435aa41160dd1255c43cb8498279738ccChris Craik 70cef7893435aa41160dd1255c43cb8498279738ccChris Craik Returns: 71cef7893435aa41160dd1255c43cb8498279738ccChris Craik Incremented value. 728d2b206a675ec20ea07100c35df34e65ee1e45e8Ruchi Kandoi """ 73cef7893435aa41160dd1255c43cb8498279738ccChris Craik with self.lock: 74cef7893435aa41160dd1255c43cb8498279738ccChris Craik val = self.dict.get(key, default_value) + inc 75cef7893435aa41160dd1255c43cb8498279738ccChris Craik self.dict[key] = val 76cef7893435aa41160dd1255c43cb8498279738ccChris Craik return val 77