How to intercept Kotlin Coroutines? - java

I'm trying to instrument Kotlin coroutines, similar to what's done here using a Javaagent. I don't want a Javaagent.
The first step is to intercept the creation, suspension and resumption of Coroutines defined in the DebugProbes. The code for that is as follows:
public class Instrumentor {
private static final Logger LOG = LoggerFactory.getLogger(Instrumentor.class);
public static void install() {
TypeDescription typeDescription = TypePool.Default.ofSystemLoader()
.describe("kotlin.coroutines.jvm.internal.DebugProbesKt")
.resolve();
new ByteBuddy()
.redefine(typeDescription, ClassFileLocator.ForClassLoader.ofSystemLoader())
.method(ElementMatchers.named("probeCoroutineCreated").and(ElementMatchers.takesArguments(1)))
.intercept(MethodDelegation.to(CoroutineCreatedAdvice.class))
.method(ElementMatchers.named("probeCoroutineResumed").and(ElementMatchers.takesArguments(1)))
.intercept(MethodDelegation.to(CoroutineResumedAdvice.class))
.method(ElementMatchers.named("probeCoroutineSuspended").and(ElementMatchers.takesArguments(1)))
.intercept(MethodDelegation.to(CoroutineSuspendedAdvice.class))
.make()
.load(ClassLoader.getSystemClassLoader(), ClassLoadingStrategy.Default.INJECTION);
DebugProbes.INSTANCE.install();
}
public static void uninstall() {
DebugProbes.INSTANCE.uninstall();
}
public static class CoroutineCreatedAdvice {
#Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
public static Continuation<Object> exit(#Advice.Return(readOnly = false) Continuation<Object> retVal) {
LOG.info("Coroutine created: {}", retVal);
return retVal;
}
}
public static class CoroutineResumedAdvice {
#Advice.OnMethodEnter(suppress = Throwable.class)
public static void enter(#Advice.Argument(0) final Continuation<Object> continuation) {
LOG.info("Coroutine resumed: {}", continuation);
}
}
public static class CoroutineSuspendedAdvice {
#Advice.OnMethodEnter(suppress = Throwable.class)
public static void enter(#Advice.Argument(0) final Continuation<Object> continuation) {
LOG.info("Coroutine suspended: {}", continuation);
}
}
}
JUnit5 test to trigger interception:
class CoroutineInstrumentationTest {
companion object {
#JvmStatic
#BeforeAll
fun beforeAll() {
Instrumentor.install()
}
#JvmStatic
#AfterAll
fun afterAll() {
Instrumentor.uninstall()
}
}
#Test
fun testInterception() {
runBlocking {
println("Test")
}
}
}
However, no interception happens (confirmed by the absence of log statements and by using a debugger). I'm new to Byte Buddy, so it's possible I'm missing something. Any ideas?
Kotlin v1.4.10, Kotlin Coroutines v1.3.9, Byte Buddy v1.10.17.

Are you sure the class is not yet loaded at this point? Try setting a breakpoint in ClassInjector.UsingReflection to see if you acutally walk through or of the injection is aborted due to a previously loaded class.
The cleaner solution would be a Java agent. You can use byte-buddy-agent to create one dynamically by ByteBuddyAgent.install() and then register an AgentBuilder on it.

Related

Kotlin Flows Java Interop Callback

