Java Code Examples for net.bytebuddy.agent.ByteBuddyAgent

@Test
@AgentAttachmentRule.Enforce
public void InputFormatTracer_should_intercept_InputFormat_direct_implementor() throws IOException, ClassNotFoundException, IllegalAccessException, InstantiationException, NoSuchMethodException, InvocationTargetException {
    assertThat(ByteBuddyAgent.install(), instanceOf(Instrumentation.class));

    //Install tracer
    ClassFileTransformer classFileTransformer = new MapReduceTracer.InputFormatTracer().installOnByteBuddyAgent();
    try {
        //Call InputFormat
        Class<?> type = classLoader.loadClass(MapReduceInputFormatTestClasses.OneLevelHierarchy.class.getName());
        invokeRecordReader(type);
        invokeListInputSplits(type);

        //Verify mock interaction
        verify(eventHandler, times(2)).accept(any(Long.class), argument.capture());
        DataAccessEventProtos.PathEvent pathEvent = DataAccessEventProtos.PathEvent
                .newBuilder()
                .setPath(inputPath)
                .setType(PathType.INPUT.name())
                .build();
        assertEquals(pathEvent, argument.getValue());
    } finally {
        ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer);
    }
}
 

@Test
public void InputFormatTracer_should_not_intercept_abstract_method() throws ClassNotFoundException, InvocationTargetException, NoSuchMethodException, InstantiationException, IllegalAccessException {
    assertThat(ByteBuddyAgent.install(), instanceOf(Instrumentation.class));

    //Install tracer
    ClassFileTransformer classFileTransformer = new MapReduceTracer.DeprecatedInputFormatTracer().installOnByteBuddyAgent();
    try {
        //Prepare JobConf
        when(jobConf.get("mapreduce.input.fileinputformat.inputdir")).thenReturn("/some/path");

        //Call InputFormat
        Class<?> type = classLoader.loadClass(MapRedInputFormatTestClasses.RealInputFormat.class.getName());
        invokeRecordReader(type);

        //We just want to test no exception because of presence of abstract method and abstract class
    } finally {
        ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer);
    }
}
 

@Test
@AgentAttachmentRule.Enforce(retransformsClasses = true)
@JavaVersionRule.Enforce(value = 8, atMost = 8)
public void testExtractionOfInflatedMethodAccessor() throws Exception {
    assertThat(ByteBuddyAgent.install(), instanceOf(Instrumentation.class));
    Method bar = Foo.class.getDeclaredMethod("bar");
    for (int i = 0; i < 20; i++) {
        bar.invoke(new Foo());
    }
    Field field = Method.class.getDeclaredField("methodAccessor");
    field.setAccessible(true);
    Object methodAccessor = field.get(bar);
    Field delegate = methodAccessor.getClass().getDeclaredField("delegate");
    delegate.setAccessible(true);
    Class<?> delegateClass = delegate.get(methodAccessor).getClass();
    ClassFileLocator classFileLocator = ClassFileLocator.AgentBased.fromInstalledAgent(delegateClass.getClassLoader());
    ClassFileLocator.Resolution resolution = classFileLocator.locate(delegateClass.getName());
    assertThat(resolution.isResolved(), is(true));
    assertThat(resolution.resolve(), notNullValue(byte[].class));
}
 

@Test
public void InputFormatTracer_should_let_getRecordReader_return_its_original_value() throws IOException, ClassNotFoundException, IllegalAccessException, InstantiationException, NoSuchMethodException, InvocationTargetException, NoSuchFieldException {
    assertThat(ByteBuddyAgent.install(), instanceOf(Instrumentation.class));

    ClassFileTransformer classFileTransformer = new MapReduceTracer.DeprecatedInputFormatTracer().installOnByteBuddyAgent();
    try {
        //Prepare JobConf
        when(jobConf.get("mapreduce.input.fileinputformat.inputdir")).thenReturn("/some/path");

        Class<?> type = classLoader.loadClass(MapRedInputFormatTestClasses.OneLevelHierarchy.class.getName());
        Object recordReader = invokeRecordReader(type);

        assertThat(recordReader, equalTo(getRecordReaderMockReflection(type)));
    } finally {
        ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer);
    }
}
 

