Set classloader for JMX endpoints to application classloader

See gh-12209

spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/endpoint/jmx/EndpointMBean.java

@@ -29,8 +29,11 @@ import javax.management.MBeanException;
 import javax.management.MBeanInfo;
 import javax.management.ReflectionException;
 
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
 import reactor.core.publisher.Mono;
 
+import org.springframework.beans.factory.BeanClassLoaderAware;
 import org.springframework.boot.actuate.endpoint.InvalidEndpointRequestException;
 import org.springframework.boot.actuate.endpoint.InvocationContext;
 import org.springframework.boot.actuate.endpoint.SecurityContext;
@@ -46,11 +49,15 @@ import org.springframework.util.ClassUtils;
  * @author Phillip Webb
  * @since 2.0.0
  */
-public class EndpointMBean implements DynamicMBean {
+public class EndpointMBean implements DynamicMBean, BeanClassLoaderAware {
 
 	private static final boolean REACTOR_PRESENT = ClassUtils.isPresent(
 			"reactor.core.publisher.Mono", EndpointMBean.class.getClassLoader());
 
+	private static final Log logger = LogFactory.getLog(EndpointMBean.class);
+
+	private ClassLoader classLoader;
+
 	private final JmxOperationResponseMapper responseMapper;
 
 	private final ExposableJmxEndpoint endpoint;
@@ -69,6 +76,11 @@ public class EndpointMBean implements DynamicMBean {
 		this.operations = getOperations(endpoint);
 	}
 
+	@Override
+	public void setBeanClassLoader(ClassLoader classLoader) {
+		this.classLoader = classLoader;
+	}
+
 	private Map<String, JmxOperation> getOperations(ExposableJmxEndpoint endpoint) {
 		Map<String, JmxOperation> operations = new HashMap<>();
 		endpoint.getOperations()
@@ -90,7 +102,28 @@ public class EndpointMBean implements DynamicMBean {
 					+ "' has no operation named " + actionName;
 			throw new ReflectionException(new IllegalArgumentException(message), message);
 		}
-		return invoke(operation, params);
+		ClassLoader previousClassLoader = overrideThreadContextClassLoaderSafe(this.classLoader);
+		try {
+			return invoke(operation, params);
+		}
+		finally {
+			overrideThreadContextClassLoaderSafe(previousClassLoader);
+		}
+	}
+
+	private static ClassLoader overrideThreadContextClassLoaderSafe(ClassLoader classLoader) {
+		if (classLoader == null) {
+			return null;
+		}
+
+		try {
+			return ClassUtils.overrideThreadContextClassLoader(classLoader);
+		}
+		catch (SecurityException exc) {
+			// can't set class loader, ignore it and proceed
+			logger.warn("Unable to override class loader for JMX endpoint.");
+		}
+		return null;
 	}
 
 	private Object invoke(JmxOperation operation, Object[] params)

spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/endpoint/jmx/JmxEndpointExporter.java

@@ -29,6 +29,7 @@ import javax.management.ObjectName;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 
+import org.springframework.beans.factory.BeanClassLoaderAware;
 import org.springframework.beans.factory.DisposableBean;
 import org.springframework.beans.factory.InitializingBean;
 import org.springframework.jmx.JmxException;
@@ -42,10 +43,12 @@ import org.springframework.util.Assert;
  * @author Phillip Webb
  * @since 2.0.0
  */
-public class JmxEndpointExporter implements InitializingBean, DisposableBean {
+public class JmxEndpointExporter implements InitializingBean, DisposableBean, BeanClassLoaderAware {
 
 	private static final Log logger = LogFactory.getLog(JmxEndpointExporter.class);
 
+	private ClassLoader classLoader;
+
 	private final MBeanServer mBeanServer;
 
 	private final EndpointObjectNameFactory objectNameFactory;
@@ -70,6 +73,11 @@ public class JmxEndpointExporter implements InitializingBean, DisposableBean {
 		this.endpoints = Collections.unmodifiableCollection(endpoints);
 	}
 
+	@Override
+	public void setBeanClassLoader(ClassLoader classLoader) {
+		this.classLoader = classLoader;
+	}
+
 	@Override
 	public void afterPropertiesSet() {
 		this.registered = register();
@@ -89,6 +97,7 @@ public class JmxEndpointExporter implements InitializingBean, DisposableBean {
 		try {
 			ObjectName name = this.objectNameFactory.getObjectName(endpoint);
 			EndpointMBean mbean = new EndpointMBean(this.responseMapper, endpoint);
+			mbean.setBeanClassLoader(this.classLoader);
 			this.mBeanServer.registerMBean(mbean, name);
 			return name;
 		}

spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/endpoint/jmx/EndpointMBeanTests.java

@@ -16,6 +16,9 @@
 
 package org.springframework.boot.actuate.endpoint.jmx;
 
+import java.net.URL;
+import java.net.URLClassLoader;
+
 import javax.management.Attribute;
 import javax.management.AttributeList;
 import javax.management.AttributeNotFoundException;
@@ -32,6 +35,7 @@ import reactor.core.publisher.Mono;
 import org.springframework.beans.FatalBeanException;
 import org.springframework.boot.actuate.endpoint.InvalidEndpointRequestException;
 import org.springframework.boot.actuate.endpoint.InvocationContext;
+import org.springframework.util.ClassUtils;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.hamcrest.CoreMatchers.instanceOf;
@@ -126,6 +130,18 @@ public class EndpointMBeanTests {
 		bean.invoke("missingOperation", NO_PARAMS, NO_SIGNATURE);
 	}
 
+	@Test
+	public void invokeShouldInvokeJmxOperationWithBeanClassLoader()
+			throws ReflectionException, MBeanException {
+		TestExposableJmxEndpoint endpoint = new TestExposableJmxEndpoint(
+				new TestJmxOperation((arguments) -> ClassUtils.getDefaultClassLoader()));
+		URLClassLoader beanClassLoader = new URLClassLoader(new URL[]{}, getClass().getClassLoader());
+		EndpointMBean bean = new EndpointMBean(this.responseMapper, endpoint);
+		bean.setBeanClassLoader(beanClassLoader);
+		Object result = bean.invoke("testOperation", NO_PARAMS, NO_SIGNATURE);
+		assertThat(result).isEqualTo(beanClassLoader);
+	}
+
 	@Test
 	public void invokeWhenOperationIsInvalidShouldThrowException()
 			throws MBeanException, ReflectionException {