Skip to content

Commit

Permalink
extract context from reactor
Browse files Browse the repository at this point in the history
  • Loading branch information
laurit committed Jan 24, 2022
1 parent d7e910a commit 11810cd
Show file tree
Hide file tree
Showing 6 changed files with 115 additions and 59 deletions.
7 changes: 4 additions & 3 deletions instrumentation/kotlinx-coroutines/javaagent/build.gradle.kts
Original file line number Diff line number Diff line change
Expand Up @@ -21,12 +21,13 @@ muzzle {
dependencies {
compileOnly("io.opentelemetry:opentelemetry-extension-kotlin")
compileOnly("org.jetbrains.kotlin:kotlin-stdlib-jdk8")
// Use first version with flow support since we have tests for it.
library("org.jetbrains.kotlinx:kotlinx-coroutines-core:1.3.0")
library("org.jetbrains.kotlinx:kotlinx-coroutines-reactor:1.3.0")
implementation(project(":instrumentation:reactor:reactor-3.1:library"))

testImplementation("io.opentelemetry:opentelemetry-extension-kotlin")
testImplementation("org.jetbrains.kotlin:kotlin-stdlib-jdk8")
// Use first version with flow support since we have tests for it.
testLibrary("org.jetbrains.kotlinx:kotlinx-coroutines-core:1.3.0")
testLibrary("org.jetbrains.kotlinx:kotlinx-coroutines-reactor:1.3.0")
}

tasks {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,11 @@

package io.opentelemetry.javaagent.instrumentation.kotlinxcoroutines;

import static java.util.Arrays.asList;
import static java.util.Collections.singletonList;

import com.google.auto.service.AutoService;
import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.instrumentation.kotlinxcoroutines.reactor.KotlinCoroutinesFluxInstrumentation;
import io.opentelemetry.javaagent.instrumentation.kotlinxcoroutines.reactor.KotlinCoroutinesMonoInstrumentation;
import java.util.List;

@AutoService(InstrumentationModule.class)
Expand All @@ -28,9 +26,6 @@ public boolean isHelperClass(String className) {

@Override
public List<TypeInstrumentation> typeInstrumentations() {
return asList(
new KotlinCoroutinesInstrumentation(),
new KotlinCoroutinesMonoInstrumentation(),
new KotlinCoroutinesFluxInstrumentation());
return singletonList(new KotlinCoroutinesInstrumentation());
}
}

This file was deleted.

Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
/*
* Copyright The OpenTelemetry Authors
* SPDX-License-Identifier: Apache-2.0
*/

package io.opentelemetry.javaagent.instrumentation.kotlinxcoroutines.reactor;

import static java.util.Arrays.asList;

import com.google.auto.service.AutoService;
import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import java.util.List;

@AutoService(InstrumentationModule.class)
public class KotlinCoroutinesReactorInstrumentationModule extends InstrumentationModule {

public KotlinCoroutinesReactorInstrumentationModule() {
super("kotlinx-coroutines", "kotlinx-coroutines-reactor");
}

@Override
public boolean isHelperClass(String className) {
return className.startsWith("io.opentelemetry.extension.kotlin.");
}

@Override
public List<TypeInstrumentation> typeInstrumentations() {
return asList(
new KotlinMonoCoroutineInstrumentation(), new KotlinPublisherCoroutineInstrumentation());
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -5,37 +5,49 @@

package io.opentelemetry.javaagent.instrumentation.kotlinxcoroutines.reactor;

import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
import static net.bytebuddy.matcher.ElementMatchers.named;
import static net.bytebuddy.matcher.ElementMatchers.namedOneOf;
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;

import io.opentelemetry.context.Context;
import io.opentelemetry.instrumentation.reactor.ContextPropagationOperator;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
import io.opentelemetry.javaagent.instrumentation.kotlinxcoroutines.KotlinCoroutinesInstrumentationHelper;
import kotlin.coroutines.CoroutineContext;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;
import reactor.core.publisher.MonoSink;

public class KotlinCoroutinesFluxInstrumentation implements TypeInstrumentation {
public class KotlinMonoCoroutineInstrumentation implements TypeInstrumentation {
@Override
public ElementMatcher<TypeDescription> typeMatcher() {
return named("kotlinx.coroutines.reactor.FluxKt");
return namedOneOf("kotlinx.coroutines.reactor.MonoCoroutine");
}

@Override
public void transform(TypeTransformer transformer) {
transformer.applyAdviceToMethod(
namedOneOf("flux").and(takesArgument(0, named("kotlin.coroutines.CoroutineContext"))),
this.getClass().getName() + "$FluxAdvice");
isConstructor()
.and(
takesArgument(0, named("kotlin.coroutines.CoroutineContext"))
.and(takesArgument(1, named("reactor.core.publisher.MonoSink")))),
this.getClass().getName() + "$MonoCoroutineAdvice");
}

@SuppressWarnings("unused")
public static class FluxAdvice {
public static class MonoCoroutineAdvice {

@Advice.OnMethodEnter
public static void enter(
@Advice.Argument(value = 0, readOnly = false) CoroutineContext coroutineContext) {
@Advice.Argument(value = 0, readOnly = false) CoroutineContext coroutineContext,
@Advice.Argument(1) MonoSink monoSink) {
Context context =
ContextPropagationOperator.getOpenTelemetryContext(
monoSink.currentContext(), Java8BytecodeBridge.currentContext());
coroutineContext =
KotlinCoroutinesInstrumentationHelper.addOpenTelemetryContext(coroutineContext);
}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
/*
* Copyright The OpenTelemetry Authors
* SPDX-License-Identifier: Apache-2.0
*/

package io.opentelemetry.javaagent.instrumentation.kotlinxcoroutines.reactor;

import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
import static net.bytebuddy.matcher.ElementMatchers.named;
import static net.bytebuddy.matcher.ElementMatchers.namedOneOf;
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;

import io.opentelemetry.context.Context;
import io.opentelemetry.instrumentation.reactor.ContextPropagationOperator;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
import io.opentelemetry.javaagent.instrumentation.kotlinxcoroutines.KotlinCoroutinesInstrumentationHelper;
import kotlin.coroutines.CoroutineContext;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;
import org.reactivestreams.Subscriber;
import reactor.core.CoreSubscriber;

public class KotlinPublisherCoroutineInstrumentation implements TypeInstrumentation {
@Override
public ElementMatcher<TypeDescription> typeMatcher() {
return namedOneOf("kotlinx.coroutines.reactive.PublisherCoroutine");
}

@Override
public void transform(TypeTransformer transformer) {
transformer.applyAdviceToMethod(
isConstructor()
.and(
takesArgument(0, named("kotlin.coroutines.CoroutineContext"))
.and(takesArgument(1, named("org.reactivestreams.Subscriber")))),
this.getClass().getName() + "$PublisherCoroutineAdvice");
}

@SuppressWarnings("unused")
public static class PublisherCoroutineAdvice {

@Advice.OnMethodEnter
public static void enter(
@Advice.Argument(value = 0, readOnly = false) CoroutineContext coroutineContext,
@Advice.Argument(1) Subscriber<?> subscriber) {
if (subscriber instanceof CoreSubscriber) {
CoreSubscriber<?> coreSubscriber = (CoreSubscriber) subscriber;
Context context =
ContextPropagationOperator.getOpenTelemetryContext(
coreSubscriber.currentContext(), Java8BytecodeBridge.currentContext());
coroutineContext =
KotlinCoroutinesInstrumentationHelper.addOpenTelemetryContext(coroutineContext);
}
}
}
}

0 comments on commit 11810cd

Please sign in to comment.