@Trace annotation does not work with kotlin suspending functions

[joedj]

As far as I can tell, this is supposed to work (https://github.com/DataDog/dd-trace-java/issues/931, https://github.com/DataDog/dd-trace-java/pull/4405), but it doesn't - traceId changes when a suspended function resumes.

I've prepared a minimal failing example here: https://github.com/joedj/datadog-coroutine-trace-test

In the example I'm using dd-java-agent 1.20.1, kotlin 1.9.10 and kotlinx-coroutines-core 1.7.3, but it fails the same way on every version I've tried.

    @Trace
    suspend fun suspending(): Pair<String, String> {
        val tracer = GlobalTracer.get()
        val traceId1 = tracer.traceId
        delay(1)
        val traceId2 = tracer.traceId
        return traceId1 to traceId2
    }

    @Test
    fun `trace annotated suspend function`(): Unit = runBlocking {
        // when
        val (traceId1, traceId2) = suspending()
        // then
        assertNotNull(traceId1)
        assertNotEquals(traceId1, "")
        assertNotEquals(traceId1, "0")
        assertEquals(traceId1, traceId2)
    }

[monosoul]

I think I have an idea of what goes wrong there.

The suspending function from the example will compile into a bytecode like this:

   @Trace
   @Nullable
   public final Object suspending(@NotNull Continuation var1) {

full decompiled bytecode example

   @Trace
   @Nullable
   public final Object suspending(@NotNull Continuation var0) {
      Object $continuation;
      label20: {
         if (var0 instanceof <undefinedtype>) {
            $continuation = (<undefinedtype>)var0;
            if ((((<undefinedtype>)$continuation).label & Integer.MIN_VALUE) != 0) {
               ((<undefinedtype>)$continuation).label -= Integer.MIN_VALUE;
               break label20;
            }
         }

         $continuation = new ContinuationImpl(var0) {
            // $FF: synthetic field
            Object result;
            int label;
            Object L$0;
            Object L$1;

            @Nullable
            public final Object invokeSuspend(@NotNull Object $result) {
               this.result = $result;
               this.label |= Integer.MIN_VALUE;
               return TheAppKt.suspending(this);
            }
         };
      }

      Object $result = ((<undefinedtype>)$continuation).result;
      Object var6 = IntrinsicsKt.getCOROUTINE_SUSPENDED();
      Tracer tracer;
      String traceId1;
      switch (((<undefinedtype>)$continuation).label) {
         case 0:
            ResultKt.throwOnFailure($result);
            tracer = GlobalTracer.get();
            Intrinsics.checkNotNullExpressionValue(tracer, "tracer");
            traceId1 = tracer.getTraceId();
            ((<undefinedtype>)$continuation).L$0 = tracer;
            ((<undefinedtype>)$continuation).L$1 = traceId1;
            ((<undefinedtype>)$continuation).label = 1;
            if (DelayKt.delay(1L, (Continuation)$continuation) == var6) {
               return var6;
            }
            break;
         case 1:
            traceId1 = (String)((<undefinedtype>)$continuation).L$1;
            tracer = (Tracer)((<undefinedtype>)$continuation).L$0;
            ResultKt.throwOnFailure($result);
            break;
         default:
            throw new IllegalStateException("call to 'resume' before 'invoke' with coroutine");
      }

      Intrinsics.checkNotNullExpressionValue(tracer, "tracer");
      String traceId2 = tracer.getTraceId();
      return TuplesKt.to(traceId1, traceId2);
   }

On each continuation this method will be invoked again, so if it creates a top-level span, then each continuation will create a separate span. I'm not sure if there's an easy way to make it work without implementing continuation handling similar to the bytecode Kotlin will generate for that function, as we'll need to keep track of if a call to that method is an initial call or if it's a continuation etc. Probably it's easier to just create spans programmatically (see this comment: https://github.com/DataDog/dd-trace-java/issues/931#issuecomment-1485146797 ).

@bantonsson what do you think?

[samboehle]

We are having this issue along with the issue that is shown in this example: https://github.com/DataDog/dd-trace-java/issues/931#issuecomment-1485146797 that causes the span to cut short of the full duration. We would prefer to have the @Trace annotation work as expected instead of using custom spans.

[HiteshGarg]

Our spans looks broken because Trace did not work properly with the suspended functions. Is there any plan for this issue?