From e8e5b1f9a6625ebf0133605890705b6db3a7e276 Mon Sep 17 00:00:00 2001 From: Jason918 Date: Mon, 12 Jul 2021 22:26:34 +0800 Subject: [PATCH] Change ContextClassLoader to NarClassLoader in ProtocolHandler (#11276) Co-authored-by: Jiang Haiting --- .../ProtocolHandlerWithClassLoader.java | 47 +++++++++-- .../ProtocolHandlerWithClassLoaderTest.java | 81 +++++++++++++++++++ 2 files changed, 121 insertions(+), 7 deletions(-) diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/protocol/ProtocolHandlerWithClassLoader.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/protocol/ProtocolHandlerWithClassLoader.java index 5adb852413f7f7..223cf81ccdb560 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/protocol/ProtocolHandlerWithClassLoader.java +++ b/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> 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); + } + } } diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/protocol/ProtocolHandlerWithClassLoaderTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/protocol/ProtocolHandlerWithClassLoaderTest.java index 0704d63ea62288..42d26d53f27c6f 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/protocol/ProtocolHandlerWithClassLoaderTest.java +++ b/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> 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); + } }