@Test
@AgentAttachmentRule.Enforce
public void OutputFormatTracer_should_intercept_InputFormat_direct_implementor() throws IOException, ClassNotFoundException, IllegalAccessException, InstantiationException, NoSuchMethodException, InvocationTargetException {
    assertThat(ByteBuddyAgent.install(), instanceOf(Instrumentation.class));

    //Install tracer
    ClassFileTransformer classFileTransformer = new MapReduceTracer.OutputFormatTracer().installOnByteBuddyAgent();
    try {
        //Call InputFormat
        Class<?> type = classLoader.loadClass(MapReduceOutputFormatTestClasses.OneLevelHierarchy.class.getName());
        invokeRecordWriter(type);

        //Verify mock interaction
        verify(eventHandler).accept(any(Long.class), argument.capture());
        DataAccessEventProtos.PathEvent pathEvent = DataAccessEventProtos.PathEvent
                .newBuilder()
                .setPath(outputPath)
                .setType(PathType.OUTPUT.name())
                .build();
        assertEquals(pathEvent, argument.getValue());
    } finally {
        ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer);
    }
}
 

@Test
@AgentAttachmentRule.Enforce(retransformsClasses = true)
public void testFromAgentClassReloadingStrategy() throws Exception {
    assertThat(ByteBuddyAgent.install(), instanceOf(Instrumentation.class));
    Foo foo = new Foo();
    assertThat(foo.foo(), is(FOO));
    ClassReloadingStrategy classReloadingStrategy = ClassReloadingStrategy.fromInstalledAgent();
    new ByteBuddy()
            .redefine(Foo.class)
            .method(named(FOO))
            .intercept(FixedValue.value(BAR))
            .make()
            .load(Foo.class.getClassLoader(), classReloadingStrategy);
    try {
        assertThat(foo.foo(), is(BAR));
    } finally {
        classReloadingStrategy.reset(Foo.class);
        assertThat(foo.foo(), is(FOO));
    }
}
 

public Statement apply(Statement base, FrameworkMethod method, Object target) {
    Enforce enforce = method.getAnnotation(Enforce.class);
    if (enforce != null) {
        if (!available) {
            return new NoOpStatement("The executing JVM does not support runtime attachment");
        }
        Instrumentation instrumentation = ByteBuddyAgent.install(ByteBuddyAgent.AttachmentProvider.DEFAULT);
        if (enforce.redefinesClasses() && !instrumentation.isRedefineClassesSupported()) {
            return new NoOpStatement("The executing JVM does not support class redefinition");
        } else if (enforce.retransformsClasses() && !instrumentation.isRetransformClassesSupported()) {
            return new NoOpStatement("The executing JVM does not support class retransformation");
        } else if (enforce.nativeMethodPrefix() && !instrumentation.isNativeMethodPrefixSupported()) {
            return new NoOpStatement("The executing JVM does not support class native method prefixes");
        }
    }
    return base;
}
 

/**
 * Initializes attachment provider, this method can only be called once as it loads native code.
 *
 * @param useEmulatedAttach {@literal true} to enable emulated attach, {@literal false} to disable
 */
public synchronized static void init(boolean useEmulatedAttach) {
    if (provider != null) {
        throw new IllegalStateException("ElasticAttachmentProvider.init() should only be called once");
    }

    ArrayList<ByteBuddyAgent.AttachmentProvider> providers = new ArrayList<>();
    if (useEmulatedAttach) {
        providers.add(ByteBuddyAgent.AttachmentProvider.ForEmulatedAttachment.INSTANCE);
    }
    providers.addAll(Arrays.asList(ByteBuddyAgent.AttachmentProvider.ForModularizedVm.INSTANCE,
        ByteBuddyAgent.AttachmentProvider.ForJ9Vm.INSTANCE,
        new CachedAttachmentProvider(ByteBuddyAgent.AttachmentProvider.ForStandardToolsJarVm.JVM_ROOT),
        new CachedAttachmentProvider(ByteBuddyAgent.AttachmentProvider.ForStandardToolsJarVm.JDK_ROOT),
        new CachedAttachmentProvider(ByteBuddyAgent.AttachmentProvider.ForStandardToolsJarVm.MACINTOSH),
        new CachedAttachmentProvider(ByteBuddyAgent.AttachmentProvider.ForUserDefinedToolsJar.INSTANCE)));

    provider = new ByteBuddyAgent.AttachmentProvider.Compound(providers);
}
 

