mirror of
https://gitee.com/mateos/mateclaw.git
synced 2026-09-13 03:13:41 +08:00
feat(task): generic one-shot Callable runner with parent-conv cancel
This commit is contained in:
matevip
parent
82594878a0
commit
d90591fa35
@ -133,6 +133,130 @@ public class AsyncTaskService implements ApplicationRunner {
|
||||
return entity;
|
||||
}
|
||||
|
||||
// ==================== One-shot Callable submission ====================
|
||||
|
||||
/**
|
||||
* Submit a one-shot {@link Callable} that runs on this service's
|
||||
* {@code pollExecutor} and persists its outcome through the standard
|
||||
* {@code mate_async_task} lifecycle.
|
||||
* <p>
|
||||
* This is the local-work counterpart to {@link #startPolling}: instead of
|
||||
* polling an external provider, the worker runs {@code work.call()} in
|
||||
* the shared executor, writes the return value to {@code resultJson} (or
|
||||
* the exception message to {@code errorMessage}), and registers itself in
|
||||
* the same {@code activePolls} / {@code pollTaskToConv} bookkeeping so
|
||||
* {@link #onConversationDeleted} can cancel both kinds uniformly.
|
||||
* <p>
|
||||
* Race closure: the worker is scheduled at 0 ms but blocks on an internal
|
||||
* {@code CountDownLatch} until the calling thread has finished
|
||||
* registering both bookkeeping entries. Without this, the executor could
|
||||
* dequeue the worker, run its {@code finally} cleanup, and return — all
|
||||
* before {@code activePolls.put} runs on the calling thread — leaving a
|
||||
* ghost entry no later event drains.
|
||||
* <p>
|
||||
* Cancellation: while {@code work.call()} runs the worker has no
|
||||
* cooperative cancel signal beyond {@link #isConversationCanceled}; it
|
||||
* checks once before invoking the body and once after, so a parent
|
||||
* conversation deleted mid-run never lands as {@code succeeded}.
|
||||
* {@link #onConversationDeleted} additionally writes {@code failed}
|
||||
* synchronously for any non-terminal {@code agent_delegate} task so the
|
||||
* DB row never lingers in {@code running} after parent deletion.
|
||||
*
|
||||
* @param taskType Discriminator written to {@code task_type}
|
||||
* (e.g. {@code "agent_delegate"}). Listeners
|
||||
* and the conversation-deleted DB write-back gate
|
||||
* on this value.
|
||||
* @param conversationId Parent conversation ID. Written to
|
||||
* {@code conversation_id} so deleting the parent
|
||||
* conversation cascade-cancels the worker. Any
|
||||
* child / detached identifiers belong in
|
||||
* {@code requestJson}, not here.
|
||||
* @param messageId Optional parent message ID.
|
||||
* @param requestJson Caller-serialized request payload.
|
||||
* @param createdBy Audit attribution; counts toward
|
||||
* {@code MAX_ACTIVE_TASKS_PER_USER}.
|
||||
* @param work Body whose return value is persisted as
|
||||
* {@code resultJson}. A thrown exception lands as
|
||||
* {@code status=failed} with the exception
|
||||
* message recorded.
|
||||
* @return the created task entity (status = pending at return time).
|
||||
*/
|
||||
public AsyncTaskEntity submitOneShot(String taskType, String conversationId,
|
||||
Long messageId, String requestJson,
|
||||
String createdBy, Callable<String> work) {
|
||||
AsyncTaskEntity entity = createTask(taskType, conversationId, messageId,
|
||||
"internal", null, requestJson, createdBy);
|
||||
final String taskId = entity.getTaskId();
|
||||
updateStatus(taskId, "running", 0, null, null);
|
||||
|
||||
// schedule(0)-vs-put race closure: see method javadoc.
|
||||
CountDownLatch enrolled = new CountDownLatch(1);
|
||||
ScheduledFuture<?> future = pollExecutor.schedule(() -> {
|
||||
try {
|
||||
try {
|
||||
enrolled.await();
|
||||
} catch (InterruptedException ie) {
|
||||
Thread.currentThread().interrupt();
|
||||
updateStatus(taskId, "failed", null, null, "worker interrupted before start");
|
||||
return;
|
||||
}
|
||||
// Pre-call cancel check: parent conversation may have been
|
||||
// deleted between submitOneShot returning and the worker
|
||||
// being dequeued.
|
||||
if (isConversationCanceled(conversationId)) {
|
||||
updateStatus(taskId, "failed", null, null, "conversation deleted before start");
|
||||
return;
|
||||
}
|
||||
String result;
|
||||
try {
|
||||
result = work.call();
|
||||
} catch (Exception e) {
|
||||
log.warn("[AsyncTask] One-shot task {} failed: {}", taskId, e.getMessage());
|
||||
updateStatus(taskId, "failed", null, null, e.getMessage());
|
||||
return;
|
||||
}
|
||||
// Post-call cancel check: parent may have been deleted while
|
||||
// work was running; avoid resurrecting a succeeded row for a
|
||||
// conversation whose DB cascade already wiped it.
|
||||
if (isConversationCanceled(conversationId)) {
|
||||
updateStatus(taskId, "failed", null, null, "conversation deleted during execution");
|
||||
} else {
|
||||
updateStatus(taskId, "succeeded", 100, result, null);
|
||||
}
|
||||
} finally {
|
||||
activePolls.remove(taskId);
|
||||
pollTaskToConv.remove(taskId);
|
||||
// Generic terminal event so SSE listeners (parent
|
||||
// conversation of an async delegation, UI badges, …) can
|
||||
// react without polling the DB. Re-fetch so the event
|
||||
// payload reflects the row we just wrote. Wrapped in a
|
||||
// try-catch because a broadcast failure on a stale
|
||||
// conversation stream must not mask the task outcome.
|
||||
try {
|
||||
AsyncTaskEntity finalEntity = findEntityByTaskId(taskId);
|
||||
if (finalEntity != null) {
|
||||
boolean success = "succeeded".equals(finalEntity.getStatus());
|
||||
broadcastTaskEventWithData(finalEntity, "async_task_completed",
|
||||
success, java.util.Map.of(),
|
||||
success ? null : finalEntity.getErrorMessage());
|
||||
}
|
||||
} catch (Exception broadcastErr) {
|
||||
log.debug("[AsyncTask] Completion broadcast failed for task {}: {}",
|
||||
taskId, broadcastErr.getMessage());
|
||||
}
|
||||
}
|
||||
}, 0, TimeUnit.MILLISECONDS);
|
||||
|
||||
activePolls.put(taskId, future);
|
||||
if (conversationId != null) {
|
||||
pollTaskToConv.put(taskId, conversationId);
|
||||
}
|
||||
enrolled.countDown();
|
||||
log.info("[AsyncTask] Submitted one-shot task {} (type={}, conv={})",
|
||||
taskId, taskType, conversationId);
|
||||
return entity;
|
||||
}
|
||||
|
||||
// ==================== 轮询管理 ====================
|
||||
|
||||
/**
|
||||
@ -145,7 +269,7 @@ public class AsyncTaskService implements ApplicationRunner {
|
||||
public void startPolling(String taskId,
|
||||
Function<String, TaskPollResult> statusChecker,
|
||||
BiConsumer<AsyncTaskEntity, TaskPollResult> onComplete) {
|
||||
AsyncTaskEntity task = findByTaskId(taskId);
|
||||
AsyncTaskEntity task = findEntityByTaskId(taskId);
|
||||
if (task == null) {
|
||||
log.warn("[AsyncTask] Cannot start polling: task {} not found", taskId);
|
||||
return;
|
||||
@ -190,7 +314,7 @@ public class AsyncTaskService implements ApplicationRunner {
|
||||
updateStatus(taskId, "failed", null, null, result.errorMessage());
|
||||
}
|
||||
// 刷新任务实体
|
||||
AsyncTaskEntity freshTask = findByTaskId(taskId);
|
||||
AsyncTaskEntity freshTask = findEntityByTaskId(taskId);
|
||||
onComplete.accept(freshTask, result);
|
||||
}
|
||||
} catch (Exception e) {
|
||||
@ -251,7 +375,22 @@ public class AsyncTaskService implements ApplicationRunner {
|
||||
int cancelled = 0;
|
||||
for (Map.Entry<String, String> entry : pollTaskToConv.entrySet()) {
|
||||
if (convId.equals(entry.getValue())) {
|
||||
cancelPolling(entry.getKey());
|
||||
String taskId = entry.getKey();
|
||||
cancelPolling(taskId);
|
||||
// One-shot tasks (taskType "agent_delegate") have no separate
|
||||
// poll loop to observe the cancel and write the terminal row:
|
||||
// cancelPolling only nukes the Future. Without this explicit
|
||||
// write-back the DB row stays "running" forever. Polling
|
||||
// tasks (video / image / ...) keep their original behavior —
|
||||
// their own poll completion or startup-recovery path is what
|
||||
// writes the terminal status.
|
||||
AsyncTaskEntity t = findEntityByTaskId(taskId);
|
||||
if (t != null
|
||||
&& "agent_delegate".equals(t.getTaskType())
|
||||
&& !"succeeded".equals(t.getStatus())
|
||||
&& !"failed".equals(t.getStatus())) {
|
||||
updateStatus(taskId, "failed", null, null, "conversation deleted");
|
||||
}
|
||||
cancelled++;
|
||||
}
|
||||
}
|
||||
@ -290,7 +429,7 @@ public class AsyncTaskService implements ApplicationRunner {
|
||||
// ==================== 查询 ====================
|
||||
|
||||
public AsyncTaskInfo getTaskInfo(String taskId) {
|
||||
AsyncTaskEntity entity = findByTaskId(taskId);
|
||||
AsyncTaskEntity entity = findEntityByTaskId(taskId);
|
||||
if (entity == null) {
|
||||
return null;
|
||||
}
|
||||
@ -307,7 +446,12 @@ public class AsyncTaskService implements ApplicationRunner {
|
||||
return entities.stream().map(this::toInfo).toList();
|
||||
}
|
||||
|
||||
private AsyncTaskEntity findByTaskId(String taskId) {
|
||||
/** Returns the persisted task entity by its public {@code taskId}, or
|
||||
* {@code null} if no row matches. Promoted from private to public so the
|
||||
* conversation-deleted listener (and overrides in tests) can resolve a
|
||||
* task's current state without going through the read-model
|
||||
* {@link #getTaskInfo}. */
|
||||
public AsyncTaskEntity findEntityByTaskId(String taskId) {
|
||||
return asyncTaskMapper.selectOne(
|
||||
new LambdaQueryWrapper<AsyncTaskEntity>()
|
||||
.eq(AsyncTaskEntity::getTaskId, taskId)
|
||||
|
||||
@ -0,0 +1,347 @@
|
||||
package vip.mate.task;
|
||||
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.DisplayName;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.springframework.test.util.ReflectionTestUtils;
|
||||
import vip.mate.channel.web.ChatStreamTracker;
|
||||
import vip.mate.task.model.AsyncTaskEntity;
|
||||
import vip.mate.task.repository.AsyncTaskMapper;
|
||||
import vip.mate.workspace.conversation.event.ConversationDeletedEvent;
|
||||
|
||||
import java.util.UUID;
|
||||
import java.util.concurrent.Callable;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
import java.util.concurrent.ConcurrentMap;
|
||||
import java.util.concurrent.CountDownLatch;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
import java.util.concurrent.atomic.AtomicBoolean;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
import java.util.function.BooleanSupplier;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.mockito.ArgumentMatchers.any;
|
||||
import static org.mockito.ArgumentMatchers.eq;
|
||||
import static org.mockito.Mockito.mock;
|
||||
import static org.mockito.Mockito.timeout;
|
||||
import static org.mockito.Mockito.verify;
|
||||
|
||||
/**
|
||||
* Contract for {@link AsyncTaskService#submitOneShot} — the one-shot Callable
|
||||
* entry point that lets a caller hand off arbitrary work to {@code pollExecutor}
|
||||
* and have the standard {@code mate_async_task} lifecycle (running → terminal
|
||||
* + automatic cancel-on-parent-conversation-deleted + bookkeeping cleanup) take
|
||||
* over.
|
||||
* <p>
|
||||
* The fixture installs an anonymous subclass that overrides {@code createTask},
|
||||
* {@code updateStatus} and {@code findEntityByTaskId} with in-memory equivalents.
|
||||
* Reason: the production methods round-trip through a MyBatis-Plus
|
||||
* LambdaUpdateWrapper that obscures the {@code (status, progress, resultJson,
|
||||
* errorMessage)} tuple under MPGENVAL placeholders. Asserting on a
|
||||
* test-controlled state map is faster, less brittle, and lets the suite focus
|
||||
* on what {@code submitOneShot} actually does — schedule the worker, register
|
||||
* bookkeeping, observe cancellation, and clean up.
|
||||
*
|
||||
* <h3>Covered paths</h3>
|
||||
* <ol>
|
||||
* <li>Success — Callable returns, status lands on succeeded with resultJson.</li>
|
||||
* <li>Exception — Callable throws, status lands on failed with errorMessage.</li>
|
||||
* <li>Conversation deletion mid-run — listener writes failed + worker's
|
||||
* second cancel-check also writes failed; both messages match.</li>
|
||||
* <li>schedule/put race stress — 200 zero-cost tasks all succeed and drain
|
||||
* both bookkeeping maps to empty (catches any ghost entry).</li>
|
||||
* <li>Latch ordering — worker observes itself enrolled in both maps the
|
||||
* moment Callable.call() begins (the latch invariant).</li>
|
||||
* </ol>
|
||||
*/
|
||||
class AsyncTaskServiceOneShotTest {
|
||||
|
||||
private AsyncTaskMapper mapper;
|
||||
private ChatStreamTracker tracker;
|
||||
|
||||
/** Test-side per-taskId snapshot — populated by the fixture's overrides
|
||||
* of createTask + updateStatus instead of going through MyBatis-Plus. */
|
||||
private static final class TaskState {
|
||||
String taskId;
|
||||
String taskType;
|
||||
String conversationId;
|
||||
String status;
|
||||
Integer progress;
|
||||
String resultJson;
|
||||
String errorMessage;
|
||||
}
|
||||
|
||||
private final ConcurrentMap<String, TaskState> states = new ConcurrentHashMap<>();
|
||||
private AsyncTaskService service;
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
mapper = mock(AsyncTaskMapper.class);
|
||||
tracker = mock(ChatStreamTracker.class);
|
||||
states.clear();
|
||||
service = new AsyncTaskService(mapper, tracker) {
|
||||
@Override
|
||||
public AsyncTaskEntity createTask(String taskType, String conversationId, Long messageId,
|
||||
String providerName, String providerTaskId,
|
||||
String requestJson, String createdBy) {
|
||||
String taskId = UUID.randomUUID().toString().replace("-", "").substring(0, 16);
|
||||
AsyncTaskEntity e = new AsyncTaskEntity();
|
||||
e.setTaskId(taskId);
|
||||
e.setTaskType(taskType);
|
||||
e.setStatus("pending");
|
||||
e.setConversationId(conversationId);
|
||||
e.setMessageId(messageId);
|
||||
e.setProviderName(providerName);
|
||||
e.setProviderTaskId(providerTaskId);
|
||||
e.setRequestJson(requestJson);
|
||||
e.setCreatedBy(createdBy);
|
||||
TaskState s = new TaskState();
|
||||
s.taskId = taskId;
|
||||
s.taskType = taskType;
|
||||
s.conversationId = conversationId;
|
||||
s.status = "pending";
|
||||
states.put(taskId, s);
|
||||
return e;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void updateStatus(String taskId, String status, Integer progress,
|
||||
String resultJson, String errorMessage) {
|
||||
states.compute(taskId, (k, old) -> {
|
||||
TaskState s = old != null ? old : new TaskState();
|
||||
s.taskId = taskId;
|
||||
s.status = status;
|
||||
if (progress != null) s.progress = progress;
|
||||
if (resultJson != null) s.resultJson = resultJson;
|
||||
if (errorMessage != null) s.errorMessage = errorMessage;
|
||||
return s;
|
||||
});
|
||||
}
|
||||
|
||||
@Override
|
||||
public AsyncTaskEntity findEntityByTaskId(String taskId) {
|
||||
TaskState s = states.get(taskId);
|
||||
if (s == null) return null;
|
||||
AsyncTaskEntity e = new AsyncTaskEntity();
|
||||
e.setTaskId(s.taskId);
|
||||
e.setTaskType(s.taskType);
|
||||
e.setStatus(s.status);
|
||||
e.setConversationId(s.conversationId);
|
||||
return e;
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
@AfterEach
|
||||
void tearDown() {
|
||||
service.shutdown();
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("Success path: status = succeeded, resultJson captures Callable return, maps drain")
|
||||
void successPath() throws Exception {
|
||||
AsyncTaskEntity entity = service.submitOneShot(
|
||||
"agent_delegate", "conv-success", null, "{}", "user-1",
|
||||
() -> "ok");
|
||||
|
||||
awaitDone(entity.getTaskId(), 5_000);
|
||||
|
||||
TaskState s = states.get(entity.getTaskId());
|
||||
assertThat(s.status).isEqualTo("succeeded");
|
||||
assertThat(s.resultJson).isEqualTo("ok");
|
||||
assertThat(s.progress).isEqualTo(100);
|
||||
assertActiveMapsEmpty();
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("Exception path: status = failed, errorMessage carries Callable's message")
|
||||
void exceptionPath() throws Exception {
|
||||
AsyncTaskEntity entity = service.submitOneShot(
|
||||
"agent_delegate", "conv-fail", null, "{}", "user-1",
|
||||
() -> { throw new RuntimeException("boom-msg"); });
|
||||
|
||||
awaitDone(entity.getTaskId(), 5_000);
|
||||
|
||||
TaskState s = states.get(entity.getTaskId());
|
||||
assertThat(s.status).isEqualTo("failed");
|
||||
assertThat(s.errorMessage).isNotNull().contains("boom-msg");
|
||||
assertActiveMapsEmpty();
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("Conversation deleted while running: terminal status is failed with deletion message")
|
||||
void cancelPathViaConversationDeleted() throws Exception {
|
||||
CountDownLatch workerStarted = new CountDownLatch(1);
|
||||
CountDownLatch workerCanProceed = new CountDownLatch(1);
|
||||
|
||||
String convId = "conv-cancel";
|
||||
AsyncTaskEntity entity = service.submitOneShot(
|
||||
"agent_delegate", convId, null, "{}", "user-1",
|
||||
() -> {
|
||||
workerStarted.countDown();
|
||||
// Hold the worker inside work.call() until the test fires
|
||||
// the deletion event. Future.cancel(false) — what the
|
||||
// listener calls — does NOT interrupt, so this await won't
|
||||
// unblock until the test's explicit countDown below.
|
||||
workerCanProceed.await(5, TimeUnit.SECONDS);
|
||||
return "should-not-be-applied";
|
||||
});
|
||||
|
||||
assertThat(workerStarted.await(5, TimeUnit.SECONDS)).isTrue();
|
||||
|
||||
// Listener runs synchronously on the test thread: it cancels the
|
||||
// future, looks up the (still-running) entity, sees taskType
|
||||
// agent_delegate and writes failed + "conversation deleted". This
|
||||
// closes the contract gap that bare cancelPolling left open before
|
||||
// this PR (DB row would otherwise stay running forever).
|
||||
service.onConversationDeleted(new ConversationDeletedEvent(convId));
|
||||
|
||||
// Release the worker so it observes isConversationCanceled = true at
|
||||
// the post-call cancel check and writes the during-execution variant
|
||||
// (also matches "conversation deleted").
|
||||
workerCanProceed.countDown();
|
||||
awaitDone(entity.getTaskId(), 5_000);
|
||||
|
||||
TaskState s = states.get(entity.getTaskId());
|
||||
assertThat(s.status).isEqualTo("failed");
|
||||
assertThat(s.errorMessage).isNotNull().contains("conversation deleted");
|
||||
assertActiveMapsEmpty();
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("Terminal status broadcasts async_task_completed on the parent conversation")
|
||||
void broadcastsCompletionEvent() throws Exception {
|
||||
AsyncTaskEntity successEntity = service.submitOneShot(
|
||||
"agent_delegate", "conv-broadcast-ok", null, "{}", "user-1",
|
||||
() -> "ok");
|
||||
AsyncTaskEntity failEntity = service.submitOneShot(
|
||||
"agent_delegate", "conv-broadcast-fail", null, "{}", "user-1",
|
||||
() -> { throw new RuntimeException("boom"); });
|
||||
|
||||
awaitDone(successEntity.getTaskId(), 5_000);
|
||||
awaitDone(failEntity.getTaskId(), 5_000);
|
||||
|
||||
// Success path → event with success=true, no errorMessage.
|
||||
verify(tracker, timeout(2_000)).broadcastObject(
|
||||
eq("conv-broadcast-ok"), eq("async_task_completed"), any());
|
||||
// Failure path → event with success=false, errorMessage carries
|
||||
// the Callable's exception message. Broadcast routes to the parent
|
||||
// conversation_id stored on the entity, which IS the parent for
|
||||
// agent-delegate one-shots per AsyncTaskService.submitOneShot.
|
||||
verify(tracker, timeout(2_000)).broadcastObject(
|
||||
eq("conv-broadcast-fail"), eq("async_task_completed"), any());
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("schedule/put race stress: 200 zero-cost tasks all succeed and bookkeeping drains")
|
||||
void scheduleAndPutRaceStress() throws Exception {
|
||||
int iterations = 200;
|
||||
AsyncTaskEntity[] entities = new AsyncTaskEntity[iterations];
|
||||
for (int i = 0; i < iterations; i++) {
|
||||
entities[i] = service.submitOneShot(
|
||||
"agent_delegate", "conv-race-" + i, null, "{}", "user-1",
|
||||
() -> "ok");
|
||||
}
|
||||
for (AsyncTaskEntity e : entities) {
|
||||
awaitDone(e.getTaskId(), 15_000);
|
||||
}
|
||||
// All terminal statuses must be succeeded.
|
||||
for (AsyncTaskEntity e : entities) {
|
||||
assertThat(states.get(e.getTaskId()).status)
|
||||
.as("task %s succeeded", e.getTaskId())
|
||||
.isEqualTo("succeeded");
|
||||
}
|
||||
// Belt-and-suspenders: any ghost (future.isDone() but key still in
|
||||
// map) would keep one or both of these non-empty.
|
||||
awaitMapsDrained(5_000);
|
||||
assertActiveMapsEmpty();
|
||||
}
|
||||
|
||||
@Test
|
||||
@DisplayName("Latch ordering: Callable.call() observes its taskId enrolled in both maps")
|
||||
void enrolledLatchOrdering() throws Exception {
|
||||
ConcurrentHashMap<String, ?> activePolls = getInternalMap("activePolls");
|
||||
ConcurrentHashMap<String, ?> pollTaskToConv = getInternalMap("pollTaskToConv");
|
||||
|
||||
AtomicReference<String> taskIdRef = new AtomicReference<>();
|
||||
AtomicBoolean observedActive = new AtomicBoolean();
|
||||
AtomicBoolean observedConvLink = new AtomicBoolean();
|
||||
CountDownLatch probed = new CountDownLatch(1);
|
||||
|
||||
// The Callable acts as a probe: if the latch invariant holds, both
|
||||
// bookkeeping maps already contain this taskId by the time the body
|
||||
// runs (calling thread put before countDown, worker awaited
|
||||
// countDown before reaching here). If the latch were removed, the
|
||||
// worker could race ahead, observe empty maps, and finish — then the
|
||||
// calling thread's put would leak a ghost entry.
|
||||
Callable<String> work = () -> {
|
||||
String tid = taskIdRef.get();
|
||||
observedActive.set(tid != null && activePolls.containsKey(tid));
|
||||
observedConvLink.set(tid != null && pollTaskToConv.containsKey(tid));
|
||||
probed.countDown();
|
||||
return "ok";
|
||||
};
|
||||
|
||||
AsyncTaskEntity entity = service.submitOneShot(
|
||||
"agent_delegate", "conv-latch", null, "{}", "user-1", work);
|
||||
taskIdRef.set(entity.getTaskId());
|
||||
|
||||
assertThat(probed.await(5, TimeUnit.SECONDS))
|
||||
.as("probe must run within timeout")
|
||||
.isTrue();
|
||||
awaitDone(entity.getTaskId(), 5_000);
|
||||
|
||||
assertThat(observedActive)
|
||||
.as("activePolls must contain taskId when work.call() begins")
|
||||
.isTrue();
|
||||
assertThat(observedConvLink)
|
||||
.as("pollTaskToConv must contain taskId when work.call() begins")
|
||||
.isTrue();
|
||||
assertActiveMapsEmpty();
|
||||
}
|
||||
|
||||
// ---------- helpers ----------
|
||||
|
||||
/** Waits until the task has reached a terminal status AND the worker's
|
||||
* finally cleanup has cleared this taskId from both bookkeeping maps. */
|
||||
private void awaitDone(String taskId, long timeoutMs) throws InterruptedException {
|
||||
awaitUntil(() -> {
|
||||
TaskState s = states.get(taskId);
|
||||
if (s == null) return false;
|
||||
if (!"succeeded".equals(s.status) && !"failed".equals(s.status)) return false;
|
||||
return !getInternalMap("activePolls").containsKey(taskId)
|
||||
&& !getInternalMap("pollTaskToConv").containsKey(taskId);
|
||||
}, timeoutMs, "task " + taskId + " never reached cleaned-up terminal state");
|
||||
}
|
||||
|
||||
private void awaitMapsDrained(long timeoutMs) throws InterruptedException {
|
||||
awaitUntil(() -> getInternalMap("activePolls").isEmpty()
|
||||
&& getInternalMap("pollTaskToConv").isEmpty(),
|
||||
timeoutMs, "activePolls / pollTaskToConv never drained");
|
||||
}
|
||||
|
||||
private void awaitUntil(BooleanSupplier cond, long timeoutMs, String message)
|
||||
throws InterruptedException {
|
||||
long deadline = System.currentTimeMillis() + timeoutMs;
|
||||
while (System.currentTimeMillis() < deadline) {
|
||||
if (cond.getAsBoolean()) return;
|
||||
Thread.sleep(20);
|
||||
}
|
||||
throw new AssertionError(message + " (waited " + timeoutMs + "ms)");
|
||||
}
|
||||
|
||||
private void assertActiveMapsEmpty() {
|
||||
assertThat(getInternalMap("activePolls"))
|
||||
.as("activePolls should be empty after all workers finish")
|
||||
.isEmpty();
|
||||
assertThat(getInternalMap("pollTaskToConv"))
|
||||
.as("pollTaskToConv should be empty after all workers finish")
|
||||
.isEmpty();
|
||||
}
|
||||
|
||||
@SuppressWarnings("unchecked")
|
||||
private ConcurrentHashMap<String, ?> getInternalMap(String name) {
|
||||
return (ConcurrentHashMap<String, ?>) ReflectionTestUtils.getField(service, name);
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user