224 lines
8.0 KiB
Python
224 lines
8.0 KiB
Python
"""并发上限单测（caps）。
|
||
|
||
覆盖：配置解析优先级与默认值、认领门的未满 / 已满 / 释放后
|
||
恢复、启动中任务计入占用、CAS 抢输后占位不泄漏、缺 key 门
|
||
提前返回后占位释放、双线程同时认领只赢一个。
|
||
|
||
只用注入的假环境与本地 SQLite，不跑真实模型、不联网。
|
||
"""
|
||
|
||
import os
|
||
import sqlite3
|
||
import sys
|
||
import tempfile
|
||
import threading
|
||
import time
|
||
import unittest
|
||
from pathlib import Path
|
||
|
||
try:
|
||
from cloud.caps import (
|
||
can_claim,
|
||
effective_capacity,
|
||
resolve_capacity,
|
||
running_count,
|
||
)
|
||
from cloud.server import init_db
|
||
from cloud.worker import Worker
|
||
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||
from cloud.caps import (
|
||
can_claim,
|
||
effective_capacity,
|
||
resolve_capacity,
|
||
running_count,
|
||
)
|
||
from cloud.server import init_db
|
||
from cloud.worker import Worker
|
||
|
||
|
||
def _insert_task(db_path, task_id, project_id="p1"):
|
||
now = time.time()
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
conn.execute(
|
||
"INSERT INTO tasks(id, project_id, prompt, status,"
|
||
" idempotency_key, depends_on, priority, created_by,"
|
||
" created_at, updated_at, last_heartbeat)"
|
||
" VALUES(?, ?, ?, 'queued', NULL, '[]', 0, 'tester', ?, ?, NULL)",
|
||
(task_id, project_id, "do work", now, now),
|
||
)
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def _status(db_path, task_id):
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
row = conn.execute(
|
||
"SELECT status FROM tasks WHERE id = ?",
|
||
(task_id,)).fetchone()
|
||
finally:
|
||
conn.close()
|
||
return row[0] if row else None
|
||
|
||
|
||
class CapsResolveTestCase(unittest.TestCase):
|
||
def test_explicit_beats_env(self):
|
||
env = {"CLOUD_WORKER_CAPACITY": "3"}
|
||
self.assertEqual(resolve_capacity(2, env=env), 2)
|
||
|
||
def test_env_used_when_no_explicit(self):
|
||
env = {"CLOUD_WORKER_CAPACITY": "3"}
|
||
self.assertEqual(resolve_capacity(None, env=env), 3)
|
||
|
||
def test_default_is_1(self):
|
||
self.assertEqual(resolve_capacity(None, env={}), 1)
|
||
self.assertEqual(resolve_capacity(), 1)
|
||
|
||
def test_invalid_values_fall_back_to_1(self):
|
||
self.assertEqual(resolve_capacity(0, env={}), 1)
|
||
self.assertEqual(resolve_capacity(-2, env={}), 1)
|
||
self.assertEqual(resolve_capacity("abc", env={}), 1)
|
||
self.assertEqual(resolve_capacity(None, env={
|
||
"CLOUD_WORKER_CAPACITY": "abc"}), 1)
|
||
self.assertEqual(resolve_capacity(None, env={
|
||
"CLOUD_WORKER_CAPACITY": "0"}), 1)
|
||
self.assertEqual(resolve_capacity(None, env={
|
||
"CLOUD_WORKER_CAPACITY": " "}), 1)
|
||
|
||
def test_effective_capacity(self):
|
||
self.assertEqual(effective_capacity(2), 2)
|
||
self.assertEqual(effective_capacity(0), 1)
|
||
self.assertEqual(effective_capacity(None), 1)
|
||
|
||
def test_can_claim_boundaries(self):
|
||
self.assertTrue(can_claim(0, 1))
|
||
self.assertFalse(can_claim(1, 1))
|
||
self.assertTrue(can_claim(1, 2))
|
||
self.assertFalse(can_claim(2, 2))
|
||
|
||
def test_running_count_counts_starting_entries(self):
|
||
# 空条目 = 已认领未起进程的启动中任务，同样占一个槽。
|
||
self.assertEqual(running_count({"a": {}, "b": {"proc": 1}}), 2)
|
||
self.assertEqual(running_count({}), 0)
|
||
|
||
|
||
class WorkerClaimGateTestCase(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||
init_db(self.db_path)
|
||
self.env = {"XCM_API_KEY": "test-placeholder"}
|
||
|
||
def tearDown(self):
|
||
self.tmp.cleanup()
|
||
|
||
def _worker(self, **kwargs):
|
||
kwargs.setdefault("env", self.env)
|
||
kwargs.setdefault("retry_backoff_seconds", 0)
|
||
kwargs.setdefault("max_retries", 0)
|
||
return Worker(self.db_path, "/fake/host.py", **kwargs)
|
||
|
||
def test_default_capacity_is_1(self):
|
||
worker = self._worker()
|
||
self.assertEqual(worker.capacity, 1)
|
||
|
||
def test_env_capacity_and_explicit_priority(self):
|
||
env = {"CLOUD_WORKER_CAPACITY": "3"}
|
||
self.assertEqual(self._worker(env=env).capacity, 3)
|
||
self.assertEqual(self._worker(env=env, capacity=2).capacity, 2)
|
||
|
||
def test_full_worker_refuses_claim_and_task_stays_queued(self):
|
||
_insert_task(self.db_path, "A")
|
||
_insert_task(self.db_path, "B")
|
||
worker = self._worker(capacity=1)
|
||
self.assertTrue(worker._claim("A"))
|
||
self.assertEqual(_status(self.db_path, "A"), "running")
|
||
self.assertFalse(worker._claim("B"))
|
||
self.assertEqual(_status(self.db_path, "B"), "queued")
|
||
self.assertEqual(set(worker._running_snapshot()), {"A"})
|
||
|
||
def test_release_restores_claim(self):
|
||
_insert_task(self.db_path, "A")
|
||
_insert_task(self.db_path, "B")
|
||
worker = self._worker(capacity=1)
|
||
self.assertTrue(worker._claim("A"))
|
||
self.assertFalse(worker._claim("B"))
|
||
worker._running_clear("A")
|
||
self.assertTrue(worker._claim("B"))
|
||
self.assertEqual(_status(self.db_path, "B"), "running")
|
||
|
||
def test_starting_task_counts_toward_capacity(self):
|
||
_insert_task(self.db_path, "B")
|
||
worker = self._worker(capacity=1)
|
||
# 空条目模拟“已认领、执行器还没起进程”的启动中任务。
|
||
worker._running_set("STARTING")
|
||
self.assertEqual(worker._running_count(), 1)
|
||
self.assertFalse(worker._claim("B"))
|
||
self.assertEqual(_status(self.db_path, "B"), "queued")
|
||
worker._running_clear("STARTING")
|
||
self.assertTrue(worker._claim("B"))
|
||
|
||
def test_capacity_2_allows_two_blocks_third(self):
|
||
for task_id in ("A", "B", "C"):
|
||
_insert_task(self.db_path, task_id)
|
||
worker = self._worker(capacity=2)
|
||
self.assertTrue(worker._claim("A"))
|
||
self.assertTrue(worker._claim("B"))
|
||
self.assertFalse(worker._claim("C"))
|
||
self.assertEqual(_status(self.db_path, "C"), "queued")
|
||
worker._running_clear("A")
|
||
self.assertTrue(worker._claim("C"))
|
||
|
||
def test_lost_cas_race_releases_slot(self):
|
||
_insert_task(self.db_path, "A")
|
||
_insert_task(self.db_path, "B")
|
||
first = self._worker(capacity=1)
|
||
second = self._worker(capacity=1)
|
||
self.assertTrue(first._claim("A"))
|
||
# second 的 CAS 抢输：返回 False，且它自己的占位已清，
|
||
# 不能留下一个空槽把自己堵死。
|
||
self.assertFalse(second._claim("A"))
|
||
self.assertEqual(second._running_snapshot(), {})
|
||
self.assertTrue(second._claim("B"))
|
||
|
||
def test_two_threads_claim_only_one_wins(self):
|
||
_insert_task(self.db_path, "A")
|
||
_insert_task(self.db_path, "B")
|
||
worker = self._worker(capacity=1)
|
||
barrier = threading.Barrier(2, timeout=5)
|
||
results = {}
|
||
|
||
def _try(task_id):
|
||
barrier.wait()
|
||
results[task_id] = worker._claim(task_id)
|
||
|
||
threads = [threading.Thread(target=_try, args=(t,))
|
||
for t in ("A", "B")]
|
||
for t in threads:
|
||
t.start()
|
||
for t in threads:
|
||
t.join(timeout=10)
|
||
self.assertEqual(sorted(results.values()), [False, True])
|
||
self.assertEqual(worker._running_count(), 1)
|
||
|
||
def test_missing_key_gate_releases_slot(self):
|
||
_insert_task(self.db_path, "A")
|
||
_insert_task(self.db_path, "B")
|
||
# env 里无 key：run_one 在缺 key 门提前返回 failed，
|
||
# 占位必须随之释放，否则 capacity=1 的 worker 会被堵死。
|
||
worker = self._worker(env={}, capacity=1)
|
||
result = worker.run_one()
|
||
self.assertTrue(result["dispatched"])
|
||
self.assertEqual(result["status"], "failed")
|
||
self.assertEqual(_status(self.db_path, "A"), "failed")
|
||
self.assertEqual(worker._running_snapshot(), {})
|
||
self.assertTrue(worker._claim("B"))
|
||
self.assertEqual(_status(self.db_path, "B"), "running")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|