fix(framework): Make TypedDict mutations thread-safe - #7980
Conversation
There was a problem hiding this comment.
Pull request overview
This PR introduces a per-instance threading.Lock to TypedDict to synchronize concurrent mutations (assignment/deletion) and ensures copied/unpickled instances get a fresh lock so serialization/copying continues to work.
Changes:
- Add
_locktoTypedDictinstances and guard__setitem__/__delitem__with it. - Ensure
copy()assigns a new lock on the copied instance. - Add
__getstate__/__setstate__to exclude the unpicklable lock from pickling and recreate it on restore.
Critical issues
copy()snapshots_datawithout holding_lock, which can still raiseRuntimeError: dictionary changed size during iterationor produce an inconsistent snapshot under concurrent writers.__getstate__exposes_datato pickling without synchronizing with_lock, so pickling can still fail withRuntimeErrorunder concurrent writers.- No tests were added for the new concurrency behavior; a small regression test would help prevent future breakage.
Simplicity/readability suggestions
- Prefer
state.pop("_lock", None)overdel state["_lock"]in__getstate__to make the method more robust to unexpected state shapes.
Consistency concerns
- None identified beyond the concurrency gaps above.
Whether the PR should be split
- No; the changes are cohesive and localized.
Brief overall verdict
The core approach is sound, but copy()/pickling still have concurrency failure modes and the change should be backed by at least a minimal concurrency test.
Suppressed comments (1)
framework/py/flwr/app/message/typeddict.py:115
__getstate__returns_databy reference and does not synchronize with_lock. Pickle will iterate_data; if another thread is mutating concurrently, this can raiseRuntimeError: dictionary changed size during iteration. Snapshot_datawhile holding_lockand remove_lockviapopfor robustness.
def __getstate__(self) -> dict[str, object]:
"""Return the state without the unpicklable lock."""
state = self.__dict__.copy()
del state["_lock"]
return state
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| new.__dict__["_check_key_fn"] = self.__dict__["_check_key_fn"] | ||
| new.__dict__["_check_value_fn"] = self.__dict__["_check_value_fn"] | ||
| new.__dict__["_data"] = cast(dict[K, V], self.__dict__["_data"]).copy() | ||
| new.__dict__["_lock"] = Lock() | ||
| return new |
| @@ -46,11 +47,13 @@ def __setitem__(self, key: K, value: V) -> None: | |||
| cast(Callable[[V], None], self.__dict__["_check_value_fn"])(value) | |||
|
|
|||
| # Set key-value pair | |||
| cast(dict[K, V], self.__dict__["_data"])[key] = value | |||
| with self.__dict__["_lock"]: | |||
| cast(dict[K, V], self.__dict__["_data"])[key] = value | |||
|
|
|||
| def __delitem__(self, key: K) -> None: | |||
| """Remove the item with the specified key.""" | |||
| del cast(dict[K, V], self.__dict__["_data"])[key] | |||
| with self.__dict__["_lock"]: | |||
| del cast(dict[K, V], self.__dict__["_data"])[key] | |||
github-actions
Bot
added
the
Maintainer
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: de43ab41f4
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| with self.__dict__["_lock"]: | ||
| cast(dict[K, V], self.__dict__["_data"])[key] = value |
There was a problem hiding this comment.
When two threads call the inherited MutableMapping.setdefault for the same missing key, both perform the unlocked lookup before reaching this lock, then each stores and returns its own default even though only one value remains in the mapping. The inherited pop and popitem have similar lookup-then-delete races. Override these compound mutators so the entire operation is protected by the instance lock rather than locking only the final assignment or deletion.
Useful? React with 👍 / 👎.
Issue
Description
Concurrent writes to TypedDict are not synchronized.
Related issues/PRs
None.
Proposal
Explanation
Add a per-instance lock around item assignment and deletion while leaving read methods unchanged. Recreate the lock when copying or unpickling a TypedDict so existing copy and serialization behavior continues to work.
Checklist
Any other comments?
No new tests were added because this is a small synchronization change. Existing validation passed: