Change ContextClassLoader to NarClassLoader in ProtocolHandler (#11276)

Co-authored-by: Jiang Haiting <jianghaiting@didichuxing.com>
(cherry picked from commit 6887005)

pulsar-broker/src/main/java/org/apache/pulsar/broker/protocol/ProtocolHandlerWithClassLoader.java

@@ -43,41 +43,74 @@ class ProtocolHandlerWithClassLoader implements ProtocolHandler {
 
     @Override
     public String protocolName() {
-        return handler.protocolName();
+        try (ClassLoaderSwitcher ignored = new ClassLoaderSwitcher(classLoader)) {
+            return handler.protocolName();
+        }
     }
 
     @Override
     public boolean accept(String protocol) {
-        return handler.accept(protocol);
+        try (ClassLoaderSwitcher ignored = new ClassLoaderSwitcher(classLoader)) {
+            return handler.accept(protocol);
+        }
     }
 
     @Override
     public void initialize(ServiceConfiguration conf) throws Exception {
-        handler.initialize(conf);
+        try (ClassLoaderSwitcher ignored = new ClassLoaderSwitcher(classLoader)) {
+            handler.initialize(conf);
+        }
     }
 
     @Override
     public String getProtocolDataToAdvertise() {
-        return handler.getProtocolDataToAdvertise();
+        try (ClassLoaderSwitcher ignored = new ClassLoaderSwitcher(classLoader)) {
+            return handler.getProtocolDataToAdvertise();
+        }
     }
 
     @Override
     public void start(BrokerService service) {
-        handler.start(service);
+        try (ClassLoaderSwitcher ignored = new ClassLoaderSwitcher(classLoader)) {
+            handler.start(service);
+        }
     }
 
     @Override
     public Map<InetSocketAddress, ChannelInitializer<SocketChannel>> newChannelInitializers() {
-        return handler.newChannelInitializers();
+        try (ClassLoaderSwitcher ignored = new ClassLoaderSwitcher(classLoader)) {
+            return handler.newChannelInitializers();
+        }
     }
 
     @Override
     public void close() {
-        handler.close();
+        try (ClassLoaderSwitcher ignored = new ClassLoaderSwitcher(classLoader)) {
+            handler.close();
+        }
+
         try {
             classLoader.close();
         } catch (IOException e) {
             log.warn("Failed to close the protocol handler class loader", e);
         }
     }
+
+    /**
+     * Help to switch the class loader of current thread to the NarClassLoader, and change it back when it's done.
+     * With the help of try-with-resources statement, the code would be cleaner than using try finally every time.
+     */
+    private static class ClassLoaderSwitcher implements AutoCloseable {
+        private final ClassLoader prevClassLoader;
+
+        ClassLoaderSwitcher(ClassLoader classLoader) {
+            prevClassLoader = Thread.currentThread().getContextClassLoader();
+            Thread.currentThread().setContextClassLoader(classLoader);
+        }
+
+        @Override
+        public void close() {
+            Thread.currentThread().setContextClassLoader(prevClassLoader);
+        }
+    }
 }

pulsar-broker/src/test/java/org/apache/pulsar/broker/protocol/ProtocolHandlerWithClassLoaderTest.java

@@ -25,8 +25,14 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNull;
 import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.expectThrows;
 
+import io.netty.channel.ChannelInitializer;
+import io.netty.channel.socket.SocketChannel;
+import java.net.InetSocketAddress;
+import java.util.Map;
 import org.apache.pulsar.broker.ServiceConfiguration;
 import org.apache.pulsar.broker.service.BrokerService;
 import org.apache.pulsar.common.nar.NarClassLoader;
@@ -68,4 +74,79 @@ public void testWrapper() throws Exception {
         verify(h, times(1)).getProtocolDataToAdvertise();
     }
 
+    public void testClassLoaderSwitcher() throws Exception {
+        NarClassLoader loader = mock(NarClassLoader.class);
+
+        String protocol = "test-protocol";
+
+        ProtocolHandler h = new ProtocolHandler() {
+            @Override
+            public String protocolName() {
+                assertEquals(Thread.currentThread().getContextClassLoader(), loader);
+                return protocol;
+            }
+
+            @Override
+            public boolean accept(String protocol) {
+                assertEquals(Thread.currentThread().getContextClassLoader(), loader);
+                return true;
+            }
+
+            @Override
+            public void initialize(ServiceConfiguration conf) throws Exception {
+                assertEquals(Thread.currentThread().getContextClassLoader(), loader);
+                throw new Exception("test exception");
+            }
+
+            @Override
+            public String getProtocolDataToAdvertise() {
+                assertEquals(Thread.currentThread().getContextClassLoader(), loader);
+                return "test-protocol-data";
+            }
+
+            @Override
+            public void start(BrokerService service) {
+                assertEquals(Thread.currentThread().getContextClassLoader(), loader);
+            }
+
+            @Override
+            public Map<InetSocketAddress, ChannelInitializer<SocketChannel>> newChannelInitializers() {
+                assertEquals(Thread.currentThread().getContextClassLoader(), loader);
+                return null;
+            }
+
+            @Override
+            public void close() {
+                assertEquals(Thread.currentThread().getContextClassLoader(), loader);
+            }
+        };
+        ProtocolHandlerWithClassLoader wrapper = new ProtocolHandlerWithClassLoader(h, loader);
+
+        ClassLoader curClassLoader = Thread.currentThread().getContextClassLoader();
+
+        assertEquals(wrapper.protocolName(), protocol);
+        assertEquals(Thread.currentThread().getContextClassLoader(), curClassLoader);
+
+        assertTrue(wrapper.accept(protocol));
+        assertEquals(Thread.currentThread().getContextClassLoader(), curClassLoader);
+
+
+        ServiceConfiguration conf = new ServiceConfiguration();
+        expectThrows(Exception.class, () -> wrapper.initialize(conf));
+        assertEquals(Thread.currentThread().getContextClassLoader(), curClassLoader);
+
+        assertEquals(wrapper.getProtocolDataToAdvertise(), "test-protocol-data");
+        assertEquals(Thread.currentThread().getContextClassLoader(), curClassLoader);
+
+        BrokerService service = mock(BrokerService.class);
+        wrapper.start(service);
+        assertEquals(Thread.currentThread().getContextClassLoader(), curClassLoader);
+
+
+        assertNull(wrapper.newChannelInitializers());
+        assertEquals(Thread.currentThread().getContextClassLoader(), curClassLoader);
+
+        wrapper.close();
+        assertEquals(Thread.currentThread().getContextClassLoader(), curClassLoader);
+    }
 }