@BeforeEach
void setUp(TestInfo testInfo) {
    MockTracer.MockInstrumentationSetup mockInstrumentationSetup = MockTracer.getOrCreateInstrumentationTracer();
    reporter = mockInstrumentationSetup.getReporter();
    ConfigurationRegistry config = mockInstrumentationSetup.getConfig();
    coreConfiguration = config.getConfig(CoreConfiguration.class);
    when(coreConfiguration.getTraceMethods()).thenReturn(Arrays.asList(
        MethodMatcher.of("private co.elastic.apm.agent.concurrent.AsyncTraceMethodInstrumentationTest$TestAsyncTraceMethodsClass#*"))
    );

    for (String tag : testInfo.getTags()) {
        TimeDuration duration = TimeDuration.of(tag.split("=")[1]);
        if (tag.startsWith("span_min_duration=")) {
            doReturn(duration).when(coreConfiguration).getSpanMinDuration();
        }
        if (tag.startsWith("trace_methods_duration_threshold=")) {
            doReturn(duration).when(coreConfiguration).getTraceMethodsDurationThreshold();
        }
    }

    tracer = mockInstrumentationSetup.getTracer();
    ElasticApmAgent.initInstrumentation(tracer, ByteBuddyAgent.install());
}
 

@Test
@AgentAttachmentRule.Enforce(retransformsClasses = true)
public void testClassRedefinitionRenamingWithStackMapFrames() throws Exception {
    assertThat(ByteBuddyAgent.install(), instanceOf(Instrumentation.class));
    ClassReloadingStrategy classReloadingStrategy = ClassReloadingStrategy.fromInstalledAgent();
    Bar bar = new Bar();
    new ByteBuddy().redefine(Qux.class)
            .name(Bar.class.getName())
            .make()
            .load(Bar.class.getClassLoader(), classReloadingStrategy);
    try {
        assertThat(bar.foo(), is(BAR));
    } finally {
        classReloadingStrategy.reset(Bar.class);
        assertThat(bar.foo(), is(FOO));
    }
}
 

@Test
public void testJaxRsTransactionNameFromPathAnnotationInheritanceEnabled() {
    when(config.getConfig(JaxRsConfiguration.class).isUseJaxRsPathForTransactionName()).thenReturn(true);
    when(config.getConfig(JaxRsConfiguration.class).isEnableJaxrsAnnotationInheritance()).thenReturn(true);

    ElasticApmAgent.initInstrumentation(tracer, ByteBuddyAgent.install());

    doRequest("test");
    doRequest("testAbstract");
    doRequest("testInterface");

    List<Transaction> actualTransactions = reporter.getTransactions();
    assertThat(actualTransactions).hasSize(3);
    assertThat(actualTransactions.get(0).getNameAsString()).isEqualTo("GET /test");
    assertThat(actualTransactions.get(1).getNameAsString()).isEqualTo("GET /testAbstract");
    assertThat(actualTransactions.get(2).getNameAsString()).isEqualTo("GET /testInterface");
}
 

public Statement apply(Statement base, FrameworkMethod method, Object target) {
    Enforce enforce = method.getAnnotation(Enforce.class);
    if (enforce != null) {
        if (!available) {
            return new NoOpStatement("The executing JVM does not support runtime attachment");
        }
        Instrumentation instrumentation = ByteBuddyAgent.install(ByteBuddyAgent.AttachmentProvider.DEFAULT);
        if (enforce.redefinesClasses() && !instrumentation.isRedefineClassesSupported()) {
            return new NoOpStatement("The executing JVM does not support class redefinition");
        } else if (enforce.retransformsClasses() && !instrumentation.isRetransformClassesSupported()) {
            return new NoOpStatement("The executing JVM does not support class retransformation");
        } else if (enforce.nativeMethodPrefix() && !instrumentation.isNativeMethodPrefixSupported()) {
            return new NoOpStatement("The executing JVM does not support class native method prefixes");
        }
    }
    return base;
}
 

@Test
@AgentAttachmentRule.Enforce(retransformsClasses = true)
@JavaVersionRule.Enforce(atMost = 10) // Wait for mechanism in sun.misc.Unsafe to define class.
public void testFromAgentClassWithAuxiliaryReloadingStrategy() throws Exception {
    assertThat(ByteBuddyAgent.install(), instanceOf(Instrumentation.class));
    Foo foo = new Foo();
    assertThat(foo.foo(), is(FOO));
    ClassReloadingStrategy classReloadingStrategy = ClassReloadingStrategy.fromInstalledAgent();
    String randomName = FOO + RandomString.make();
    new ByteBuddy()
            .redefine(Foo.class)
            .method(named(FOO))
            .intercept(FixedValue.value(BAR))
            .make()
            .include(new ByteBuddy().subclass(Object.class).name(randomName).make())
            .load(Foo.class.getClassLoader(), classReloadingStrategy);
    try {
        assertThat(foo.foo(), is(BAR));
    } finally {
        classReloadingStrategy.reset(Foo.class);
        assertThat(foo.foo(), is(FOO));
    }
    assertThat(Class.forName(randomName), notNullValue(Class.class));
}
 

