Task may be unregistered outside of the trace context in exceptional cases. (#137865) (#138004)

This change makes sure we're stopping traces in these cases so we don't risk
leaking spans in APMTracer.

Author: mosche

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 <Request extends ActionRequest, Response extends ActionResponse> Task registerAndExecute(
         String type,
         TransportAction<Request, Response> 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);
             }

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() {