Creating custom context key for use in custom sampler

[cirpene]

I have defined custom sampler using io.opentelemetry.sdk.autoconfigure.spi.traces.ConfigurableSamplerProvider SPI and it seem to be working, however context key (public static final ContextKey<Boolean> KEY = ContextKey.named("suppress-telemetry")) set on the application side (Context.current().with(Suppress.KEY, true).makeCurrent()) seem to be a different instance from one that debug of sampler reveals, hence the parentContext.get(Suppress.KEY) always returns null making it not possible to make a sampling decision based on a context key. This in turn is required to disable sampling for specific code path, as outlined in links to open-telemetry/opentelemetry-java#1717

Are there any additional undocumented steps required to make context keys available to custom samplers? Or reflecting on context key array and context key name is the only option for now, since references differ?

[laurit]

Context keys use identity comparison so you need to use the same key instance. Since agent and application code is isolated from each other and they don't share the opentelemetry api classes using the same key is not straightforward. Instead of using the key directly in you application code it may be easier to use the extension to instrument your application. Instead of setting the key in your application call a method like runWithSuppressed(Runnable runnable) in your application this method would just call runnable.run(). Instrumentation could modify this method and rewrite it to try (Scope ignore = Context.current().with(Suppress.KEY, true).makeCurrent()) { runnable.run() }

[cirpene]

Looks like instrumented code shares the application classloader, while sampler is using extension classloader, making the key set in this case no different from the plainly set one on application side, just with extra hoops.

[laurit]

Yeah, I didn't think this through completely. To make the instrumentation work you could build a custom distro. There you can place classes on boot class loader, like in https://github.com/open-telemetry/opentelemetry-java-instrumentation/blob/main/examples/distro/bootstrap/src/main/java/com/example/javaagent/bootstrap/AgentApi.java The classes in boot loader are visible to both the agent and application code and could be used as an api that allows calling from application code into the agent code.
It is probably possible to achieve something like that with an extension but you'd need to be more creative. One option would be to use byte buddy to define a class in boot loader, you could use that class with reflection from both application and agent code. Another option that could work is finding something in the jdk code that could let you share info between agent and application code. For example java.lang.System has getProperties and setProperties that use java.util.Properties which extends Hashtable<Object,Object>. Although system properties are supposed to be strings you could probably sneak in an object from some other type (I haven't tried this).

[cirpene]

Thanks, defining in bootstrap with ByteBuddy from agent extension works well.

Implementation for reference

+ java/com/example/tracing/extension/suppress/

AgentApiUtil.java

package com.example.tracing.extension.suppress;

import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.ContextKey;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Scope;

@SuppressWarnings("unchecked")
public class AgentApiUtil {
    private static final ContextKey<Boolean> KEY; // Context, Scope and ContextKey are shaded

    static {
        ContextKey<Boolean> key = null;

        try {
            Class<?> clazz = Class.forName("com.example.tracing.AgentApi");
            Field field = clazz.getDeclaredField("KEY");
            key = (ContextKey<Boolean>) field.get(null);
            if (key == null) {
                key = ContextKey.named("suppress-telemetry");
                field.set(null, key);
            }
        } catch (Throwable ignore) {
        }

        KEY = key;
    }

    public static ContextKey<Boolean> getKey() {
        return KEY;
    }

    public static Scope withSuppressed() {
        return Context.current().with(KEY, true).makeCurrent();
    }

    // Untyped Object to work-around shaded/unshaded difference in Sampler API and runtime
    public static boolean extractVlaue(Object obj) {
        if (obj instanceof Context ctx && ctx.get(KEY) instanceof Boolean b && b) {
            return b;
        }

        return false;
    }
}

Suppress.java

package com.example.tracing.extension.suppress;


import java.util.function.Supplier;

public class Suppress {
    // this will be rewritten by instrumenting agent if one is present
    public static void telemetry(Runnable run) {
        run.run();
    }

    // this will be rewritten by instrumenting agent if one is present
    public static <T> T telemetry(Supplier<T> run) {
        return run.get();
    }
}

SuppressableSampler.java

package com.example.tracing.extension.suppress;

