Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
import java.util.List;
import java.util.logging.Level;
import java.util.logging.Logger;
import javax.annotation.Nullable;

/**
* A registered callback.
Expand All @@ -29,11 +30,13 @@ public final class CallbackRegistration {
private final Runnable callback;
private final List<InstrumentDescriptor> instrumentDescriptors;
private final boolean hasStorages;
@Nullable private final ClassLoader contextClassLoader;
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think strong ref to the context class loader here is ok, given that we already have a strong ref to the callback's class loader

could make it a weakref, but not sure there's a foolproof behavior when it gets cleared


private CallbackRegistration(
List<SdkObservableMeasurement> observableMeasurements, Runnable callback) {
this.observableMeasurements = observableMeasurements;
this.callback = callback;
this.contextClassLoader = Thread.currentThread().getContextClassLoader();
this.instrumentDescriptors =
observableMeasurements.stream()
.map(SdkObservableMeasurement::getInstrumentDescriptor)
Expand Down Expand Up @@ -80,13 +83,18 @@ public void invokeCallback(RegisteredReader reader, long startEpochNanos, long e
observableMeasurements.forEach(
observableMeasurement ->
observableMeasurement.setActiveReader(reader, startEpochNanos, epochNanos));
// Restore the context class loader that was active when the callback was registered.
Thread currentThread = Thread.currentThread();
ClassLoader previousContextClassLoader = currentThread.getContextClassLoader();
currentThread.setContextClassLoader(contextClassLoader);
try {
callback.run();
} catch (Throwable e) {
propagateIfFatal(e);
throttlingLogger.log(
Level.WARNING, "An exception occurred invoking callback for " + this + ".", e);
} finally {
currentThread.setContextClassLoader(previousContextClassLoader);
observableMeasurements.forEach(SdkObservableMeasurement::unsetActiveReader);
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
import java.util.Arrays;
import java.util.Collections;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
Expand Down Expand Up @@ -219,6 +220,59 @@ void invokeCallback_NoStorage() {
assertThat(counter.get()).isEqualTo(0);
}

@Test
void invokeCallback_RestoresContextClassLoader() {
// Simulate the context class loader at registration time
ClassLoader registrationClassLoader = new ClassLoader() {};
ClassLoader originalClassLoader = Thread.currentThread().getContextClassLoader();

Thread.currentThread().setContextClassLoader(registrationClassLoader);
AtomicReference<ClassLoader> observedClassLoader = new AtomicReference<>();
Runnable callback =
() -> observedClassLoader.set(Thread.currentThread().getContextClassLoader());
CallbackRegistration callbackRegistration =
CallbackRegistration.create(Collections.singletonList(measurement2), callback);
Thread.currentThread().setContextClassLoader(originalClassLoader);

// Simulate invocation on a thread with null context class loader (like DaemonThreadFactory)
Thread.currentThread().setContextClassLoader(null);
callbackRegistration.invokeCallback(registeredReader, 0, 1);

// Callback should have seen the registration-time classloader
assertThat(observedClassLoader.get()).isSameAs(registrationClassLoader);

// After invocation, the thread's context classloader should be restored to null
assertThat(Thread.currentThread().getContextClassLoader()).isNull();

// Clean up
Thread.currentThread().setContextClassLoader(originalClassLoader);
}

@Test
void invokeCallback_RestoresContextClassLoaderOnException() {
ClassLoader registrationClassLoader = new ClassLoader() {};
ClassLoader originalClassLoader = Thread.currentThread().getContextClassLoader();

Thread.currentThread().setContextClassLoader(registrationClassLoader);
Runnable callback =
() -> {
throw new RuntimeException("Error!");
};
CallbackRegistration callbackRegistration =
CallbackRegistration.create(Collections.singletonList(measurement2), callback);
Thread.currentThread().setContextClassLoader(originalClassLoader);

// Simulate invocation on a thread with null context class loader
Thread.currentThread().setContextClassLoader(null);
callbackRegistration.invokeCallback(registeredReader, 0, 1);

// Context classloader should still be restored even after exception
assertThat(Thread.currentThread().getContextClassLoader()).isNull();

// Clean up
Thread.currentThread().setContextClassLoader(originalClassLoader);
}

@Test
void invokeCallback_MultipleMeasurements_ThrowsException() {
Runnable callback =
Expand Down
Loading