diff --git a/server/src/main/java/org/elasticsearch/tasks/TaskManager.java b/server/src/main/java/org/elasticsearch/tasks/TaskManager.java index 2ed347c226870..b6216d4b7db12 100644 --- a/server/src/main/java/org/elasticsearch/tasks/TaskManager.java +++ b/server/src/main/java/org/elasticsearch/tasks/TaskManager.java @@ -191,12 +191,6 @@ void maybeStartTrace(ThreadContext threadContext, Task task) { tracer.startTrace(threadContext, task, task.getAction(), attributes); } - void maybeStopTrace(ThreadContext threadContext, Task task) { - if (threadContext.hasTraceContext()) { - tracer.stopTrace(task); - } - } - public Task registerAndExecute( String type, TransportAction action, @@ -358,7 +352,7 @@ public Task unregister(Task task) { return removedTask; } } finally { - maybeStopTrace(threadPool.getThreadContext(), task); + tracer.stopTrace(task); // stop trace if started / known by tracer for (RemovedTaskListener listener : removedTaskListeners) { listener.onRemoved(task); } diff --git a/server/src/test/java/org/elasticsearch/tasks/TaskManagerTests.java b/server/src/test/java/org/elasticsearch/tasks/TaskManagerTests.java index 86436a0852f58..f38af77be6150 100644 --- a/server/src/test/java/org/elasticsearch/tasks/TaskManagerTests.java +++ b/server/src/test/java/org/elasticsearch/tasks/TaskManagerTests.java @@ -69,6 +69,7 @@ import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.when; public class TaskManagerTests extends ESTestCase { @@ -311,6 +312,9 @@ public TaskId getParentTask() { ? Map.of(Tracer.AttributeKeys.TASK_ID, task.getId(), Tracer.AttributeKeys.PARENT_TASK_ID, parentTask.toString()) : Map.of(Tracer.AttributeKeys.TASK_ID, task.getId()); verify(mockTracer).startTrace(any(), eq(task), eq("testAction"), eq(attributes)); + + taskManager.unregister(task); + verify(mockTracer).stopTrace(task); // always attempt stopping to guard against leaks } } @@ -393,7 +397,8 @@ public TaskId getParentTask() { // no trace context taskManager.unregister(task); - verifyNoInteractions(mockTracer); + verify(mockTracer).stopTrace(task); // always attempt stopping to guard against leaks + verifyNoMoreInteractions(mockTracer); } /** @@ -438,6 +443,7 @@ public TaskId getParentTask() { ); verify(mockTracer).startTrace(any(), eq(task), eq("actionName"), anyMap()); + verify(mockTracer).stopTrace(task); // always attempt stopping to guard against leaks } /** @@ -480,7 +486,8 @@ public TaskId getParentTask() { ActionTestUtils.assertNoFailureListener(r -> {}) ); - verifyNoInteractions(mockTracer); + verify(mockTracer).stopTrace(task); // always attempt stopping to guard against leaks + verifyNoMoreInteractions(mockTracer); } public void testRegisterWithEnabledDisabledTracing() {