import io.opentelemetry.api.common.Attributes;
import io.opentelemetry.api.trace.SpanKind;
import io.opentelemetry.context.Context;
import io.opentelemetry.sdk.trace.data.LinkData;
import io.opentelemetry.sdk.trace.samplers.Sampler;
import io.opentelemetry.sdk.trace.samplers.SamplingResult;

import java.util.List;

public class SuppressableSampler implements Sampler {
    private final Sampler delegate = Sampler.parentBased(Sampler.alwaysOn());

    @Override
    public SamplingResult shouldSample(
            Context parentContext,
            String traceId,
            String name,
            SpanKind spanKind,
            Attributes attributes,
            List<LinkData> parentLinks
    ) {
        if (AgentApiUtil.extractVlaue(parentContext)) {
            return SamplingResult.drop();
        }

        return delegate.shouldSample(
                parentContext,
                traceId,
                name,
                spanKind,
                attributes,
                parentLinks
        );
    }

    @Override
    public String getDescription() {
        return "suppressable";
    }
}

SuppressableSamplerProvider.java

package com.example.tracing.extension.suppress;

import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
import io.opentelemetry.sdk.autoconfigure.spi.traces.ConfigurableSamplerProvider;
import io.opentelemetry.sdk.trace.samplers.Sampler;

public class SuppressableSamplerProvider implements ConfigurableSamplerProvider {
    @Override
    public Sampler createSampler(ConfigProperties configProperties) {
        return new SuppressableSampler();
    }

    @Override
    public String getName() {
        return "suppressable";
    }
}

SuppressInstrumentation.java

package com.example.tracing.extension.suppress;

import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.ContextKey;
import net.bytebuddy.ByteBuddy;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.modifier.Ownership;
import net.bytebuddy.description.modifier.Visibility;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.dynamic.loading.ClassInjector;
import net.bytebuddy.implementation.bytecode.assign.Assigner;
import net.bytebuddy.matcher.ElementMatcher;

import java.util.Map;
import java.util.function.Supplier;

import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
import static net.bytebuddy.matcher.ElementMatchers.isPublic;
import static net.bytebuddy.matcher.ElementMatchers.isStatic;
import static net.bytebuddy.matcher.ElementMatchers.named;
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;

public class SuppressInstrumentation implements TypeInstrumentation {
    public static class Skip {

    }

    public static Skip SKIP = new Skip();

    @Override
    public ElementMatcher<ClassLoader> classLoaderOptimization() {
        return hasClassesNamed("com.example.tracing.extension.suppress.Suppress");
    }

    @Override
    public ElementMatcher<TypeDescription> typeMatcher() {
        return named("com.example.tracing.extension.suppress.Suppress");
    }

    @Override
    public void transform(TypeTransformer transformer) {
        Map<TypeDescription, byte[]> descriptor = new ByteBuddy().subclass(Object.class)
                .name("com.example.tracing.AgentApi")
                .defineField("KEY", ContextKey.class, Ownership.STATIC, Visibility.PUBLIC)
                .make()
                .getAllTypes();

        ClassInjector.UsingUnsafe.ofBootLoader().inject(descriptor);

        // init
        AgentApiUtil.getKey();

        transformer.applyAdviceToMethod(
                isMethod().and(isPublic()).and(isStatic()).and(named("telemetry")).and(takesArguments(Runnable.class)),
                this.getClass().getName() + "$SuppressRunnableAdvice"
        );

        transformer.applyAdviceToMethod(
                isMethod().and(isPublic()).and(isStatic()).and(named("telemetry")).and(takesArguments(Supplier.class)),
                this.getClass().getName() + "$SuppressSupplierAdvice"
        );
    }

    @SuppressWarnings({"unused"})
    public static class SuppressRunnableAdvice {
        @Advice.OnMethodEnter(suppress = Throwable.class, skipOn = Skip.class)
        public static Object onEnter(@Advice.Argument(0) Runnable run) {
            try (var ignore = AgentApiUtil.withSuppressed()) {
                run.run();
            }

            return SKIP;
        }
    }