@Test
@AgentAttachmentRule.Enforce(redefinesClasses = true)
public void testTutorialGettingStartedClassReloading() throws Exception {
    ByteBuddyAgent.install();
    FooReloading foo = new FooReloading();
    try {
        new ByteBuddy()
                .redefine(BarReloading.class)
                .name(FooReloading.class.getName())
                .make()
                .load(FooReloading.class.getClassLoader(), ClassReloadingStrategy.fromInstalledAgent());
        assertThat(foo.m(), is("bar"));
    } finally {
        ClassReloadingStrategy.fromInstalledAgent().reset(FooReloading.class); // Assure repeatability.
    }
    assertThat(foo.m(), is("foo"));
}
 

@Test
@JavaVersionRule.Enforce(value = 8, j9 = false)
@AgentAttachmentRule.Enforce
@IntegrationRule.Enforce
public void testNonCapturingLambda() throws Exception {
    assertThat(ByteBuddyAgent.install(), instanceOf(Instrumentation.class));
    ClassLoader classLoader = lambdaSamples();
    ClassFileTransformer classFileTransformer = new AgentBuilder.Default()
            .with(poolStrategy)
            .ignore(none())
            .with(AgentBuilder.LambdaInstrumentationStrategy.ENABLED)
            .type(isSubTypeOf(Callable.class)).transform(new SingleMethodReplacer("call"))
            .installOn(ByteBuddyAgent.getInstrumentation());
    try {
        Class<?> sampleFactory = classLoader.loadClass(LAMBDA_SAMPLE_FACTORY);
        @SuppressWarnings("unchecked")
        Callable<String> instance = (Callable<String>) sampleFactory.getDeclaredMethod("nonCapturing").invoke(sampleFactory.getDeclaredConstructor().newInstance());
        assertThat(instance.call(), is(BAR));
    } finally {
        assertThat(ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer), is(true));
        AgentBuilder.LambdaInstrumentationStrategy.release(classFileTransformer, ByteBuddyAgent.getInstrumentation());
    }
}
 

@Test
@JavaVersionRule.Enforce(value = 8, j9 = false)
@AgentAttachmentRule.Enforce
@IntegrationRule.Enforce
public void testNonCapturingLambdaIsConstant() throws Exception {
    assertThat(ByteBuddyAgent.install(), instanceOf(Instrumentation.class));
    ClassLoader classLoader = lambdaSamples();
    ClassFileTransformer classFileTransformer = new AgentBuilder.Default()
            .with(poolStrategy)
            .ignore(none())
            .with(AgentBuilder.LambdaInstrumentationStrategy.ENABLED)
            .type(isSubTypeOf(Callable.class)).transform(new SingleMethodReplacer("call"))
            .installOn(ByteBuddyAgent.getInstrumentation());
    try {
        Class<?> sampleFactory = classLoader.loadClass(LAMBDA_SAMPLE_FACTORY);
        assertThat(sampleFactory.getDeclaredMethod("nonCapturing").invoke(sampleFactory.getDeclaredConstructor().newInstance()),
                sameInstance(sampleFactory.getDeclaredMethod("nonCapturing").invoke(sampleFactory.getDeclaredConstructor().newInstance())));
    } finally {
        assertThat(ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer), is(true));
        AgentBuilder.LambdaInstrumentationStrategy.release(classFileTransformer, ByteBuddyAgent.getInstrumentation());
    }
}
 

@Test
@JavaVersionRule.Enforce(value = 8, j9 = false)
@AgentAttachmentRule.Enforce
@IntegrationRule.Enforce
public void testLambdaFactoryIsReset() throws Exception {
    assertThat(ByteBuddyAgent.install(), instanceOf(Instrumentation.class));
    ClassLoader classLoader = lambdaSamples();
    ClassFileTransformer classFileTransformer = new AgentBuilder.Default()
            .with(poolStrategy)
            .ignore(none())
            .with(AgentBuilder.LambdaInstrumentationStrategy.ENABLED)
            .installOn(ByteBuddyAgent.getInstrumentation());
    ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer);
    AgentBuilder.LambdaInstrumentationStrategy.release(classFileTransformer, ByteBuddyAgent.getInstrumentation());
    Class<?> sampleFactory = classLoader.loadClass(LAMBDA_SAMPLE_FACTORY);
    @SuppressWarnings("unchecked")
    Callable<String> instance = (Callable<String>) sampleFactory.getDeclaredMethod("nonCapturing").invoke(sampleFactory.getDeclaredConstructor().newInstance());
    assertThat(instance.call(), is(FOO));
}
 

