Intercepting OutOfMemoryErrors with JVM agent

I am trying to write an agent with ByteBuddy that would intercept the construction of java.lang.OutOfMemoryError and call my static method on that occasion. The target intention is to propagate information about Java errors in Prometheus metrics in a reliable way.

Example code:

public class OOMAgent {
    private static final Method HANDLE_OOM;

    static {
        try {
            HANDLE_OOM = OOMAgent.class.getDeclaredMethod("handleOOM");
        } catch (NoSuchMethodException e) {
            throw new EvitaInternalError("!!! OOMAgent initialization failed !!!", e);
        }
    }

    public static void premain(String agentArgs, Instrumentation inst) {
        if (HANDLE_OOM != null) {
            new AgentBuilder.Default()
                .type(named("java.lang.OutOfMemoryError"))
                .transform(
                    (builder, typeDescription, classLoader, javaModule, protectionDomain) -> builder
                        .constructor(any())
                        .intercept(SuperMethodCall.INSTANCE.andThen(MethodCall.invoke(HANDLE_OOM)))
                )
                .installOn(inst);
        }
    }

    public static void handleOOM() {
        System.out.println("!!! OOM !!!");
    }

}

I'm using Java, ByteBuddy 1.14.14, and have the following information in my MANIFEST.MF file

Can-Redefine-Classes: true
Can-Retransform-Classes: true
Premain-Class: io.evitadb.externalApi.observability.agent.OOMAgent

The OOMAgent class is called correctly and the agent is "installed" (no exception is thrown). But when the exception occurs, the handleOOM method is not called.

Can you spot something obviously wrong or give me some hints on how I could "debug" such an agent to find out what's wrong?

Maybe the Java classes are somehow protected? I am also using the Java 9 module system - could there be some kind of "visibility" problem?

Final solution:

I have working solution in this class.

In order to propagate information to my own classes, I was forced to create intermediary class ErrorMonitor that connects the callback from error / exception construction with Prometheus metrics and my liveness probe via. later injected consumer lambda. The class needs to be injected into the bootstrap classloader in order to be visible by the error constructors.

Because I want to intercept multiple errors by single advice and propagate the type of the exception to the metrics label, I need to use @OnMethodExit annotation where the @Advice.This annotation on input arguments is working (this is not possible for @OnMethodEnter).

The ErrorMonitor class ends up being loaded twice - once in bootstrap classloader and second in the application classloader. This means I had to inject consumers into both of them, which is visible in class ObservabilityManager.java.

Solution

Did you mean something like this?

import net.bytebuddy.agent.ByteBuddyAgent;
import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.asm.Advice.OnMethodEnter;
import net.bytebuddy.dynamic.loading.ClassInjector;

import java.lang.instrument.Instrumentation;

import static net.bytebuddy.matcher.ElementMatchers.*;

public class BootstrapOOMAgent {
  public static void main(String[] args) {
    premain("dummy", ByteBuddyAgent.install());
    throw new OutOfMemoryError("uh-oh!");
  }

  public static void premain(String arg, Instrumentation instrumentation) {
    ClassInjector.UsingUnsafe.Factory factory = ClassInjector.UsingUnsafe.Factory.resolve(instrumentation);
    AgentBuilder agentBuilder = new AgentBuilder.Default();
    agentBuilder = agentBuilder.with(new AgentBuilder.InjectionStrategy.UsingUnsafe.OfFactory(factory));

    agentBuilder
      .disableClassFormatChanges()
      .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
      .ignore(none())
      .ignore(nameStartsWith("net.bytebuddy."))
      .type(is(OutOfMemoryError.class))
      .transform((builder, typeDescription, classLoader, module, protectionDomain) -> builder
        .visit(
          Advice
            .to(MyAdvice.class)
            .on(isConstructor())
        ))
      .installOn(instrumentation);
    }

  public static class MyAdvice {
    @OnMethodEnter
    public static boolean before() {
      System.out.println("!!! OOM !!!");
      return true;
    }
  }
}

Console log:

!!! OOM !!!
Exception in thread "main" java.lang.OutOfMemoryError: uh-oh!
  at BootstrapOOMAgent.main(BootstrapOOMAgent.java:37)

Quickly adapted from my answer here.

Try it on JDoodle.