    @SuppressWarnings("unused")
    public static class SuppressSupplierAdvice {
        @Advice.OnMethodEnter(suppress = Throwable.class, skipOn = Skip.class)
        public static Object onEnter(@Advice.Argument(0) Supplier<?> run, @Advice.Local("retval") Object retval) {
            try (var ignore = AgentApiUtil.withSuppressed()) {
                retval = run.get();

                return SKIP;
            }
        }

        @SuppressWarnings("UnusedAssignment")
        @Advice.OnMethodExit(suppress = Throwable.class)
        public static void onExit(
                @Advice.Return(readOnly = false, typing = Assigner.Typing.DYNAMIC) Object ret,
                @Advice.Local("retval") Object retval
        ) {
            ret = retval;
        }
    }
}

SuppressModule.java

package com.example.tracing.extension.suppress;

import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;

import java.util.List;

public class SuppressModule extends InstrumentationModule {
    public SuppressModule() {
        super("rms-suppress-module");
    }

    @Override
    public List<TypeInstrumentation> typeInstrumentations() {
        return List.of(
                new SuppressInstrumentation()
        );
    }

    @Override
    public List<String> getAdditionalHelperClassNames() {
        return List.of(
                AgentApiUtil.class.getName()
        );
    }

    @Override
    public int order() {
        return 1;
    }
}

+ resources/META-INF/services/

io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule

com.example.tracing.extension.suppress.SuppressModule

io.opentelemetry.sdk.autoconfigure.spi.traces.ConfigurableSamplerProvider

com.example.tracing.extension.suppress.SuppressableSamplerProvider

[zerocooldog]

@laurit I'm experiencing a similar issue.

Even after registering the helper class, there is an issue where the extension class loader and the application class loader are different, causing the newly added values to not be read after object creation.

It seems like I need to load my custom Java code into the bootstrap class loader. Is there another way to do this?

I found that the AgentInstaller.setBootstrapPackages method is available in the opentelemetry-java-instrumentation source code. Is this the correct way to load classes from a package into the bootstrap class loader?

Also, I noticed that the Constants.BOOTSTRAP_PACKAGE_PREFIXES class defines properties like io.opentelemetry.javaagent.bootstrap. If this is correct, could you please explain how to use it?

Alternatively, I want to call the Instrumentation.appendToBootstrapClassLoaderSearch method to directly register a JAR file with the bootstrap class loader. Is there a way to obtain the Instrumentation object when writing an extension module?

[cirpene]

Agent extension can contribute to instrumentation process via InstrumentationModule SPI, you can add a resource named io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule with the reference to your InstrumentationModule implementation within the extension jar.

In the collapsed example implementation above I add SuppressModule (implementing InstrumentationModule) that creates SuppressInstrumentation (implementing TypeInstrumentation). SuppressInstrumentation then uses bytebuddy to define a new class with a single static field, render it and then load it with bootstrap classloader just before defining the transformation set:

        Map<TypeDescription, byte[]> descriptor = new ByteBuddy().subclass(Object.class)
                .name("com.example.tracing.AgentApi")
                .defineField("KEY", ContextKey.class, Ownership.STATIC, Visibility.PUBLIC)
                .make()
                .getAllTypes();

        ClassInjector.UsingUnsafe.ofBootLoader().inject(descriptor);

Which then is loaded/used by both application (in instrumented methods of Suppress) and extension via AgentApiUtil.

Here com.example.tracing.AgentApi is a class that is built entirely by bytebuddy, but I suppose you can also use some existing class as a template to be re-loaded via bootstrap classloader, if you want to.

Also not sure if TypeInstrumentation.transform is a good place to create/inject the class, but at least it seem to be called once and early enough to preceed any sampler calls and, naturally, the instrumentation calls.

Also technically once you injected the AgentApi class into bootstrap, you don't need the instrumentation, Suppress/any part of application can reference the AgentApiUtil directly and make same try-final as the onEnter methods, but then it will become tightly coupled with the extension and will not work without one. Instrumentation on other hand allows you to place Suppress in some different [application] jar, unrelated to the extension and enjoy the Suppress.telemetry being simply a noop callback when agent (or at least it's extension) is not loaded.