# Copyright (c) 2015-2019 The Botogram Authors (see AUTHORS) # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER # DEALINGS IN THE SOFTWARE. import collections import multiprocessing import multiprocessing.managers class OverrideableDict(dict): pass class SharedMemoryCommands: """Definition of IPC commands for the shared memory""" def __init__(self): self._memories = {} self._manager = multiprocessing.managers.SyncManager() self._locks = set() self._locks_queues = {} def start(self): """Start the shared memory manager""" self._manager.start() def get(self, memory_id, reply): """Get the shared memory which has the provided ID""" new = False if memory_id not in self._memories: self._memories[memory_id] = self._manager.dict() new = True # Send the shared memory to the process which requested it reply((self._memories[memory_id], new)) def list(self, memory_id, reply): """Get all the shared memories available""" reply(list(self._memories.keys())) def lock_acquire(self, lock_id, reply): """Acquire a lock""" # If the lock isn't acquired acquire it if lock_id not in self._locks: self._locks.add(lock_id) return reply(None) # Else ignore the request, and add the reply function to the queue if lock_id not in self._locks_queues: self._locks_queues[lock_id] = collections.deque() self._locks_queues[lock_id].appendleft(reply) def lock_release(self, lock_id, reply): """Release a lock""" # If the lock wasn't acquired, just return if lock_id not in self._locks: return reply(None) self._locks.remove(lock_id) # If there are processes waiting for this lock, wake up one of them if lock_id in self._locks_queues: self._locks_queues[lock_id].pop()(None) # And clear up the queue if it's empty if not len(self._locks_queues[lock_id]): del self._locks_queues[lock_id] reply(None) def lock_status(self, lock_id, reply): """Check if a lock was acquired""" reply(lock_id in self._locks) def lock_import(self, locks, reply): """Bulk import all the locks""" self._locks = set(locks) self._locks_queues = {} reply(None) def lock_export(self, __, reply): """Export all the locks""" reply(self._locks) class MultiprocessingDriver: """This is a multiprocessing-ready driver for the shared memory""" def __init__(self): self._memories = {} def __reduce__(self): return rebuild_driver, tuple() def _command(self, command, arg): """Send a command""" ipc = multiprocessing.current_process().ipc return ipc.command(command, arg) def get(self, memory_id): # Create the shared memory if it doens't exist is_new = False if memory_id not in self._memories: memory, is_new = self._command("shared.get", memory_id) self._memories[memory_id] = memory return self._memories[memory_id], is_new def lock_acquire(self, lock_id): # This automagically blocks if the lock is already acquired self._command("shared.lock_acquire", lock_id) def lock_release(self, lock_id): self._command("shared.lock_release", lock_id) def lock_status(self, lock_id): return self._command("shared.lock_status", lock_id) def import_data(self, data): # This will merge the provided component with the shared memory for memory_id, memory in data["storage"].items(): memory = self.get(memory_id) memory.update(data) if len(data["locks"]): self._command("shared.lock_import", data["locks"]) def export_data(self): result = {"storage": {}} for memory_id, data in self._memories.items(): result["storage"][memory_id] = dict(data) result["locks"] = self._command("shared.lock_export", None) return result def rebuild_driver(): return MultiprocessingDriver()