This article will delve into the parts of the AREX agent's source code that implement full stack tracking and mock data I/O.
AREX is an open-source automation testing tool that uses Java Agent bytecode injection technology to record and store request and response data in production environments. It then replays requests and injects mock data in testing environments, storing new responses to achieve automatic recording, replay, and comparison, providing convenience for interface regression testing. When collecting data, multiple data entries (such as Request/Response, and other service invocation requests and responses) will be collected for the same request. AREX uses link tracking to link these data entries together as a complete test case.
AREX's trace tracking is similar to OpenTelemetry. Here we will briefly introduce how full trace tracking is implemented in OpenTelemetry.
The implementation of OpenTelemetry's full-link tracking
OpenTelemetry is an open-source observability tool for distributed systems, and its implementation of full-stack tracing relies on the propagation of context.
Data propagation can be divided into two categories: in-process propagation and distributed propagation. In in-process propagation, the context object passes the trace within a single service, which is relatively straightforward. However, in distributed propagation, context propagation passes context information between different services.
Context
In OpenTelemetry, context is a data structure that contains key-value pairs, such as a thread or a loop, used to pass data during request processing. OpenTelemetry provides a context object in each programming language. For example, in Java, OpenTelemetry uses ThreadLocal to store context; in Go, OpenTelemetry uses the context package; in Node.js, OpenTelemetry uses the async_hooks package; in Python, OpenTelemetry uses threading.local to store context. In C++, OpenTelemetry uses the Boost.Context library to manage context.
Context can be used to store signals such as tracing, logging, and metrics data, which can be accessed through APIs. By calling these APIs, the entire context object can be accessed, which means that tracing, logging, and metrics signals are integrated and share data throughout the context. For example, if both tracing and metrics signals are enabled, recording a metric can automatically create a tracing span. The same goes for logging: if available, logging is automatically bound to the current tracing.
Context propagation
Intra-process propagation can be either implicit or explicit, depending on the programming language being used. Implicit propagation is done automatically by storing the active context in a thread-local variable (Java, Python, Ruby, NodeJS). Explicit propagation requires passing the active context as a parameter from one function to another (Go).
For distributed propagation of traces, the tracer generates a unique transaction id for the first request and adds it to the context of the request. In subsequent requests, this context can be used to retrieve the transaction id and associate the entire request chain. For correlating database access, OpenTelemetry provides database integration libraries, such as the JDBC integration library in OpenTelemetry Java Instrumentation. This integration library automatically adds the transaction id to the database request and adds information about the database access to the span, for better understanding of the performance and behavior of the entire request chain.
The link tracking implemented by AREX
ArexThreadLocal
ArexThreadLocal is the base class for storing context in AREX, which inherits from the InheritableThreadLocal class. It is used to store tracing, logging, and metrics signals, among other data.
public class ArexThreadLocal<T> extends InheritableThreadLocal<T> {}
InheritableThreadLocal is a thread-local storage class in Java that allows child threads to inherit thread-local variables from their parent thread. Unlike regular ThreadLocal, InheritableThreadLocal allows child threads to access thread-local variables set in their parent thread.
The characteristics of InheritableThreadLocal include:
InheritableThreadLocal allows accessing thread-local variables set in the parent thread in a child thread, which is very useful in scenarios where data needs to be shared among multiple threads.
InheritableThreadLocal is thread-safe, and multiple threads can simultaneously access thread-local variables in the same InheritableThreadLocal instance without any thread-safety issues.
InheritableThreadLocal can be inherited, and child threads can inherit thread-local variables from the InheritableThreadLocal instance set in the parent thread. This feature enables child threads to inherit context information from the parent thread, making task processing more convenient.
It should be noted that the use of InheritableThreadLocal needs to be cautious, as it may lead to memory leak issues. If the objects stored in InheritableThreadLocal are not cleaned up in time, these objects will continue to exist in memory until the application exits.
Therefore, when using InheritableThreadLocal, it is important to be mindful of cleaning up the objects stored in it to avoid memory leak issues.
TraceContextManager
TraceContextManager is the management object of the tracing context in AREX, which includes a static variable TRACE_CONTEXT object (ArexThreadLocal) for storing and reading the TraceID. In addition, TraceContextManager also includes an IDGenerator for generating IDs with the prefix "AREX-".
Setting the static variable TRACE_CONTEXT through TraceContextManager can be understood as the entry point of tracing, which can set the TransactionID and perform context tracing.
ContextManager
These two functions should be noted, currentContext() is an internal dependent call, while currentContext(boolean createIfAbsent, String caseId) is a call for replaying and entry for recording. Understanding these two functions is the key to understanding the code.
public class ContextManager {
...
public static ArexContext currentContext() {
return currentContext(false, null);
}
/**
* agent will call this method
*/
public static ArexContext currentContext(boolean createIfAbsent, String caseId) {
// replay scene
if (StringUtil.isNotEmpty(caseId)) {
TraceContextManager.set(caseId);
ArexContext context = ArexContext.of(caseId, TraceContextManager.generateId());
// Each replay init generates the latest context(maybe exist previous recorded context)
RECORD_MAP.put(caseId, context);
return context;
}
...
ContextManager calls Set
When the currentContext() function in the ContextManager is called, if the passed caseID is not empty, it indicates that the current scenario is for replay, and set(caseID) is set.
When the currentContext() function is called with an empty caseID parameter, the TRACE_CONTEXT.get() method is invoked. If it cannot obtain a value, the following code is executed to generate a new transactionID using the IDGenerator and store it in the ThreadLocal variable, indicating that the current scenario is a recording (record) scenario.
messageId = idGenerator.next();
TRACE_CONTEXT.set(messageId);
Finally, the generated transactionID and its corresponding ArexContext will be stored in a ConcurrentHashMap with the transactionID as the key and the ArexContext as the value, which will be used for subsequent calls. The ContextManager manages all the ArexContexts and stores them in a Map with the caseID as the key.
public class ContextManager {
public static Map<String, ArexContext> RECORD_MAP = new ConcurrentHashMap<>();
}
...
Where ArexContext stores caseID,replayID and other information.
public class ArexContext {
private final String caseId;
private final String replayId;
private final long createTime;
private final SequenceProvider sequence;
private final List<Integer> methodSignatureHashList = new ArrayList<>();
private final Map<String, Object> cachedReplayResultMap = new ConcurrentHashMap<>();
private Map<String, Set<String>> excludeMockTemplate;
...
}
The ContextManager.currentContext() function is used to query the current context in the Agent injection script.
When the ContextManager.currentContext() function is called with parameters ContextManager.currentContext(true, id), the following actions take place:
- In the EventProcessor, the initContext function calls ContextManager.currentContext(true, id), and onCreate calls initContext.
- In the CaseEventDispatcher, onEvent(Create) calls the initContext function described above.
- In the ServletAdviceHelper, the onEvent function is called.
- In the FilterInstrumentationV3, the onEvent function is called. These classes and functions are visible in the code where injection takes place.
public class FilterInstrumentationV3 extends TypeInstrumentation {
@Override
public ElementMatcher<TypeDescription> typeMatcher() {
return not(isInterface()).and(hasSuperType(named("javax.servlet.Filter")));
}
@Override
public List<MethodInstrumentation> methodAdvices() {
ElementMatcher<MethodDescription> matcher = named("doFilter")
.and(takesArgument(0, named("javax.servlet.ServletRequest")))
.and(takesArgument(1, named("javax.servlet.ServletResponse")));
return Collections.singletonList(new MethodInstrumentation(matcher, FilterAdvice.class.getName()));
}
...
}
ServletAdviceHelper
In ServletAdviceHelper, the shouldSkip method checks whether the frequency limit (RecordLimiter.acquire) has been exceeded. If the request exceeds the limit, it will not be traced and the method will return directly. If the request does not exceed the limit, a TraceID will be generated and the onEvent method in CaseEventDispatcher will be called with a CreateEvent object (CaseEventDispatcher.onEvent(CaseEvent.ofCreateEvent)), indicating the creation of a new Trace.
CaseEventDispatcher.onEvent(CaseEvent.ofEnterEvent());
if (shouldSkip(adapter, httpServletRequest)) {
return null;
}
The actual call process
In the injection code of AREX, the functions ContextManager.needReplay() and ContextManager.needRecord() are called. In these two functions, the AREX context object is obtained by calling the currentContext() function, and based on the data in the context object, it is determined whether the current mode is replay or record. If context.isReplay(), it is replay mode, otherwise it is record mode.
The recording of the entry request is triggered when the doFilter() method of javax.servlet.Filter (as well as several other classes and methods) receives a request, and if it passes the recording frequency check, it will start recording the request.
Implementation of recording and replay in AREX
About ByteBuddy annotations
AREX injection is implemented using ByteBuddy, a powerful and easy-to-use bytecode manipulation library. ByteBuddy provides many annotations for commenting and configuring bytecode generation or modification, including:
| Annotations | Vale | Description | |
|---|---|---|---|
@OnMethodEnter | This annotation indicates that the annotated method will be called when entering the target method. The method declared with this annotation must be static. When the target method is a constructor, the @This annotation can only be written as a field and cannot be read or called as a method. | skipOn() | OnDefaultValue is an annotation that indicates if the return value of the advice method is false for boolean, 0 for byte, short, char, int, long, float, double, or null for reference types, then skip the target method. In other words, if the return value of the advice method is the default value of the primitive type or null for reference type, then the target method will be skipped. OnNonDefaultValue is the opposite of OnDefaultValue. It skips the target method if the return value of the advice method is not the default value of the primitive type or null for reference type. VoidDefault represents not skipping any method. The disadvantage of using primitive types is that it cannot retain additional information, and only relies on judging 0 or non-zero. Custom types can carry additional information. |
| prependLineNumber() | |||
| If set to true, the line number of the target method will be modified. | |||
| inline() | This annotation identifies that the method should be inlined into the target method. | ||
| suppress() | Ignore certain exceptions, handle warnings, and the default behavior is to suppress warnings. | ||
@OnMethodExit | This annotation indicates that the annotated method will be called when the target method finishes executing. The annotated method must be static. If the target method is terminated prematurely, then this method will not be called. | repeatOn() | This annotation marks whether the target method should be executed repeatedly. |
| onThrowable() | Translation: If the method being woven throws certain exceptions, they can be handled by a corresponding handler. | ||
| backupArguments() | Backup the types of all executed methods. Initially, it may impact efficiency. Backup parameters, defaulting to true, will create a dedicated copy of the original parameters. | ||
| inline() | The method should be inlined into the target method. | ||
| suppress() | Ignore some exceptions. | ||
@This | The annotated parameter should be a reference to the modified object and cannot be used on static methods and constructors. | optional() | = false is the default value. It cannot be used on constructors and static methods, otherwise, an error will be thrown: Exception in thread "main" java.lang.IllegalStateException: Cannot map this reference for static method or constructor start. If optional() = true, This can be null when encountering objects without an instance, such as constructors and static methods. |
| readOnly() | The annotated parameter is a read-only reference to the modified object. This is equivalent to declaring the parameter as final, and signals that the method cannot modify the object passed in. | ||
| typing() | Type conversion, Assigner.Typing.STATIC; does not perform forced type conversion. If the type does not match, an error will be reported directly. Assigner.Typing.DYNAMIC; performs forced type conversion. | ||
@Argument | When annotated on a parameter of the target type, it means that the annotated parameter will be retrieved using the index represented by value(). | The annotation is used to retrieve the incoming parameter. | |
| readOnly() | read only | ||
| typing() | The conversion strategy used for this type is static by default (the type will not be modified). Dynamic conversion can be achieved by using void.class, which can be changed to another class. | ||
| optional() | The alternative value to be used if the index does not exist. It is disabled by default. | ||
@AllArguments | Use an array to contain the parameters of the target method. The parameters of the target must be an array. | All input parameters have been obtained. | |
| readOnly() | read only | ||
| typing() | Type switch. | ||
@Return | Annotating a parameter to reference the return value of the target method. It can only be used with @Advice.OnMethodExit to capture the return value. It can be understood as a pointer that points to the return value. | readOnly() | read only |
typing() | Type conversion, by default, is static conversion. | ||
@Thrown | Get the thrown exception. | Catch the exception. | |
| readOnly() | read only | ||
| typing() | By default, dynamic type transformation is used, which allows changing the type. | ||
@FieldValue | 被The parameters of the annotation can refer to local variables defined within the target method, which means they can access the variables defined in the target class. | String value() | The name of the local variable. |
| declaringType() | The declared type of the local variable. | ||
| readOnly() | read only | ||
| typing() | The default behavior is static transformation. | ||
@Origin | Using a String to represent the target method, reflection is used to convert the signature of the target string into method and class formats, and then invoke it. | String value() default DEFAULT | The default value is "". |
@Enter | Annotated on a parameter, it refers to the return value of the advice method annotated with @OnMethodEnter. | Retrieve the return value. | |
| readOnly() | read only | ||
| typing() | Conversion | ||
@Exit | |||
| Annotated on a parameter, it refers to the return value of the advice method annotated with @OnMethodExit. | Retrieve the return value. | ||
| readOnly() | read only | ||
| typing() | Conversion | ||
@Local | The annotation used to declare a parameter as a local variable that is woven into the target method by Byte Buddy. The local variable can be read and written by both @OnMethodEnter and @OnMethodExit advices. However, if the local variable is referenced by an exit advice, it must also be declared in the enter advice. It is used to create a value in the local variable and create a local variable for a method. A common use case is to create a local variable within a method that can be accessed by both @Advice.OnMethodEnter and @Advice.OnMethodExit. | String value() | name |
@StubValue | mock value, which always returns a set value | It is necessary to use an Object to receive the return value, which may be a default value such as null or 0. | |
@Unused | The marked parameter is always returned with a default value, such as 0 for int or null for other types. | ||
Taking apache-httpclient-v4 as an example
public class SyncClientModuleInstrumentation extends ModuleInstrumentation
Define a class named SyncClientModuleInstrumentation that extends the ModuleInstrumentation class. Add the @AutoService(ModuleInstrumentation.class) annotation and implement the InternalHttpClientInstrumentation class, which extends the TypeInstrumentation class.
public class InternalHttpClientInstrumentation extends TypeInstrumentation
The InternalHttpClientInstrumentation class implements the typeMatcher function, which is used to match the target class for injection. In this case, "org.apache.http.impl.client.InternalHttpClient" is used as the target class name.
If you need to inject multiple classes, you can use helper functions like nameContains(), nameEndsWith(), and nameStartsWith() to match the class names:
nameContains(): Matches class names that contain a specified string.nameEndsWith(): Matches class names that end with a specified string.nameStartsWith(): Matches class names that start with a specified string.
The InternalHttpClientInstrumentation class also implements the methodAdvice function, which is used to obtain the target methods for injection.
return singletonList(new MethodInstrumentation(
isMethod().and(named("doExecute"))
.and(takesArguments(3))
.and(takesArgument(0, named("org.apache.http.HttpHost")))
.and(takesArgument(1, named("org.apache.http.HttpRequest")))
.and(takesArgument(2, named("org.apache.http.protocol.HttpContext"))),
this.getClass().getName() + "$ExecuteAdvice"));
Implement the ExecuteAdvice class, which is associated with the previously mentioned $ExecuteAdvice.
public static class ExecuteAdvice
Inject the code to retrieve the Request request from the parameters and create local variables extractor and mockResult when entering the injected function.
@Advice.OnMethodEnter(skipOn = Advice.OnNonDefaultValue.class, suppress = Throwable.class)
public static boolean onEnter(
@Advice.Argument(1) HttpRequest request,
@Advice.Local("extractor") HttpClientExtractor<HttpRequest, HttpResponse> extractor,
@Advice.Local("mockResult") MockResult mockResult) { ...
Check the Request request to see if it meets the ignore conditions. If it does, exit the method.
Check the current state in the context to determine if it is in recording or replay mode (ContextManager.needRecordOrReplay()).
If it is in replay mode, retrieve the required MOCK data for replay using mockResult = extractor.replay().
If the MOCK data retrieved from the database is of type Throwable, return success along with the Object.
If the retrieved MOCK data is of type HttpResponseWrapper, return success along with the processed response message (i.e., the MOCK data).
Inject the code to retrieve the exception and return value Request when exiting the injected function.
@Advice.OnMethodExit(onThrowable = Exception.class, suppress = Throwable.class)
public static void onExit(
@Advice.Thrown(readOnly = false) Exception throwable,
@Advice.Return(readOnly = false) CloseableHttpResponse response,
@Advice.Local("extractor") HttpClientExtractor<HttpRequest, HttpResponse> extractor,
@Advice.Local("mockResult") MockResult mockResult) {...
After processing the MockResult in the entered function, check if it is not empty. If it is not empty and of type Throwable, assign it to the throwable variable; otherwise, return the response.
If the MockResult is empty, check if it is in recording mode. If it is in recording mode, start recording the data (either throwable or response) to the database.
Take Jedis v4 as an example
The JedisModuleInstrumentation class is a subclass of ModuleInstrumentation.
The JedisFactoryInstrumentation class is a subclass of TypeInstrumentation.
The matching of class and method names follows the previous description, so it will not be repeated here.
The makeObject method is an injected function. When entering the function, it creates a class called JedisWrapper and returns it to the original class's jedisSocketFactory field.
The clientConfig function is similar to the previous steps and is also an injected function.
In the injected makeObject function, when exiting the function, the jedis is called and the result is returned.
Difficulties of AREX implementation
Multi-threaded
When AREX collects data, we need to concatenate these data to make a complete test case for the same request (Request/Response, request response of other service calls, etc.). Our applications often use asynchronous frameworks and a lot of multi-threading, which makes it very difficult to concatenate the data.
- FutureTaskInstrumentation
FutureTaskInstrumentation contains two static internal classes, $CallableAdvice and $RunnableAdvice.
public static class CallableAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static void methodEnter(
@Advice.Argument(value = 0, readOnly = false) Callable<?> callable) {
callable = CallableWrapper.get(callable);
} FutureTask }
@SuppressWarnings("unused")
public static class RunnableAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static void methodEnter(
@Advice.Argument(value = 0, readOnly = false) Runnable runnable) {
runnable = RunnableWrapper.get(runnable);
}
}...
- ThreadPoolInstrumentation
Contains two classes $ExecutorRunnableAdvice and $ExecutorCallableAdvice, implemented as above.
- ThreadInstrumentation
Contains the $StartAdvice class. It contains a static method named methodEnter that uses the Advice annotation in the Java Agent. The purpose of this method is to intercept the run method before it is executed and to perform some operations. Specifically, this method takes the argument runnable of the run method and gets it through the FieldValue annotation, then checks if the ArexContext exists, and if it does, wraps the runnable using the RunnableWrapper.get() method, and then wraps the runnable with the runnable is assigned back.
public static class StartAdvice {
@Advice.OnMethodEnter(suppress = Throwable.class)
public static void methodEnter(
@Advice.FieldValue(value = "target", readOnly = false) Runnable runnable) {
ArexContext context = ContextManager.currentContext();
if (context != null) {
runnable = RunnableWrapper.get(runnable);
}
}
}
...
Asynchronous
Many asynchronous frameworks and libraries exist in the Java ecosystem, such as Reactor and RxJava, and there are also many libraries that provide asynchronous implementations, such as lettuce which provides synchronous/asynchronous access to Redis. Since different scenarios usually require different solutions, different approaches need to be used to solve the cross-thread tracking problem in asynchronous programming.
In the case of ApacheAsyncClient, it listens for responses and initiates callbacks through a fixed thread. Therefore, multiple trace passes across threads need to be ensured throughout the calling, listening, and callback process.
To solve this problem, you can inject the execute function of the org.apache.http.impl.nio.client.InternalHttpAsyncClient class as follows and use the TraceTransmitter in the FutureCallbackWrapper to pass the Trace.
public class InternalHttpAsyncClientInstrumentation extends TypeInstrumentation {
@Override
public ElementMatcher<TypeDescription> typeMatcher() {
return named("org.apache.http.impl.nio.client.InternalHttpAsyncClient");
}
@Override
public List<MethodInstrumentation> methodAdvices() {
return singletonList(new MethodInstrumentation(
isMethod().and(named("execute"))
.and(takesArguments(4))
.and(takesArgument(0, named("org.apache.http.nio.protocol.HttpAsyncRequestProducer")))
.and(takesArgument(1, named("org.apache.http.nio.protocol.HttpAsyncResponseConsumer")))
.and(takesArgument(2, named("org.apache.http.protocol.HttpContext")))
.and(takesArgument(3, named("org.apache.http.concurrent.FutureCallback"))),
this.getClass().getName() + "$ExecuteAdvice"));
}...
Code isolation and interoperability
For system stability, the framework code of AREX agent is loaded in a separate Class loader, which does not interoperate with the application code. To ensure that the injected code can be accessed correctly at runtime, we have also made a simple modification to the ClassLoader to ensure that the runtime code will be loaded by the correct ClassLoader.
Similar to SpringBoot's LaunchedURLClassLoader, it is a class loader that is mainly responsible for loading application classes and resources and creating a URL array based on the application classpath and JAR file at application startup, and then using this URL array to initialize the ClassLoader.
When an application needs to load a class or resource, LaunchedURLClassLoader will first look in its own cache, and if it cannot find it, it will load the class or resource from a URL in the URL array.
Custom URLClassLoader and the system's own ClassLoader may conflict because the class loader in Java uses a two-parent delegation model.
In this model, each class loader has a parent class loader, and when a class needs to be loaded, it will first delegate it to its parent class loader, and only if the parent class loader cannot load it will it try to load it itself. But in this case, if both the custom URLClassLoader and the system's own ClassLoader can load the same class, then there will be two different class instances, which will cause problems for the program.
To avoid this conflict, you can override the findClass() method in the custom URLClassLoader so that it only loads its own class instead of delegating it to the parent class loader. This ensures that the custom URLClassLoader does not conflict with the system's own ClassLoader.
public class AgentClassLoader extends URLClassLoader
Classes starting with package io.arex.inst.runtime need to be loaded using a custom URLClassLoader instead of the system's own ClassLoader. this is because these classes are run as injected code, and to ensure that the code is accessed correctly, they must be loaded using the custom URLClassLoader to load them. Therefore, these classes must go through the user-state ClassLoader to be loaded and not use the system's own ClassLoader.
@Override
protected Class<?> findClass(String name) throws ClassNotFoundException {
if (StringUtil.startWithFrom(name, "runtime", 13)) {
return null;
}
JarEntryInfo jarEntryInfo = findJarEntry(name.replace('.', '/') + ".class");
if (jarEntryInfo != null && jarEntryInfo.getJarEntry() != null) {
byte[] bytes;
try {
bytes = getJarEntryBytes(jarEntryInfo);
} catch (IOException exception) {
throw new ClassNotFoundException(name, exception);
}
definePackageIfNeeded(jarEntryInfo, name);
return defineClass(name, bytes);
}
return null;
}
...
The AgentInitializer class calls initialize() in the Premain method and accesses the AgentClassLoader to load its own jar package. As follows:
public class ArexJavaAgent {
public static void premain(String agentArgs, Instrumentation inst) {
agentmain(agentArgs, inst);
}
public static void agentmain(String agentArgs, Instrumentation inst) {
init(inst, agentArgs);
}
private static void init(Instrumentation inst, String agentArgs) {
try {
installBootstrapJar(inst);
// those services must load by app class loader
//ServiceInitializer.start(agentArgs);
AgentInitializer.initialize(inst, getJarFile(ArexJavaAgent.class), agentArgs);
System.out.println("ArexJavaAgent installed.");
} catch (Exception ex) {
System.out.println("ArexJavaAgent start failed.");
ex.printStackTrace();
}
}
...