447 lines
13 KiB
Python
447 lines
13 KiB
Python
"""
|
||
多进程数据共享测试 - 验证跨进程数据同步能力
|
||
"""
|
||
|
||
import multiprocessing
|
||
import time
|
||
import pytest
|
||
from mpsp.mpsp import MultiProcessingSharedPool
|
||
|
||
|
||
# ==================== 辅助函数（需要在模块级别定义以便多进程使用）====================
|
||
|
||
|
||
def worker_put_data(key, value):
|
||
"""子进程：往共享池写入数据"""
|
||
pool = MultiProcessingSharedPool()
|
||
pool.put(key, value)
|
||
return True
|
||
|
||
|
||
def worker_get_data(key, result_queue):
|
||
"""子进程：从共享池读取数据并放入结果队列"""
|
||
pool = MultiProcessingSharedPool()
|
||
value = pool.get(key)
|
||
result_queue.put(value)
|
||
|
||
|
||
def worker_check_exists(key, result_queue):
|
||
"""子进程：检查 key 是否存在"""
|
||
pool = MultiProcessingSharedPool()
|
||
result_queue.put(pool.exists(key))
|
||
|
||
|
||
def worker_modify_data(key, new_value, result_queue):
|
||
"""子进程：修改数据并返回旧值"""
|
||
pool = MultiProcessingSharedPool()
|
||
old_value = pool.get(key)
|
||
pool.put(key, new_value)
|
||
result_queue.put(old_value)
|
||
|
||
|
||
def worker_wait_and_get(key, wait_time, result_queue):
|
||
"""子进程：等待一段时间后读取数据"""
|
||
time.sleep(wait_time)
|
||
pool = MultiProcessingSharedPool()
|
||
result_queue.put(pool.get(key))
|
||
|
||
|
||
def worker_increment_counter(key, iterations):
|
||
"""子进程：对计数器进行递增"""
|
||
pool = MultiProcessingSharedPool()
|
||
for _ in range(iterations):
|
||
# 注意：这不是原子操作，仅用于测试并发访问
|
||
current = pool.get(key, 0)
|
||
pool.put(key, current + 1)
|
||
|
||
|
||
def worker_pop_data(key, result_queue):
|
||
"""子进程：弹出数据"""
|
||
pool = MultiProcessingSharedPool()
|
||
value = pool.pop(key)
|
||
result_queue.put(value)
|
||
|
||
|
||
# ==================== 测试类 ====================
|
||
|
||
|
||
class TestParentChildProcess:
|
||
"""测试父子进程间数据传递"""
|
||
|
||
def test_parent_write_child_read(self):
|
||
"""测试父进程写入，子进程读取"""
|
||
pool = MultiProcessingSharedPool()
|
||
pool.clear()
|
||
|
||
# 父进程写入数据
|
||
pool.put("shared_key", "shared_value")
|
||
|
||
# 子进程读取数据
|
||
result_queue = multiprocessing.Queue()
|
||
p = multiprocessing.Process(
|
||
target=worker_get_data, args=("shared_key", result_queue)
|
||
)
|
||
p.start()
|
||
p.join()
|
||
|
||
result = result_queue.get()
|
||
assert result == "shared_value"
|
||
|
||
def test_child_write_parent_read(self):
|
||
"""测试子进程写入，父进程读取"""
|
||
pool = MultiProcessingSharedPool()
|
||
pool.clear()
|
||
|
||
# 子进程写入数据
|
||
p = multiprocessing.Process(
|
||
target=worker_put_data, args=("child_key", "child_value")
|
||
)
|
||
p.start()
|
||
p.join()
|
||
|
||
# 父进程读取数据（需要短暂等待以确保数据同步）
|
||
time.sleep(0.1)
|
||
result = pool.get("child_key")
|
||
assert result == "child_value"
|
||
|
||
def test_parent_child_data_isolation(self):
|
||
"""测试父子进程数据隔离 - 验证 manager.dict 的同步机制"""
|
||
pool = MultiProcessingSharedPool()
|
||
pool.clear()
|
||
|
||
# 父进程写入初始数据
|
||
pool.put("isolation_test", "parent_value")
|
||
|
||
# 子进程读取并修改
|
||
result_queue = multiprocessing.Queue()
|
||
p = multiprocessing.Process(
|
||
target=worker_modify_data,
|
||
args=("isolation_test", "child_value", result_queue),
|
||
)
|
||
p.start()
|
||
p.join()
|
||
|
||
# 验证子进程读取到了父进程的值
|
||
child_read = result_queue.get()
|
||
assert child_read == "parent_value"
|
||
|
||
# 验证父进程可以看到子进程的修改
|
||
time.sleep(0.1)
|
||
parent_read = pool.get("isolation_test")
|
||
assert parent_read == "child_value"
|
||
|
||
|
||
class TestMultipleChildProcesses:
|
||
"""测试多个子进程间的数据共享"""
|
||
|
||
def test_multiple_children_write_same_key(self):
|
||
"""测试多个子进程写入同一个 key（后面的覆盖前面的）"""
|
||
pool = MultiProcessingSharedPool()
|
||
pool.clear()
|
||
|
||
processes = []
|
||
for i in range(5):
|
||
p = multiprocessing.Process(
|
||
target=worker_put_data, args=("shared_key", f"value_{i}")
|
||
)
|
||
processes.append(p)
|
||
p.start()
|
||
|
||
for p in processes:
|
||
p.join()
|
||
|
||
# 短暂等待确保所有写入完成
|
||
time.sleep(0.1)
|
||
|
||
# 验证有一个值被成功写入
|
||
result = pool.get("shared_key")
|
||
assert result.startswith("value_")
|
||
|
||
def test_multiple_children_write_different_keys(self):
|
||
"""测试多个子进程写入不同的 key"""
|
||
pool = MultiProcessingSharedPool()
|
||
pool.clear()
|
||
|
||
num_processes = 5
|
||
processes = []
|
||
|
||
for i in range(num_processes):
|
||
p = multiprocessing.Process(
|
||
target=worker_put_data, args=(f"key_{i}", f"value_{i}")
|
||
)
|
||
processes.append(p)
|
||
p.start()
|
||
|
||
for p in processes:
|
||
p.join()
|
||
|
||
# 短暂等待确保所有写入完成
|
||
time.sleep(0.1)
|
||
|
||
# 验证所有值都被成功写入
|
||
for i in range(num_processes):
|
||
assert pool.get(f"key_{i}") == f"value_{i}"
|
||
|
||
def test_multiple_children_read_same_key(self):
|
||
"""测试多个子进程读取同一个 key"""
|
||
pool = MultiProcessingSharedPool()
|
||
pool.clear()
|
||
|
||
# 父进程写入数据
|
||
pool.put("shared_key", "shared_value")
|
||
|
||
# 多个子进程读取
|
||
num_processes = 5
|
||
result_queue = multiprocessing.Queue()
|
||
processes = []
|
||
|
||
for _ in range(num_processes):
|
||
p = multiprocessing.Process(
|
||
target=worker_get_data, args=("shared_key", result_queue)
|
||
)
|
||
processes.append(p)
|
||
p.start()
|
||
|
||
for p in processes:
|
||
p.join()
|
||
|
||
# 收集所有结果
|
||
results = []
|
||
for _ in range(num_processes):
|
||
results.append(result_queue.get())
|
||
|
||
# 所有子进程都应该读取到相同的值
|
||
assert all(r == "shared_value" for r in results)
|
||
|
||
def test_concurrent_exists_check(self):
|
||
"""测试并发检查 key 是否存在"""
|
||
pool = MultiProcessingSharedPool()
|
||
pool.clear()
|
||
|
||
# 写入一些数据
|
||
for i in range(5):
|
||
pool.put(f"key_{i}", f"value_{i}")
|
||
|
||
# 多个子进程并发检查
|
||
result_queue = multiprocessing.Queue()
|
||
processes = []
|
||
|
||
for i in range(10):
|
||
key = f"key_{i % 7}" # 有些 key 存在，有些不存在
|
||
p = multiprocessing.Process(
|
||
target=worker_check_exists, args=(key, result_queue)
|
||
)
|
||
processes.append(p)
|
||
p.start()
|
||
|
||
for p in processes:
|
||
p.join()
|
||
|
||
# 收集结果
|
||
results = []
|
||
for _ in range(10):
|
||
results.append(result_queue.get())
|
||
|
||
# 前 5 个应该存在 (key_0 到 key_4)，后 5 个不存在 (key_5, key_6 重复检查)
|
||
assert sum(results) >= 5 # 至少 5 个存在
|
||
|
||
|
||
# 进程池测试需要在模块级别定义 worker 函数
|
||
def _pool_worker_map(args):
|
||
"""进程池 map 操作的 worker"""
|
||
idx, key = args
|
||
shared_pool = MultiProcessingSharedPool()
|
||
value = shared_pool.get(key)
|
||
return idx, value
|
||
|
||
|
||
def _pool_worker_apply(key):
|
||
"""进程池 apply_async 的 worker"""
|
||
shared_pool = MultiProcessingSharedPool()
|
||
return shared_pool.get(key)
|
||
|
||
|
||
class TestProcessPool:
|
||
"""测试在进程池中使用"""
|
||
|
||
def test_pool_map_with_shared_data(self):
|
||
"""测试在进程池 map 操作中使用共享数据"""
|
||
pool = MultiProcessingSharedPool()
|
||
pool.clear()
|
||
|
||
# 写入测试数据
|
||
for i in range(5):
|
||
pool.put(f"input_{i}", i * 10)
|
||
|
||
# 使用进程池
|
||
with multiprocessing.Pool(processes=3) as process_pool:
|
||
results = process_pool.map(
|
||
_pool_worker_map, [(i, f"input_{i}") for i in range(5)]
|
||
)
|
||
|
||
# 验证结果
|
||
results_dict = {idx: val for idx, val in results}
|
||
for i in range(5):
|
||
assert results_dict[i] == i * 10
|
||
|
||
def test_pool_apply_async_with_shared_data(self):
|
||
"""测试在进程池 apply_async 中使用共享数据"""
|
||
pool = MultiProcessingSharedPool()
|
||
pool.clear()
|
||
|
||
pool.put("async_key", "async_value")
|
||
|
||
with multiprocessing.Pool(processes=2) as process_pool:
|
||
result = process_pool.apply_async(_pool_worker_apply, ("async_key",))
|
||
assert result.get(timeout=5) == "async_value"
|
||
|
||
|
||
class TestDataVisibility:
|
||
"""测试数据可见性和同步时机"""
|
||
|
||
def test_immediate_visibility_after_put(self):
|
||
"""测试写入后立即可见"""
|
||
pool = MultiProcessingSharedPool()
|
||
pool.clear()
|
||
|
||
pool.put("immediate_key", "immediate_value")
|
||
# 同一进程内应该立即可见
|
||
assert pool.exists("immediate_key")
|
||
assert pool.get("immediate_key") == "immediate_value"
|
||
|
||
def test_cross_process_visibility_with_delay(self):
|
||
"""测试跨进程可见性（带延迟）"""
|
||
pool = MultiProcessingSharedPool()
|
||
pool.clear()
|
||
|
||
# 父进程写入
|
||
pool.put("delayed_key", "delayed_value")
|
||
|
||
# 子进程延迟后读取
|
||
result_queue = multiprocessing.Queue()
|
||
p = multiprocessing.Process(
|
||
target=worker_wait_and_get, args=("delayed_key", 0.2, result_queue)
|
||
)
|
||
p.start()
|
||
p.join()
|
||
|
||
result = result_queue.get()
|
||
assert result == "delayed_value"
|
||
|
||
|
||
class TestConcurrentModifications:
|
||
"""测试并发修改场景"""
|
||
|
||
def test_concurrent_counter_increments(self):
|
||
"""测试并发计数器递增（非原子操作，预期会有竞争条件）"""
|
||
pool = MultiProcessingSharedPool()
|
||
pool.clear()
|
||
|
||
# 初始化计数器
|
||
pool.put("counter", 0)
|
||
|
||
num_processes = 4
|
||
iterations_per_process = 10
|
||
|
||
processes = []
|
||
for _ in range(num_processes):
|
||
p = multiprocessing.Process(
|
||
target=worker_increment_counter,
|
||
args=("counter", iterations_per_process),
|
||
)
|
||
processes.append(p)
|
||
p.start()
|
||
|
||
for p in processes:
|
||
p.join()
|
||
|
||
time.sleep(0.1)
|
||
|
||
# 由于竞争条件，实际值可能小于期望值
|
||
# 这个测试主要是为了验证并发访问不会崩溃
|
||
final_count = pool.get("counter")
|
||
assert isinstance(final_count, int)
|
||
assert 0 <= final_count <= num_processes * iterations_per_process
|
||
|
||
def test_concurrent_pop_operations(self):
|
||
"""测试并发 pop 操作"""
|
||
pool = MultiProcessingSharedPool()
|
||
pool.clear()
|
||
|
||
# 初始化多个 key
|
||
num_keys = 5
|
||
for i in range(num_keys):
|
||
pool.put(f"pop_key_{i}", f"pop_value_{i}")
|
||
|
||
result_queue = multiprocessing.Queue()
|
||
processes = []
|
||
|
||
for i in range(num_keys):
|
||
p = multiprocessing.Process(
|
||
target=worker_pop_data, args=(f"pop_key_{i}", result_queue)
|
||
)
|
||
processes.append(p)
|
||
p.start()
|
||
|
||
for p in processes:
|
||
p.join()
|
||
|
||
# 收集所有 pop 的结果
|
||
popped_values = []
|
||
for _ in range(num_keys):
|
||
popped_values.append(result_queue.get())
|
||
|
||
# 验证所有值都被正确 pop
|
||
assert len(popped_values) == num_keys
|
||
for i in range(num_keys):
|
||
assert f"pop_value_{i}" in popped_values
|
||
|
||
# 验证所有 key 都被移除
|
||
time.sleep(0.1)
|
||
for i in range(num_keys):
|
||
assert not pool.exists(f"pop_key_{i}")
|
||
|
||
|
||
class TestComplexDataTypes:
|
||
"""测试复杂数据类型的多进程共享"""
|
||
|
||
def test_share_nested_dict(self):
|
||
"""测试共享嵌套字典"""
|
||
pool = MultiProcessingSharedPool()
|
||
pool.clear()
|
||
|
||
nested_data = {
|
||
"level1": {"level2": {"level3": [1, 2, 3]}},
|
||
"list_of_dicts": [{"a": 1}, {"b": 2}],
|
||
}
|
||
|
||
pool.put("nested_key", nested_data)
|
||
|
||
result_queue = multiprocessing.Queue()
|
||
p = multiprocessing.Process(
|
||
target=worker_get_data, args=("nested_key", result_queue)
|
||
)
|
||
p.start()
|
||
p.join()
|
||
|
||
result = result_queue.get()
|
||
assert result == nested_data
|
||
|
||
def test_share_large_list(self):
|
||
"""测试共享大型列表"""
|
||
pool = MultiProcessingSharedPool()
|
||
pool.clear()
|
||
|
||
large_list = list(range(10000))
|
||
pool.put("large_list_key", large_list)
|
||
|
||
result_queue = multiprocessing.Queue()
|
||
p = multiprocessing.Process(
|
||
target=worker_get_data, args=("large_list_key", result_queue)
|
||
)
|
||
p.start()
|
||
p.join()
|
||
|
||
result = result_queue.get()
|
||
assert result == large_list
|