Dictionary from multiprocessing handles to StorageWeakRef.
| 62 | |
| 63 | |
| 64 | class SharedCache(dict): |
| 65 | """Dictionary from multiprocessing handles to StorageWeakRef.""" |
| 66 | |
| 67 | def __init__(self) -> None: |
| 68 | # free_dead_references() is called if the len exceeds the current |
| 69 | # limit. The limit scales with the number of remaining live objects. |
| 70 | self.limit = 128 |
| 71 | # `fork` inherits lock state, so in case we fork when the lock is held, |
| 72 | # we register a function to reset the lock to a new object to avoid |
| 73 | # possible deadlocks, following python multiprocessing library design. |
| 74 | self._after_fork() |
| 75 | register_after_fork(self, SharedCache._after_fork) |
| 76 | |
| 77 | def _after_fork(self): |
| 78 | self.lock = threading.Lock() |
| 79 | |
| 80 | def get(self, key): # type: ignore[override] |
| 81 | with self.lock: |
| 82 | return dict.get(self, key) |
| 83 | |
| 84 | def __setitem__(self, key, storage_ref): |
| 85 | with self.lock: |
| 86 | dict.__setitem__(self, key, storage_ref) |
| 87 | if len(self) > self.limit: |
| 88 | self.free_dead_references() |
| 89 | |
| 90 | def free_dead_references(self): |
| 91 | live = 0 |
| 92 | for key, storage_ref in list(self.items()): |
| 93 | if storage_ref.expired(): |
| 94 | del self[key] |
| 95 | else: |
| 96 | live += 1 |
| 97 | self.limit = max(128, live * 2) |
| 98 | |
| 99 | |
| 100 | # mapping from handles to StorageWeakRef objects |