Handle JMX operations that return proxied MXBeans.

strategy/src/main/java/com/msd/gin/halyard/util/MBeanManager.java

@@ -5,19 +5,23 @@
 import java.lang.reflect.InvocationHandler;
 import java.lang.reflect.Method;
 import java.lang.reflect.Proxy;
-import java.util.ArrayList;
+import java.util.IdentityHashMap;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 
 import javax.management.JMException;
+import javax.management.JMX;
 import javax.management.MBeanRegistration;
 import javax.management.MBeanServer;
 import javax.management.ObjectInstance;
 import javax.management.ObjectName;
 
+import org.apache.commons.lang3.tuple.Pair;
+
 public abstract class MBeanManager<T> {
 	private final MBeanServer mbs;
-	private List<ObjectInstance> mbeanInsts;
+	private Map<Object,Pair<ObjectInstance,Object>> wrappers;
 
 	public static String getId(Object o) {
 		return Integer.toHexString(o.hashCode());
@@ -29,12 +33,13 @@ public MBeanManager() {
 
 	public final void register(T t) {
 		List<MBeanDetails> mbeanObjs = mbeans(t);
-		mbeanInsts = new ArrayList<>(mbeanObjs.size());
+		wrappers = new IdentityHashMap<>(mbeanObjs.size() + 1);
 		try {
 			for (MBeanDetails mbeanObj : mbeanObjs) {
 				Class<?> intf = mbeanObj.getInterface();
 				Object weakMBean = Proxy.newProxyInstance(intf.getClassLoader(), new Class[] {intf, MBeanRegistration.class}, new WeakMBean(mbeanObj.getMBean()));
-				mbeanInsts.add(mbs.registerMBean(weakMBean, mbeanObj.getName()));
+				ObjectInstance inst = mbs.registerMBean(weakMBean, mbeanObj.getName());
+				wrappers.put(mbeanObj.getMBean(), Pair.of(inst, weakMBean));
 			}
 		} catch (JMException e) {
 			throw new AssertionError(e);
@@ -44,8 +49,8 @@ public final void register(T t) {
 	protected abstract List<MBeanDetails> mbeans(T t);
 
 	public final void unregister() {
-		for (Iterator<ObjectInstance> iter = mbeanInsts.iterator(); iter.hasNext(); ) {
-			ObjectInstance inst = iter.next();
+		for (Iterator<Pair<ObjectInstance,Object>> iter = wrappers.values().iterator(); iter.hasNext(); ) {
+			ObjectInstance inst = iter.next().getKey();
 			try {
 				mbs.unregisterMBean(inst.getObjectName());
 			} catch (JMException e) {
@@ -74,7 +79,15 @@ public Object invoke(Object proxy, Method method, Object[] args) throws Throwabl
 			if (method.getDeclaringClass() == MBeanRegistration.class) {
 				return method.invoke(this, args);
 			} else {
-				return method.invoke(mbean, args);
+				Object result = method.invoke(mbean, args);
+				if (JMX.isMXBeanInterface(method.getReturnType())) {
+					// result maybe proxied
+					Pair<ObjectInstance,Object> wrapper = wrappers.get(result);
+					if (wrapper != null) {
+						return wrapper.getValue();
+					}
+				}
+				return result;
 			}
 		}
 

strategy/src/test/java/com/msd/gin/halyard/strategy/AsyncPullPusherTest.java

@@ -1,5 +1,10 @@
 package com.msd.gin.halyard.strategy;
 
+import javax.management.JMException;
+import javax.management.MBeanServer;
+import javax.management.MBeanServerFactory;
+import javax.management.ObjectName;
+
 import org.apache.hadoop.conf.Configuration;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
@@ -22,9 +27,15 @@ public void tearDown() {
 
 	@Test
 	public void testThreadPool() {
-		int tasks = pullPusher.getThreadPoolExecutor().getActiveCount();
-		assertEquals(0, tasks);
-		assertEquals(tasks, pullPusher.getThreadPoolExecutor().getThreadDump().length);
+		int taskCount = pullPusher.getThreadPoolExecutor().getActiveCount();
+		assertEquals(0, taskCount);
+		assertEquals(taskCount, pullPusher.getThreadPoolExecutor().getThreadDump().length);
 		assertEquals(0, pullPusher.getThreadPoolExecutor().getQueueDump().length);
 	}
+
+	@Test
+	public void testMXBean() throws JMException {
+		MBeanServer mbs = MBeanServerFactory.newMBeanServer();
+		mbs.registerMBean(pullPusher.getThreadPoolExecutor(), ObjectName.getInstance("foo:type=test"));
+	}
 }