fix: synchronize in-memory store access
This commit is contained in:
2 files changed
+172
-95
No files matched your search
@@ -1,4 +1,7 @@
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
from agentmsgs.stores.memory import InMemoryStore
|
||||
|
||||
@@ -86,3 +89,57 @@ def test_memory_thread_info_for_empty_thread():
|
||||
assert state.unread_from_others_count == 0
|
||||
assert s.get_thread_info(uuid.uuid4()) is None
|
||||
assert s.get_thread_read_state(uuid.uuid4(), a) is None
|
||||
|
||||
|
||||
def test_memory_concurrent_get_or_create_agent_is_idempotent():
|
||||
store = InMemoryStore()
|
||||
start = threading.Barrier(2)
|
||||
|
||||
class SlowNameIndex(dict):
|
||||
def __contains__(self, key):
|
||||
present = super().__contains__(key)
|
||||
if not present:
|
||||
time.sleep(0.01)
|
||||
return present
|
||||
|
||||
store._name_idx = SlowNameIndex()
|
||||
|
||||
def create_agent():
|
||||
start.wait()
|
||||
return store.get_or_create_agent("Alice")
|
||||
|
||||
with ThreadPoolExecutor(max_workers=2) as executor:
|
||||
agents = list(executor.map(lambda _: create_agent(), range(2)))
|
||||
|
||||
assert len({agent.id for agent in agents}) == 1
|
||||
assert len(store.list_agents()) == 1
|
||||
|
||||
|
||||
def test_memory_concurrent_append_assigns_unique_sequences():
|
||||
store = InMemoryStore()
|
||||
alice = store.get_or_create_agent("Alice")
|
||||
bob = store.get_or_create_agent("Bob")
|
||||
thread = store.create_thread({alice, bob})
|
||||
gate = threading.Barrier(2, timeout=0.2)
|
||||
local = threading.local()
|
||||
|
||||
class SlowMessageList(list):
|
||||
def __len__(self):
|
||||
length = super().__len__()
|
||||
if not hasattr(local, "waited"):
|
||||
try:
|
||||
gate.wait()
|
||||
except threading.BrokenBarrierError:
|
||||
pass
|
||||
local.waited = True
|
||||
return length
|
||||
|
||||
store._msgs[thread.id] = SlowMessageList()
|
||||
|
||||
def append_message(index):
|
||||
return store.append_message(thread.id, alice, str(index))
|
||||
|
||||
with ThreadPoolExecutor(max_workers=2) as executor:
|
||||
messages = list(executor.map(append_message, range(2)))
|
||||
|
||||
assert sorted(message.seq for message in messages) == [1, 2]
|
||||
Reference in new issue
Block a user