@Test
@JavaVersionRule.Enforce(value = 8, j9 = false)
@AgentAttachmentRule.Enforce
@IntegrationRule.Enforce
public void testArgumentCapturingLambda() throws Exception {
    assertThat(ByteBuddyAgent.install(), instanceOf(Instrumentation.class));
    ClassLoader classLoader = lambdaSamples();
    ClassFileTransformer classFileTransformer = new AgentBuilder.Default()
            .with(poolStrategy)
            .ignore(none())
            .with(AgentBuilder.LambdaInstrumentationStrategy.ENABLED)
            .type(isSubTypeOf(Callable.class)).transform(new SingleMethodReplacer("call"))
            .installOn(ByteBuddyAgent.getInstrumentation());
    try {
        Class<?> sampleFactory = classLoader.loadClass(LAMBDA_SAMPLE_FACTORY);
        @SuppressWarnings("unchecked")
        Callable<String> instance = (Callable<String>) sampleFactory.getDeclaredMethod("argumentCapturing", String.class).invoke(sampleFactory.getDeclaredConstructor().newInstance(), FOO);
        assertThat(instance.call(), is(BAR));
    } finally {
        assertThat(ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer), is(true));
        AgentBuilder.LambdaInstrumentationStrategy.release(classFileTransformer, ByteBuddyAgent.getInstrumentation());
    }
}
 

@Test
@JavaVersionRule.Enforce(value = 8, j9 = false)
@AgentAttachmentRule.Enforce
@IntegrationRule.Enforce
public void testReturnTypeTransformingLambda() throws Exception {
    assertThat(ByteBuddyAgent.install(), instanceOf(Instrumentation.class));
    ClassLoader classLoader = lambdaSamples();
    ClassFileTransformer classFileTransformer = new AgentBuilder.Default()
            .with(poolStrategy)
            .ignore(none())
            .with(AgentBuilder.LambdaInstrumentationStrategy.ENABLED)
            .type(isSubTypeOf(Callable.class)).transform(new SingleMethodReplacer("call"))
            .installOn(ByteBuddyAgent.getInstrumentation());
    try {
        Class<?> sampleFactory = classLoader.loadClass(LAMBDA_SAMPLE_FACTORY);
        Runnable instance = (Runnable) sampleFactory.getDeclaredMethod("returnTypeTransforming").invoke(sampleFactory.getDeclaredConstructor().newInstance());
        instance.run();
    } finally {
        assertThat(ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer), is(true));
        AgentBuilder.LambdaInstrumentationStrategy.release(classFileTransformer, ByteBuddyAgent.getInstrumentation());
    }
}
 

@Test
@JavaVersionRule.Enforce(value = 8, j9 = false)
@AgentAttachmentRule.Enforce
@IntegrationRule.Enforce
public void testInstanceCapturingLambda() throws Exception {
    assertThat(ByteBuddyAgent.install(), instanceOf(Instrumentation.class));
    ClassLoader classLoader = lambdaSamples();
    ClassFileTransformer classFileTransformer = new AgentBuilder.Default()
            .with(poolStrategy)
            .ignore(none())
            .with(AgentBuilder.LambdaInstrumentationStrategy.ENABLED)
            .type(isSubTypeOf(Callable.class)).transform(new SingleMethodReplacer("call"))
            .installOn(ByteBuddyAgent.getInstrumentation());
    try {
        Class<?> sampleFactory = classLoader.loadClass(LAMBDA_SAMPLE_FACTORY);
        @SuppressWarnings("unchecked")
        Callable<String> instance = (Callable<String>) sampleFactory.getDeclaredMethod("instanceCapturing").invoke(sampleFactory.getDeclaredConstructor().newInstance());
        assertThat(instance.call(), is(BAR));
    } finally {
        assertThat(ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer), is(true));
        AgentBuilder.LambdaInstrumentationStrategy.release(classFileTransformer, ByteBuddyAgent.getInstrumentation());
    }
}
 

