Move onFinishCallback before span or transaction is finished (#3459)

* moved span and transaction finish callback before they are actually finished
* moved performance collection before transaction is finished, fixing frame measurements not being set

Author: stefanosiano

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Fixes
 
+- Move onFinishCallback before span or transaction is finished ([#3459](https://github.com/getsentry/sentry-java/pull/3459))
 - Add timestamp when a profile starts ([#3442](https://github.com/getsentry/sentry-java/pull/3442))
 - Move fragment auto span finish to onFragmentStarted ([#3424](https://github.com/getsentry/sentry-java/pull/3424))
 - Remove profiling timeout logic and disable profiling on API 21 ([#3478](https://github.com/getsentry/sentry-java/pull/3478))

sentry-android-core/src/test/java/io/sentry/android/core/ActivityLifecycleIntegrationTest.kt

@@ -303,7 +303,7 @@ class ActivityLifecycleIntegrationTest {
         sut.ttidSpanMap.values.first().finish()
         sut.ttfdSpanMap.values.first().finish()
 
-        // then transaction should not be immediatelly finished
+        // then transaction should not be immediately finished
         verify(fixture.hub, never())
             .captureTransaction(
                 anyOrNull(),

sentry/src/main/java/io/sentry/SentryTracer.java

@@ -200,26 +200,45 @@ public void finish(
     if (!root.isFinished()
         && (!transactionOptions.isWaitForChildren() || hasAllChildrenFinished())) {
 
+      final @NotNull AtomicReference<List<PerformanceCollectionData>> performanceCollectionData =
+          new AtomicReference<>();
+      // We set the new spanFinishedCallback here instead of creation time, calling the old one to
+      // avoid the user overwrites it by setting a custom spanFinishedCallback on the root span
+      final @Nullable SpanFinishedCallback oldCallback = this.root.getSpanFinishedCallback();
+      this.root.setSpanFinishedCallback(
+          span -> {
+            if (oldCallback != null) {
+              oldCallback.execute(span);
+            }
+
+            // Let's call the finishCallback here, when the root span has a finished date but it's
+            // not finished, yet
+            final @Nullable TransactionFinishedCallback finishedCallback =
+                transactionOptions.getTransactionFinishedCallback();
+            if (finishedCallback != null) {
+              finishedCallback.execute(this);
+            }
+
+            if (transactionPerformanceCollector != null) {
+              performanceCollectionData.set(transactionPerformanceCollector.stop(this));
+            }
+          });
+
       // any un-finished childs will remain unfinished
       // as relay takes care of setting the end-timestamp + deadline_exceeded
       // see
       // https://github.com/getsentry/relay/blob/40697d0a1c54e5e7ad8d183fc7f9543b94fe3839/relay-general/src/store/transactions/processor.rs#L374-L378
       root.finish(finishStatus.spanStatus, finishTimestamp);
 
-      List<PerformanceCollectionData> performanceCollectionData = null;
-      if (transactionPerformanceCollector != null) {
-        performanceCollectionData = transactionPerformanceCollector.stop(this);
-      }
-
       ProfilingTraceData profilingTraceData = null;
       if (Boolean.TRUE.equals(isSampled()) && Boolean.TRUE.equals(isProfileSampled())) {
         profilingTraceData =
             hub.getOptions()
                 .getTransactionProfiler()
-                .onTransactionFinish(this, performanceCollectionData, hub.getOptions());
+                .onTransactionFinish(this, performanceCollectionData.get(), hub.getOptions());
       }
-      if (performanceCollectionData != null) {
-        performanceCollectionData.clear();
+      if (performanceCollectionData.get() != null) {
+        performanceCollectionData.get().clear();
       }
 
       hub.configureScope(
@@ -232,11 +251,6 @@ public void finish(
                 });
           });
       final SentryTransaction transaction = new SentryTransaction(this);
-      final TransactionFinishedCallback finishedCallback =
-          transactionOptions.getTransactionFinishedCallback();
-      if (finishedCallback != null) {
-        finishedCallback.execute(this);
-      }
 
       if (timer != null) {
         synchronized (timerLock) {
@@ -605,7 +619,9 @@ private boolean hasAllChildrenFinished() {
     final List<Span> spans = new ArrayList<>(this.children);
     if (!spans.isEmpty()) {
       for (final Span span : spans) {
-        if (!span.isFinished()) {
+        // This is used in the spanFinishCallback, when the span isn't finished, but has a finish
+        // date
+        if (!span.isFinished() && span.getFinishDate() == null) {
           return false;
         }
       }

sentry/src/main/java/io/sentry/Span.java

@@ -37,7 +37,9 @@ public final class Span implements ISpan {
 
   private final @NotNull IHub hub;
 
-  private final @NotNull AtomicBoolean finished = new AtomicBoolean(false);
+  private boolean finished = false;
+
+  private final @NotNull AtomicBoolean isFinishing = new AtomicBoolean(false);
 
   private final @NotNull SpanOptions options;
 
@@ -122,7 +124,7 @@ public Span(
       final @Nullable SentryDate timestamp,
       final @NotNull Instrumenter instrumenter,
       @NotNull SpanOptions spanOptions) {
-    if (finished.get()) {
+    if (finished) {
       return NoOpSpan.getInstance();
     }
 
@@ -133,7 +135,7 @@ public Span(
   @Override
   public @NotNull ISpan startChild(
       final @NotNull String operation, final @Nullable String description) {
-    if (finished.get()) {
+    if (finished) {
       return NoOpSpan.getInstance();
     }
 
@@ -143,7 +145,7 @@ public Span(
   @Override
   public @NotNull ISpan startChild(
       @NotNull String operation, @Nullable String description, @NotNull SpanOptions spanOptions) {
-    if (finished.get()) {
+    if (finished) {
       return NoOpSpan.getInstance();
     }
     return transaction.startChild(context.getSpanId(), operation, description, spanOptions);
@@ -192,7 +194,7 @@ public void finish(@Nullable SpanStatus status) {
   @Override
   public void finish(final @Nullable SpanStatus status, final @Nullable SentryDate timestamp) {
     // the span can be finished only once
-    if (!finished.compareAndSet(false, true)) {
+    if (finished || !isFinishing.compareAndSet(false, true)) {
       return;
     }
 
@@ -235,6 +237,7 @@ public void finish(final @Nullable SpanStatus status, final @Nullable SentryDate
     if (spanFinishedCallback != null) {
       spanFinishedCallback.execute(this);
     }
+    finished = true;
   }
 
   @Override
@@ -284,7 +287,7 @@ public void setTag(final @NotNull String key, final @NotNull String value) {
 
   @Override
   public boolean isFinished() {
-    return finished.get();
+    return finished;
   }
 
   public @NotNull Map<String, Object> getData() {
@@ -409,6 +412,11 @@ void setSpanFinishedCallback(final @Nullable SpanFinishedCallback callback) {
     this.spanFinishedCallback = callback;
   }
 
+  @Nullable
+  SpanFinishedCallback getSpanFinishedCallback() {
+    return spanFinishedCallback;
+  }
+
   private void updateStartDate(@NotNull SentryDate date) {
     this.startTimestamp = date;
   }

sentry/src/test/java/io/sentry/SentryTracerTest.kt

@@ -1378,4 +1378,15 @@ class SentryTracerTest {
 
         assertEquals(5, tracer.children.size)
     }
+
+    @Test
+    fun `tracer is not finished when finishCallback is called`() {
+        val transaction = fixture.getSut(transactionFinishedCallback = {
+            assertFalse(it.isFinished)
+            assertNotNull(it.finishDate)
+        })
+        assertFalse(transaction.isFinished)
+        assertNull(transaction.finishDate)
+        transaction.finish()
+    }
 }

sentry/src/test/java/io/sentry/SpanTest.kt

@@ -495,6 +495,19 @@ class SpanTest {
         assertSame(span.localMetricsAggregator, span.localMetricsAggregator)
     }
 
+    // test to ensure that the span is not finished when the finishCallback is called
+    @Test
+    fun `span is not finished when finishCallback is called`() {
+        val span = fixture.getSut()
+        span.setSpanFinishedCallback {
+            assertFalse(span.isFinished)
+            assertNotNull(span.finishDate)
+        }
+        assertFalse(span.isFinished)
+        assertNull(span.finishDate)
+        span.finish()
+    }
+
     private fun getTransaction(transactionContext: TransactionContext = TransactionContext("name", "op")): SentryTracer {
         return SentryTracer(transactionContext, fixture.hub)
     }