I´ve been looking for a suitable solution or best practice when I want to use Kotlin Flows with ordinary callbacks. My use case is that I write a kotlin library that uses Kotlin Flow internally and i have to assume that the users will use Java for instance. So I thought that the best solution is to overload a basic callback interface to my flow method and call it in collect something like this:
class KotlinClass {
interface Callback {
fun onResult(result: Int)
}
private fun foo() = flow {
for (i in 1..3) {
emit(i)
}
}
fun bar(callback: Callback) {
runBlocking {
foo().collect { callback.onResult(it) }
}
}
private fun main() {
bar(object : Callback {
override fun onResult(result: Int) {
TODO("Not yet implemented")
}
})
}
and in my Java Application i can simply use it like that:
public class JavaClass {
public void main() {
KotlinClass libraryClass = new KotlinClass();
libraryClass.bar(new KotlinClass.Callback() {
#Override
public void onResult(int result) {
// TODO("Not yet implemented")
}
});
}
}
I am not sure whats the way to go because I would like to have my Kotlin library that uses Flows usable in a good fashion for Java and Kotlin.
I came across callbackFlow but that seems to be only if I want to let´s call it flow-ify a callback-based API? Because I am quite new to Kotlin and Flows please apologise if my question is flawed in cause of missing some basic concepts of kotlin.
I would give the Java client more control over the flow. I would add a onStart and onCompletion method to your callback interface. Beside this I would use an own CoroutineScope - maybe customizable from the Java client. And I would not block the calling thread from within the Kotlin function - no runBlocking.
#InternalCoroutinesApi
class KotlinClass {
val coroutineScope = CoroutineScope(Dispatchers.Default)
interface FlowCallback {
#JvmDefault
fun onStart() = Unit
#JvmDefault
fun onCompletion(thr: Throwable?) = Unit
fun onResult(result: Int)
}
private fun foo() = flow {
for (i in 1..3) {
emit(i)
}
}
fun bar(flowCallback: FlowCallback) {
coroutineScope.launch {
foo().onStart { flowCallback.onStart() }
.onCompletion { flowCallback.onCompletion(it) }
.collect { flowCallback.onResult(it) }
}
}
fun close() {
coroutineScope.cancel()
}
}
Now the Java client is in full control how to start, collect and cancel the flow. For example you could use a latch to wait for completion, set an timeout and cancel the couroutine scope. This looks in the first place like a lot of code, but typically you will need this kind of flexibility.
public class JavaClass {
public static void main(String[] args) throws InterruptedException {
CountDownLatch latch = new CountDownLatch(1);
KotlinClass libraryClass = new KotlinClass();
libraryClass.bar(new KotlinClass.FlowCallback() {
#Override
public void onCompletion(#Nullable Throwable thr) {
latch.countDown();
}
#Override
public void onResult(int result) {
System.out.println(result);
}
});
try {
latch.await(5, TimeUnit.SECONDS);
} finally {
libraryClass.close();
}
}
}
You don't need to create a interface in the Kotlin code. You can define bar like that:
fun bar(callback: (Int) -> Unit) {
runBlocking {
foo().collect { callback(it) }
}
}
From the Java code you can call the function like that:
public class JavaClass {
public static void main(String[] args) {
KotlinClass libraryClass = new KotlinClass();
libraryClass.bar(v -> { System.out.println(v); return Unit.INSTANCE; });
}
}
In case anyone wondering for a general solution. Here's our version of enhancement from #rene answer here.
Accept a generic type
A configurable coroutineScope
// JavaFlow.kt
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.InternalCoroutinesApi
import kotlinx.coroutines.cancel
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.collect
import kotlinx.coroutines.flow.onCompletion
import kotlinx.coroutines.flow.onStart
import kotlinx.coroutines.launch
#InternalCoroutinesApi
class JavaFlow<T>(
private val coroutineScope: CoroutineScope = CoroutineScope(Dispatchers.Default)
) {
interface OperatorCallback <T> {
#JvmDefault
fun onStart() = Unit
#JvmDefault
fun onCompletion(thr: Throwable?) = Unit
fun onResult(result: T)
}
fun collect(
flow: Flow<T>,
operatorCallback: OperatorCallback<T>,
) {
coroutineScope.launch {
flow
.onStart { operatorCallback.onStart() }
.onCompletion { operatorCallback.onCompletion(it) }
.collect { operatorCallback.onResult(it) }
}
}
fun close() {
coroutineScope.cancel()
}
}
Java caller-side:
// code omitted...
new JavaFlow<File>().collect(
// compressImageAsFlow is our actual kotlin flow extension
FileUtils.compressImageAsFlow(file, activity),
new JavaFlow.OperatorCallback<File>() {
#Override
public void onResult(File result) {
// do something with the result here
SafeSingleton.setFile(result);
}
}
);
// or using lambda with method references
// new JavaFlow<File>().collect(
// FileUtils.compressImageAsFlow(file, activity),
// SafeSingleton::setFile
// );
// Change coroutineScope to Main
// new JavaFlow<File>(CoroutineScopeKt.MainScope()).collect(
// FileUtils.compressImageAsFlow(file, activity),
// SafeSingleton::setFile
// );
OperatorCallback.onStart and OperatorCallback.onCompletion is optional, override it as needed.

Mocking builders in JUnit test

Let's say if I have a test that uses builders to construct objects. The problem is that the builder() method in the builder class is static.
Generally, mocking a static method is already an indicator of bad design. However, in the case of builders, the builder() methods are always static. What's the best approach to unit testing methods using builders()? Should the builders be refactored into a separate class to facilitate mocking?
class Service {
private SqsClient sqsClient;
private String sqsQueueUrl;
public Service(String sqsQueueUrl) {
this.sqsClient = SqsClient.builder().build();
this.sqsQueueUrl = sqsQueueUrl;
}
public SqsClient getClient() {
return this.client;
}
public SqsClient setClient(SqsClient client) {
this.client = client;
}
public String getSqsQueueUrl() {
return this.sqsQueueUrl;
}
public void setSqsQueueUrl(String sqsQueueUrl) {
this.sqsQueueUrl = sqsQueueUrl;
}
public void onEvent(Activity activity) {
// How to mock builders in unit test?
DeleteMessageRequest deleteRequest = DeleteMessageRequest.builder().queueUrl(this.sqsQueueUrl).receiptHandle(activity.getReceiptHandle()).build();
DeleteMessageResponse deleteMessageResponse = this.sqsClient.deleteMessage(deleteRequest);
}
}
class ServiceTest {
#Test
public void testEvent() {
String sqsQueueUrl = "http://127.0.0.1";
String receiptHandle = "asdasd";
SqsClient sqsClient = EasyMock.mock(SqsClient.class);
Service service = EasyMock.mock(Service.class);
// EasyMock expect and replay here.
service.setClient(sqsClient);
service.setSqsQueueUrl(sqsQueueUrl);
Activity activity1 = new Activity();
activity.setReceiptHandle(receiptHandle);
service.onEvent(activity);
}
}

How do I test Function's code when it's passed as method parameter?

Is it possible to test code that is written in lambda function that is passed inside the method process?
#AllArgsConstructor
public class JsonController {
private final JsonElementProcessingService jsonElementProcessingService;
private final JsonObjectProcessingService jsonObjectProcessingService;
private final JsonArrayProcessingService jsonArrayProcessingService;
public void process(String rawJson) {
jsonElementProcessingService.process(json -> {
JsonElement element = new JsonParser().parse(json);
if (element.isJsonArray()) {
return jsonArrayProcessingService.process(element.getAsJsonArray());
} else {
return jsonObjectProcessingService.process(element.getAsJsonObject());
}
}, rawJson);
}
}
Since the lambda is lazy the function is not invoked (Function::apply) when I call JsonController::process so is there any way to check that jsonArrayProcessingService::process is called?
#RunWith(JMockit.class)
public class JsonControllerTest {
#Injectable
private JsonElementProcessingService jsonElementProcessingService;
#Injectable
private JsonObjectProcessingService jsonObjectProcessingService;
#Injectable
private JsonArrayProcessingService jsonArrayProcessingService;
#Tested
private JsonController jsonController;
#Test
public void test() {
jsonController.process("[{\"key\":1}]");
// how check here that jsonArrayProcessingService was invoked?
}
}
Just make it testable (and readable) by converting it to a method:
public void process(String rawJson) {
jsonElementProcessingService.process(this::parse, rawJson);
}
Object parse(String json) {
JsonElement element = new JsonParser().parse(json);
if (element.isJsonArray()) {
return jsonArrayProcessingService.process(element.getAsJsonArray());
} else {
return jsonObjectProcessingService.process(element.getAsJsonObject());
}
}
The relevant guiding principles I personally follow are:
anytime my lambdas require curly brackets, convert them to a method
organise code so that it can be unit tested
You may need to change the return type of the parse method to match whatever your processing services (which you didn’t show) return.
Given its relatively-basic redirection logic, don't you just want to confirm which of the #Injectables got called:
#Test
public void test() {
jsonController.process("[{\"key\":1}]");
new Verifications() {{
jsonArrayProcessingService.process(withInstanceOf(JsonArray.class));
}};
}

Instrumentation using ByteBuddy is not working on 3-rd party classes

Instrumentation using ByteBuddy is not working when I tried to instrument 3-rd party classes
I manage to create code which instrument my own code and everything worked as expected.
When I tried to use the same code for class which are part of 3-rd party dependencies the instrumentation didn't work.
This code is working for me:
public class A {
public void print(){
System.out.println("in class A method 'print'");
}
}
public class AgentLoad {
public static void agentmain(String agentArgs, Instrumentation inst) throws Exception {
final ElementMatcher.Junction<NamedElement> matcher = ElementMatchers.named("com.instrumentation.A");
new AgentBuilder.Default()
.type(matcher)
.transform( new AgentBuilder.Transformer.ForAdvice()
.include(AgentLoad.class.getClassLoader())
.advice(named("print"), AAdvice.class.getName()))
.installOn(inst);
}
public static void main(String[] args) throws Exception{
AgentLoader.loadAgentClass(AgentLoad.class.getName(), null);
}
public static class AAdvice {
#Advice.OnMethodEnter
public static void enter() {
System.out.println("Enter Yes!!!!");
}
#Advice.OnMethodExit
public static void exit() {
System.out.println("Exist Yes!!!!");
}
}
}
This code is not working for me:
public static void agentmain(String agentArgs, Instrumentation inst) {
AgentBuilder builder = new AgentBuilder.Default();
builder.type(ElementMatchers.named("com.amazonaws.http.AmazonHttpClient"))
.transform( new AgentBuilder.Transformer.ForAdvice()
.include(AgentLoad.class.getClassLoader())
.advice(ElementMatchers.named("execute").and(isAnnotatedWith(SdkInternalApi.class)),
AmazonHttpClientAdvice.class.getName()));
builder.installOn(inst);
}
public static class AmazonHttpClientAdvice {
#Advice.OnMethodEnter
public static void executeEnter(#Advice.Argument(0) Request<?> request) {
System.out.println("Eenter !!!" + request);
}
#Advice.OnMethodExit
public static void exit(#Advice.Return(readOnly = false, typing = DYNAMIC) Object returned) {
System.out.println("Exist !!!! " + returned);
}
}
}
The expected result is a print of enter and exit when calling the 'execute' in class AmazonHttpClient.
Note: The agent is dynamically attached using the next lib: https://github.com/electronicarts/ea-agent-loader
Byte Buddy's API is immutable, all method calls are without side-effects but only return a new builder instance. This means that:
AgentBuilder builder = new AgentBuilder.Default();
builder.type(ElementMatchers.named("com.amazonaws.http.AmazonHttpClient"))
.transform(new AgentBuilder.Transformer.ForAdvice()
.include(AgentLoad.class.getClassLoader())
.advice(ElementMatchers.named("execute").and(isAnnotatedWith(SdkInternalApi.class)),
AmazonHttpClientAdvice.class.getName()));
builder.installOn(inst);
does not do anything. You need to reassign the result of the builder chain to the builder variable again or call installOn within the chain.