@Test
@JavaVersionRule.Enforce(value = 8, j9 = false)
@AgentAttachmentRule.Enforce
@IntegrationRule.Enforce
public void testNonCapturingLambdaWithArguments() throws Exception {
    assertThat(ByteBuddyAgent.install(), instanceOf(Instrumentation.class));
    ClassLoader classLoader = lambdaSamples();
    ClassFileTransformer classFileTransformer = new AgentBuilder.Default()
            .with(poolStrategy)
            .ignore(none())
            .with(AgentBuilder.LambdaInstrumentationStrategy.ENABLED)
            .type(isSubTypeOf(Class.forName("java.util.function.Function"))).transform(new SingleMethodReplacer("apply"))
            .installOn(ByteBuddyAgent.getInstrumentation());
    try {
        Class<?> sampleFactory = classLoader.loadClass(LAMBDA_SAMPLE_FACTORY);
        Object instance = sampleFactory.getDeclaredMethod("nonCapturingWithArguments").invoke(sampleFactory.getDeclaredConstructor().newInstance());
        assertThat(instance.getClass().getMethod("apply", Object.class).invoke(instance, FOO), is((Object) BAR));
    } finally {
        assertThat(ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer), is(true));
        AgentBuilder.LambdaInstrumentationStrategy.release(classFileTransformer, ByteBuddyAgent.getInstrumentation());
    }
}
 

@Test
@JavaVersionRule.Enforce(value = 8, j9 = false)
@AgentAttachmentRule.Enforce
@IntegrationRule.Enforce
public void testCapturingLambdaWithArguments() throws Exception {
    assertThat(ByteBuddyAgent.install(), instanceOf(Instrumentation.class));
    ClassLoader classLoader = lambdaSamples();
    ClassFileTransformer classFileTransformer = new AgentBuilder.Default()
            .with(poolStrategy)
            .ignore(none())
            .with(AgentBuilder.LambdaInstrumentationStrategy.ENABLED)
            .type(isSubTypeOf(Class.forName("java.util.function.Function"))).transform(new SingleMethodReplacer("apply"))
            .installOn(ByteBuddyAgent.getInstrumentation());
    try {
        Class<?> sampleFactory = classLoader.loadClass(LAMBDA_SAMPLE_FACTORY);
        Object instance = sampleFactory.getDeclaredMethod("capturingWithArguments", String.class).invoke(sampleFactory.getDeclaredConstructor().newInstance(), FOO);
        assertThat(instance.getClass().getMethod("apply", Object.class).invoke(instance, FOO), is((Object) BAR));
    } finally {
        assertThat(ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer), is(true));
        AgentBuilder.LambdaInstrumentationStrategy.release(classFileTransformer, ByteBuddyAgent.getInstrumentation());
    }
}
 

@Test
@AgentAttachmentRule.Enforce
public void testAvailable() {
    assertThat(ClassInjector.UsingInstrumentation.isAvailable(), is(true));
    assertThat(ClassInjector.UsingInstrumentation.of(folder,
                    ClassInjector.UsingInstrumentation.Target.SYSTEM,
                    ByteBuddyAgent.install()).isAlive(), is(true));
}
 

@Test
@AgentAttachmentRule.Enforce
public void ContainerResourceMonitoringModule_should_attach_to_recordCpuUsage() throws ClassNotFoundException, NoSuchMethodException, IllegalAccessException, InstantiationException, InvocationTargetException, NoSuchFieldException {
    assertThat(ByteBuddyAgent.install(), instanceOf(Instrumentation.class));

    final Header[] header = new Header[1];
    final Object[] event = new Object[1];
    TriConsumer<Long, Header, Object> cons = (t, h, o) -> {
        header[0] = h;
        event[0] = o;
    };
    ReflectionHelper.setField(null, classLoader.loadClass(ContainerResourceMonitoringTracer.class.getName()), "eventHandler", cons);

    ClassFileTransformer classFileTransformer = new ContainerResourceMonitoringTracer.VcoreUsageTracer().installOnByteBuddyAgent();

    try {
        Class<?> clazz = classLoader.loadClass(ContainerMetrics.class.getName());
        Constructor<?> constructor = clazz.getDeclaredConstructor(MetricsSystem.class, ContainerId.class, long.class, long.class);
        constructor.setAccessible(true);
        Object inFormat = constructor.newInstance(mock(MetricsSystem.class), ContainerId.fromString("container_e16940_1541383784275_24407_01_000023"), 5L, 5L);

        Method m = clazz.getDeclaredMethod("recordCpuUsage", int.class, int.class);
        m.invoke(inFormat, 10, 10000);

        assertNotNull(header[0]);
        assertNotNull(event[0]);

    } finally {
        ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer);
    }
}
 

