Skip to content

Commit

Permalink
Fix flaky finagle test
Browse files Browse the repository at this point in the history
  • Loading branch information
laurit committed May 23, 2024
1 parent a1400dd commit 9028dae
Show file tree
Hide file tree
Showing 3 changed files with 73 additions and 2 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
/*
* Copyright The OpenTelemetry Authors
* SPDX-License-Identifier: Apache-2.0
*/

package io.opentelemetry.javaagent.instrumentation.finaglehttp.v23_11;

import io.opentelemetry.context.Context;
import io.opentelemetry.context.Scope;
import scala.Function1;

public final class Function1Wrapper {

public static <T1, R> Function1<T1, R> wrap(Function1<T1, R> function1) {
Context context = Context.current();
return value -> {
try (Scope ignored = context.makeCurrent()) {
return function1.apply(value);
}
};
}

private Function1Wrapper() {}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
/*
* Copyright The OpenTelemetry Authors
* SPDX-License-Identifier: Apache-2.0
*/

package io.opentelemetry.javaagent.instrumentation.finaglehttp.v23_11;

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

import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;
import scala.Function1;

public class PromiseMonitoredInstrumentation implements TypeInstrumentation {

@Override
public ElementMatcher<TypeDescription> typeMatcher() {
return named("com.twitter.util.Promise$Monitored");
}

@Override
public void transform(TypeTransformer transformer) {
transformer.applyAdviceToMethod(
isConstructor().and(takesArgument(1, named("scala.Function1"))),
this.getClass().getName() + "$WrapFunctionAdvice");
}

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

@Advice.OnMethodEnter(suppress = Throwable.class)
public static void wrap(
@Advice.Argument(value = 1, readOnly = false) Function1<?, ?> function1) {
if (function1 == null) {
return;
}

function1 = Function1Wrapper.wrap(function1);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@

package io.opentelemetry.javaagent.instrumentation.finaglehttp.v23_11;

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

import com.google.auto.service.AutoService;
import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
Expand All @@ -21,6 +21,7 @@ public TwitterUtilCoreInstrumentationModule() {

@Override
public List<TypeInstrumentation> typeInstrumentations() {
return singletonList(new LocalSchedulerActivationInstrumentation());
return asList(
new LocalSchedulerActivationInstrumentation(), new PromiseMonitoredInstrumentation());
}
}

0 comments on commit 9028dae

Please sign in to comment.