How to decorate an existing Java object's method?

Question

EDIT: I've described our solution at https://stackoverflow.com/a/60235242/3236516

I have a java object. It is an instance of one of many subclasses that extend an abstract class. I would like to modify one of its methods such that it runs some additional code before calling the original method. My goal is conceptually the same as a pointcut in AspectJ.

It is fine if I create some modified version of the original object rather than mutating the original. It is also fine if the solution involves bytecode manipulation.

Prior Work

I've considered creating a proxy via JavaAssist. The trouble is that ProxyFactory's create method expects that I know the constructor input types in advance. I don't. I can create my object without calling the constructor via Objenesis, but then the resulting proxy object will have null values for any values set by the constructor. This means my resulting object will behave differently from the original whenever a value set by the constructor is directly referenced.

Context

We are using Flink via AWS Kinesis Data Analytics to transform some streaming data. We would like to include some common code at the beginning of all of our StreamOperator's open() methods without having to modify each operator. One use case for this is to ensure a custom metrics agent is running on each instance an operator is running on.

Answer 1

Answer from original asker: We solved the problem by creating a ByteBuddy proxy for StreamExecutionEnvironment that intercepted calls to getStreamGraph and recast (using reflection) each node's jobVertexClass to a class that extended the original class type, but included our custom logic. Because different classes require different parameters, we instantiated the proxy without calling a constructor by using Objenesis. To solve the problem of private fields normally set in the constructor being left null, we used reflection to alter the visibility of all private fields and then copied every field value from the original object to the proxy object.

We did not pursue the agent solution proposed by Rafael Winterhalter because it requires the ability to run the agent setup code on every worker instance, which is analogous to the original problem of wanting to start a metrics agent on each worker machine. Though I did not state this in my original question, the code creating the proxy objects occurs on the Flink job management machine, not the worker machines.

Answer 2

First of all, I'd file a feature request on AWS to support your use case. That would be the cleanest solution.

Second, I would refrain from finding any way to overwrite open(). Since you are in an environment where you don't have much control, I'd imagine approaches either do not work at all or are fragile and break with an update of the environment.

I would do a lazy initialization of in the respective UDF methods and of course factor that out in some common utility method.

private Counter counter;

@Override
public Integer map(String value) {
    if (counter == null) {
        RuntimeContext ctx = getRuntimeContext();
        counter = ctx.getMetricGroup().counter("outputs");
    }
    counter.inc();
    return Integer.parseInt(value);
}

Answer 3

A Flink-specific solution might be to implement custom versions of the Flink operators you are using. I'm not convinced this would lead you to a good place; just sharing the idea in case it's helpful.

There isn't much documentation on how to implement custom operators, but there has been a Flink Forward talk on this topic.

Answer 4

With Byte Buddy, you can either create a wrapper or a Java agent which can both achieve this goal. If you struggle with constructor invocation of a wrapper class, the same problem would however occure using Byte Buddy as any library is bound to the constraints given by the JVM.

To create a Java agent, use the AgentBuilder. You can then specify all types to intercept using the type step, for example all types that implement a certain interface or extend a class. For transform, Byte Buddy offers a method decoraction API called Advice, it allows you to add additional code such as:

class MyAdvice {
  @Advice.OnMethodEnter
  static void enter() { System.out.println("Hello"); }
}

by

builder = builder.visit(Advice.to(MyAdvice.class).on(named("foo")));

you can for example print hello world at the start of all methods named "foo" for the types you specified. You can find out more about Java agents in the package documentation for the java.instrument package.