@Test
@AgentAttachmentRule.Enforce
public void ContainerResourceMonitoringModule_should_attach_to_isProcessTreeOverLimit() throws ClassNotFoundException, NoSuchMethodException, IllegalAccessException, InstantiationException, InvocationTargetException {
    assertThat(ByteBuddyAgent.install(), instanceOf(Instrumentation.class));

    final Header[] header = new Header[1];
    final Object[] event = new Object[1];
    ContainerResourceMonitoringTracer.initEventHandler((t, h, o) -> {
        header[0] = h;
        event[0] = o;
    });
    ClassFileTransformer classFileTransformer = new ContainerResourceMonitoringTracer.MemorySizeTracer().installOnByteBuddyAgent();

    try {
        ContainerExecutor exec = mock(ContainerExecutor.class);
        AsyncDispatcher dispatcher = mock(AsyncDispatcher.class);
        Context ctx = mock(Context.class);

        Class<?> clazz = classLoader.loadClass(ContainersMonitorImpl.class.getName());
        Method m = clazz.getDeclaredMethod("isProcessTreeOverLimit", String.class, long.class, long.class, long.class);
        Object inFormat = clazz.getConstructor(ContainerExecutor.class, AsyncDispatcher.class, Context.class).newInstance(exec, dispatcher, ctx);

        m.setAccessible(true);
        m.invoke(inFormat, "container_e600_1516870220739_0069_01_000032", 2000, 1000, 3000);

        assertNotNull(header[0]);
        assertNotNull(event[0]);

    } finally {
        ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer);
    }
}
 

@Test
@AgentAttachmentRule.Enforce
public void RMAppTracer_should_not_failed_attaching_RMContextImpl_constructior_due_to_TriConsumer_visibility() throws ClassNotFoundException, NoSuchMethodException, IllegalAccessException, InvocationTargetException, URISyntaxException, IOException, InstantiationException, NoSuchFieldException, InterruptedException {
    assertThat(ByteBuddyAgent.install(), instanceOf(Instrumentation.class));

    final Header[] header = new Header[1];
    final Object[] event = new Object[1];
    TriConsumer<Long, Header, Object> cons = (t, h, o) -> {
        header[0] = h;
        event[0] = o;
    };
    ReflectionHelper.setField(null, classLoader.loadClass(RMAppTracer.class.getName()), "eventHandler", cons);

    ClassFileTransformer classFileTransformer = new RMAppTracer.RMContextImplThread().installOnByteBuddyAgent();

    try {
        Class<?> clazz = classLoader.loadClass(RMContextImpl.class.getName());
        Constructor<?> constructor = clazz.getDeclaredConstructor(Dispatcher.class, ContainerAllocationExpirer.class,
                AMLivelinessMonitor.class, AMLivelinessMonitor.class, DelegationTokenRenewer.class, AMRMTokenSecretManager.class,
                RMContainerTokenSecretManager.class, NMTokenSecretManagerInRM.class, ClientToAMTokenSecretManagerInRM.class,
                RMApplicationHistoryWriter.class);

        constructor.newInstance(mock(Dispatcher.class),
                mock(ContainerAllocationExpirer.class),
                mock(AMLivelinessMonitor.class),
                mock(AMLivelinessMonitor.class),
                mock(DelegationTokenRenewer.class),
                mock(AMRMTokenSecretManager.class),
                mock(RMContainerTokenSecretManager.class),
                mock(NMTokenSecretManagerInRM.class),
                mock(ClientToAMTokenSecretManagerInRM.class),
                mock(RMApplicationHistoryWriter.class));
    } finally {
        ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer);
    }
}
 

@Test
@AgentAttachmentRule.Enforce
public void InputFormatTracer_should_intercept_InputFormat_direct_implementor() throws IOException, ClassNotFoundException, IllegalAccessException, InstantiationException, NoSuchMethodException, InvocationTargetException {
    assertThat(ByteBuddyAgent.install(), instanceOf(Instrumentation.class));

    //Install tracer
    ClassFileTransformer classFileTransformer = new MapReduceTracer.DeprecatedInputFormatTracer().installOnByteBuddyAgent();
    try {
        //Prepare JobConf
        when(jobConf.get("mapreduce.input.fileinputformat.inputdir")).thenReturn(inputPath);

        //Call InputFormat
        Class<?> type = classLoader.loadClass(MapRedInputFormatTestClasses.OneLevelHierarchy.class.getName());
        invokeRecordReader(type);

        //Verify mock interaction
        verify(eventHandler).accept(any(Long.class), argument.capture());
        DataAccessEventProtos.PathEvent pathEvent = DataAccessEventProtos.PathEvent
                .newBuilder()
                .setPath(inputPath)
                .setType(PathType.INPUT.name())
                .build();
        assertEquals(pathEvent, argument.getValue());
    } finally {
        ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer);
    }
}
 

