diff --git a/src/main/java/org/apache/rocketmq/dashboard/aspect/admin/MQAdminAspect.java b/src/main/java/org/apache/rocketmq/dashboard/aspect/admin/MQAdminAspect.java index 4b5bd218..9e47f3e0 100644 --- a/src/main/java/org/apache/rocketmq/dashboard/aspect/admin/MQAdminAspect.java +++ b/src/main/java/org/apache/rocketmq/dashboard/aspect/admin/MQAdminAspect.java @@ -82,14 +82,22 @@ public Object aroundMQAdminMethod(ProceedingJoinPoint joinPoint) throws Throwabl try { if (isPoolConfigIsolatedByUser(rmqConfigure.isLoginRequired(), rmqConfigure.getAuthMode(), methodName)) { currentUserInfo = (UserInfo) UserInfoContext.get(WebUtil.USER_NAME); - // 2. Borrow the user-specific MQAdminExt instance. - // currentUser.getName() is assumed to be the AccessKey, and currentUser.getPassword() is SecretKey. - mqAdminExt = userMQAdminPoolManager.borrowMQAdminExt(currentUserInfo.getUsername(), currentUserInfo.getPassword()); - - // 3. Set the borrowed MQAdminExt instance into the ThreadLocal for MQAdminInstance. - // This makes it available to MQAdminExtImpl methods. - MQAdminInstance.setCurrentMQAdminExt(mqAdminExt); - log.debug("MQAdminExt borrowed for user {} and set in ThreadLocal.", currentUserInfo.getUsername()); + if (currentUserInfo == null) { + // The method is invoked outside of a request thread (e.g. a scheduled task) or + // the login context is missing. Fall back to the default pool instead of NPE. + log.warn("No user info in context for method {}, falling back to the default pool.", methodName); + mqAdminExt = mqAdminExtPool.borrowObject(); + MQAdminInstance.setCurrentMQAdminExt(mqAdminExt); + } else { + // 2. Borrow the user-specific MQAdminExt instance. + // currentUser.getName() is assumed to be the AccessKey, and currentUser.getPassword() is SecretKey. + mqAdminExt = userMQAdminPoolManager.borrowMQAdminExt(currentUserInfo.getUsername(), currentUserInfo.getPassword()); + + // 3. Set the borrowed MQAdminExt instance into the ThreadLocal for MQAdminInstance. + // This makes it available to MQAdminExtImpl methods. + MQAdminInstance.setCurrentMQAdminExt(mqAdminExt); + log.debug("MQAdminExt borrowed for user {} and set in ThreadLocal.", currentUserInfo.getUsername()); + } } else { mqAdminExt = mqAdminExtPool.borrowObject(); // Fallback to a default MQAdminExt if no user is provided MQAdminInstance.setCurrentMQAdminExt(mqAdminExt); diff --git a/src/main/java/org/apache/rocketmq/dashboard/service/impl/MessageServiceImpl.java b/src/main/java/org/apache/rocketmq/dashboard/service/impl/MessageServiceImpl.java index 69a2b2ca..49adfa30 100644 --- a/src/main/java/org/apache/rocketmq/dashboard/service/impl/MessageServiceImpl.java +++ b/src/main/java/org/apache/rocketmq/dashboard/service/impl/MessageServiceImpl.java @@ -191,6 +191,8 @@ public int compare(MessageView o1, MessageView o2) { } catch (Exception e) { Throwables.throwIfUnchecked(e); throw new RuntimeException(e); + } finally { + autoCloseConsumerWrapper.releaseConsumer(); } } @@ -395,6 +397,8 @@ private MessagePageTask queryFirstMessagePage(MessageQueryByPage query) { } catch (Exception e) { Throwables.throwIfUnchecked(e); throw new RuntimeException(e); + } finally { + autoCloseConsumerWrapper.releaseConsumer(); } } @@ -461,6 +465,8 @@ private Page queryMessageByTaskPage(MessageQueryByPage query, List< } catch (Exception e) { Throwables.throwIfUnchecked(e); throw new RuntimeException(e); + } finally { + autoCloseConsumerWrapper.releaseConsumer(); } } diff --git a/src/main/java/org/apache/rocketmq/dashboard/support/AutoCloseConsumerWrapper.java b/src/main/java/org/apache/rocketmq/dashboard/support/AutoCloseConsumerWrapper.java index 967b6f6b..ee5a1a1d 100644 --- a/src/main/java/org/apache/rocketmq/dashboard/support/AutoCloseConsumerWrapper.java +++ b/src/main/java/org/apache/rocketmq/dashboard/support/AutoCloseConsumerWrapper.java @@ -31,6 +31,7 @@ import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; @Component @@ -42,6 +43,9 @@ public class AutoCloseConsumerWrapper { private final AtomicBoolean isTaskScheduled = new AtomicBoolean(false); private final AtomicBoolean isClosing = new AtomicBoolean(false); private static volatile Instant lastUsedTime = Instant.now(); + // Number of threads currently using the consumer. The idle-close task must never + // shut the consumer down while any thread is still using it. + private final AtomicInteger inUseCount = new AtomicInteger(0); private static final ScheduledExecutorService SCHEDULER = @@ -55,25 +59,35 @@ public AutoCloseConsumerWrapper() { public DefaultMQPullConsumer getConsumer(RPCHook rpcHook, Boolean useTLS) { lastUsedTime = Instant.now(); - DefaultMQPullConsumer consumer = CONSUMER_REF.get(); - if (consumer == null) { - synchronized (this) { - consumer = CONSUMER_REF.get(); - if (consumer == null) { - consumer = createNewConsumer(rpcHook, useTLS); - CONSUMER_REF.set(consumer); - } + // The whole acquisition must happen while holding the lock: close() also runs under + // this lock, so the consumer we return cannot be shut down by the idle-close task in + // between (without the lock, close() could run right after the null-check and shut + // down a consumer that is about to be handed to the caller). + synchronized (this) { + DefaultMQPullConsumer consumer = CONSUMER_REF.get(); + if (consumer == null) { + consumer = createNewConsumer(rpcHook, useTLS); + CONSUMER_REF.set(consumer); try { consumer.start(); } catch (MQClientException e) { consumer.shutdown(); CONSUMER_REF.set(null); throw new RuntimeException("Failed to start consumer", e); - } } + inUseCount.incrementAndGet(); + return consumer; } - return consumer; + } + + /** + * Marks one usage of the consumer as finished. Must be called from a finally block after + * every successful {@link #getConsumer} call, otherwise the in-use count never drops to + * zero and the consumer will never be closed by the idle-close task again. + */ + public void releaseConsumer() { + inUseCount.decrementAndGet(); } @@ -104,9 +118,11 @@ private void startIdleCheckTask() { } public void checkAndCloseIdleConsumer() { - if (shouldClose()) { + if (inUseCount.get() == 0 && shouldClose()) { synchronized (this) { - if (shouldClose()) { + // Re-check under the lock: a thread may have acquired the consumer while we + // were waiting for the lock. + if (inUseCount.get() == 0 && shouldClose()) { close(); } }