Debugging techniques for Byte Buddy?

I am trying to generate a very simple code with Byte Buddy.
I have a POJO class where some fields are annotated with #SecureAttribute, For such fields I would like to override getter implementation and redirect the call to a SecurityService.getSecureValue() implementation.
Original class:
public class Properties {
#SecureAttribute
protected String password;
public String getPassword() {
return password;
}
public void setPassword(String password) {
this.password = password;
}
}
Desired Proxy:
public class PropertiesProxy {
private SecurityService securityService;
public void setSecurityService(SecurityService var1) {
this.securityService = var1;
}
public SecurityService getSecurityService() {
return this.securityService;
}
#Override
public String getPassword() {
return securityService.getSecureValue(password);
}
}
Emitting a field was easy but overriding a method becomes complicated. I have found a number of samples relative to my task which I try to apply but do not seem to get the required result.
So my major question is: how do I trace and debug the code generator? First thing I've learned was to print the class to file:
DynamicType.Unloaded<?> unloadedType = byteBuddy.make();
unloadedType.saveIn(new File("d:/temp/bytebuddy"));
This gives me an output where the extra field was added but not a glance of the getter override (disassembled from .class file):
public class PropertiesImpl$ByteBuddy$OLlyZYNY extends PropertiesImpl {
private SecurityService securityService;
public void setSecurityService(SecurityService var1) {
this.securityService = var1;
}
public SecurityService getSecurityService() {
return this.securityService;
}
public PropertiesImpl$ByteBuddy$OLlyZYNY() {
}
}
Here I do not exactly understand how to look for the error. Does it mean that I used totally wrong method implementation and Byte Buddy simply skipped it? Or am I wrong with ElementMatchers? Is there some trace or whatever that will give me a clue how to fix my code?
Current implementation:
private Class<?> wrapProperties() throws IOException {
DynamicType.Builder<?> byteBuddy = new ByteBuddy()
.subclass(PropertiesImpl.class)
.defineProperty("securityService", SecurityService.class);
Arrays.stream(PropertiesImpl.class.getDeclaredFields())
.filter(item -> item.getAnnotation(SecureAttribute.class) != null)
.forEach(item -> byteBuddy
.method(ElementMatchers.named(getGetterBeanName(item)))
.intercept(new GetterWrapperImplementation(item)));
DynamicType.Unloaded<?> unloadedType = byteBuddy.make();
unloadedType.saveIn(new File("d:/temp/bytebuddy"));
Class<?> wrapperClass = unloadedType.load(PropertiesImpl.class.getClassLoader(), ClassLoadingStrategy.Default.WRAPPER)
.getLoaded();
return wrapperClass;
}
public static class GetterWrapperImplementation implements Implementation {
public static final TypeDescription SS_TYPE;
public static final MethodDescription SS_GET_SECURE_VALUE;
private final Field filed;
static {
try {
SS_TYPE = new TypeDescription.ForLoadedType(SecurityService.class);
SS_GET_SECURE_VALUE = new MethodDescription.ForLoadedMethod(SecurityService.class.getDeclaredMethod("getSecureValue", String.class));
}
catch (final NoSuchMethodException | SecurityException e) {
throw new RuntimeException(e);
}
}
public GetterWrapperImplementation(Field filed) {
this.filed = filed;
}
#Override
public InstrumentedType prepare(final InstrumentedType instrumentedType) {
return instrumentedType;
}
#Override
public ByteCodeAppender appender(final Target implementationTarget) {
final TypeDescription thisType = implementationTarget.getInstrumentedType();
return new ByteCodeAppender.Simple(Arrays.asList(
TypeCreation.of(SS_TYPE),
// get securityService field
MethodVariableAccess.loadThis(),
FieldAccess.forField(thisType.getDeclaredFields()
.filter(ElementMatchers.named("securityService"))
.getOnly()
).read(),
// get secure field
MethodVariableAccess.loadThis(),
FieldAccess.forField(thisType.getDeclaredFields()
.filter(ElementMatchers.named(filed.getName()))
.getOnly()
).read(),
MethodInvocation.invoke(SS_GET_SECURE_VALUE),
MethodReturn.of(TypeDescription.STRING)
));
}
}
What I know for the fact is that breakpoints inside ByteCodeAppender appender(final Target implementationTarget) do not get hit, but again not sure how to interpret this.
Thanks.
The Byte Buddy DSL is immutable. This means that you always have to call:
builder = builder.method(...).intercept(...);
Your forEach does not do what you expect for this reason.
As for your implementation, you can just use MethodCall on a field and define the other field as an argument.

Categories

Resources