@Test
@AgentAttachmentRule.Enforce
public void InputFormatTracer_should_intercept_InputFormat_direct_implementor_for_deprecated_inputdir() throws IOException, ClassNotFoundException, IllegalAccessException, InstantiationException, NoSuchMethodException, InvocationTargetException {
    assertThat(ByteBuddyAgent.install(), instanceOf(Instrumentation.class));

    //Install tracer
    ClassFileTransformer classFileTransformer = new MapReduceTracer.DeprecatedInputFormatTracer().installOnByteBuddyAgent();
    try {
        //Prepare JobConf
        when(jobConf.get("mapred.input.dir")).thenReturn(deprecatedInputPath);

        //Call InputFormat
        Class<?> type = classLoader.loadClass(MapRedInputFormatTestClasses.OneLevelHierarchy.class.getName());
        invokeRecordReader(type);

        //Verify mock interaction
        verify(eventHandler).accept(any(Long.class), argument.capture());
        DataAccessEventProtos.PathEvent pathEvent = DataAccessEventProtos.PathEvent
                .newBuilder()
                .setPath(deprecatedInputPath)
                .setType(PathType.INPUT.name())
                .build();
        assertEquals(pathEvent, argument.getValue());
    } finally {
        ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer);
    }
}
 

@Test
public void InputFormatTracer_should_intercept_child_class_if_not_calling_super() throws IOException, ClassNotFoundException, InvocationTargetException, NoSuchMethodException, InstantiationException, IllegalAccessException, NoSuchFieldException {
    assertThat(ByteBuddyAgent.install(), instanceOf(Instrumentation.class));

    //Install tracer
    ClassFileTransformer classFileTransformer = new MapReduceTracer.DeprecatedInputFormatTracer().installOnByteBuddyAgent();
    try {
        //JobConf is prepared by Level3NotCallingSuper class
        //Testing PathEvent value shows Level3NotCallingSuper
        //was intercepted

        when(jobConf.get("mapreduce.input.fileinputformat.inputdir")).thenReturn("/not_calling_super");

        //Call InputFormat
        Class<?> type = classLoader.loadClass(MapRedInputFormatTestClasses.Level3NotCallingSuper.class.getName());
        Object recordReader = invokeRecordReader(type);

        //Verify mock interaction
        verify(eventHandler).accept(any(Long.class), argument.capture());
        DataAccessEventProtos.PathEvent pathEvent = DataAccessEventProtos.PathEvent
                .newBuilder()
                .setPath("/not_calling_super")
                .setType(PathType.INPUT.name())
                .build();
        assertEquals(pathEvent, argument.getValue());
        assertThat("", (boolean) type.getDeclaredField("isAccessed").get(null));
    } finally {
        ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer);
    }
}
 

@Test
@AgentAttachmentRule.Enforce
public void OutputFormatTracer_should_intercept_InputFormat_direct_implementor() throws IOException, ClassNotFoundException, IllegalAccessException, InstantiationException, NoSuchMethodException, InvocationTargetException {
    assertThat(ByteBuddyAgent.install(), instanceOf(Instrumentation.class));

    // Prepare jobConf
    when(jobConf.get("mapreduce.output.fileoutputformat.outputdir"))
            .thenReturn(outputPath);

    //Install tracer
    ClassFileTransformer classFileTransformer = new MapReduceTracer.DeprecatedOutputFormatTracer().installOnByteBuddyAgent();
    try {
        //Call OnputFormat
        Class<?> type = classLoader.loadClass(MapRedOutputFormatTestClasses.OneLevelHierarchy.class.getName());
        invokeRecordWriter(type);

        //Verify mock interaction
        verify(eventHandler).accept(any(Long.class), argument.capture());
        DataAccessEventProtos.PathEvent pathEvent = DataAccessEventProtos.PathEvent
                .newBuilder()
                .setPath(outputPath)
                .setType(PathType.OUTPUT.name())
                .build();
        assertEquals(pathEvent, argument.getValue());
    } finally {
        ByteBuddyAgent.getInstrumentation().removeTransformer(classFileTransformer);
    }
}