Skip to content

fix(framework): Make TypedDict mutations thread-safe #7980

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
panh99 wants to merge 3 commits into from
Closed

fix(framework): Make TypedDict mutations thread-safe #7980

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
fccc24e
fix(framework): Make TypedDict mutations thread-safe
panh99 Aug 22, 2026
3c37e47
Merge branch 'main' into codex/typeddict-thread-safe
panh99 Aug 22, 2026
de43ab4
Merge branch 'main' into codex/typeddict-thread-safe
danieljanes Aug 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 18 additions & 3 deletions framework/py/flwr/app/message/typeddict.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@
# ==============================================================================
"""Typed dict base class for *Records."""


from collections.abc import (
Callable,
ItemsView,
Expand All @@ -23,6 +22,7 @@
MutableMapping,
ValuesView,
)
from threading import Lock
from typing import Generic, Self, TypeVar, cast

K = TypeVar("K") # Key type
Expand All @@ -38,6 +38,7 @@ def __init__(
self.__dict__["_check_key_fn"] = check_key_fn
self.__dict__["_check_value_fn"] = check_value_fn
self.__dict__["_data"] = {}
self.__dict__["_lock"] = Lock()

def __setitem__(self, key: K, value: V) -> None:
"""Set the given key to the given value after type checking."""
Expand All @@ -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
Comment on lines +50 to +51

@chatgpt-codex-connector chatgpt-codex-connector Bot Aug 22, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Protect compound mutations with the same lock

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 👍 / 👎.


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]
Comment on lines 43 to +56

def __getitem__(self, item: K) -> V:
"""Return the value for the specified key."""
Expand Down Expand Up @@ -102,4 +105,16 @@ def copy(self) -> Self:
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
Comment on lines 105 to 109

def __getstate__(self) -> dict[str, object]:
"""Return the state without the unpicklable lock."""
state = self.__dict__.copy()
del state["_lock"]
return state

def __setstate__(self, state: dict[str, object]) -> None:
"""Restore the state and create a new lock."""
self.__dict__.update(state)
self.__dict__["_lock"] = Lock()
Loading