[fix] Fix flaky test MultiBundlesTest

.github/workflows/pr-test.yml

@@ -2,13 +2,17 @@ name: aop mvn build check and ut
 
 on:
   pull_request:
+    # PR target branches (e.g. master, branch-4.0)
     branches:
       - master
       - branch-*
   push:
+    # Only maintenance / main lines. Exclude PR head branches like branch-4.0.10.2
+    # so "branch-4.0.10.2 -> branch-4.0" PRs are not double-triggered by push.
     branches:
       - master
       - branch-*
+      - '!branch-*.*.*'
 
 jobs:
   build:

amqp-impl/src/main/java/io/streamnative/pulsar/handlers/amqp/AmqpBrokerService.java

@@ -16,15 +16,19 @@
 
 import io.netty.util.concurrent.DefaultThreadFactory;
 import io.streamnative.pulsar.handlers.amqp.admin.AmqpAdmin;
+import java.io.IOException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import lombok.Getter;
+import lombok.extern.slf4j.Slf4j;
 import org.apache.pulsar.broker.PulsarService;
 import org.apache.pulsar.broker.authentication.AuthenticationService;
+import org.apache.pulsar.broker.resources.MetadataStoreCacheLoader;
 
 /**
  * AMQP broker related.
  */
+@Slf4j
 public class AmqpBrokerService {
     @Getter
     private AmqpTopicManager amqpTopicManager;
@@ -42,6 +46,8 @@ public class AmqpBrokerService {
     private PulsarService pulsarService;
     @Getter
     private AmqpAdmin amqpAdmin;
+    @Getter
+    private MetadataStoreCacheLoader metadataStoreCacheLoader;
 
     public AmqpBrokerService(PulsarService pulsarService, AmqpServiceConfiguration config) {
         this.pulsarService = pulsarService;
@@ -53,6 +59,15 @@ public AmqpBrokerService(PulsarService pulsarService, AmqpServiceConfiguration c
         this.queueService = new QueueServiceImpl(exchangeContainer, queueContainer);
         this.connectionContainer = new ConnectionContainer(pulsarService, exchangeContainer, queueContainer);
         this.amqpAdmin = new AmqpAdmin("localhost", config.getAmqpAdminPort());
+        try {
+            // Used by admin ownership redirects to resolve the owner broker's amqpAdminPort.
+            this.metadataStoreCacheLoader = new MetadataStoreCacheLoader(pulsarService.getPulsarResources(),
+                    30_000);
+        } catch (Exception e) {
+            // Unit tests may mock PulsarResources without load-report store; keep service usable.
+            log.warn("Failed to init MetadataStoreCacheLoader for AoP, admin redirects may fail", e);
+            this.metadataStoreCacheLoader = null;
+        }
     }
 
     private ExecutorService initRouteExecutor(AmqpServiceConfiguration config) {
@@ -67,4 +82,14 @@ public boolean isAuthenticationEnabled() {
     public AuthenticationService getAuthenticationService() {
         return pulsarService.getBrokerService().getAuthenticationService();
     }
+
+    public void close() {
+        if (metadataStoreCacheLoader != null) {
+            try {
+                metadataStoreCacheLoader.close();
+            } catch (IOException e) {
+                log.warn("Failed to close MetadataStoreCacheLoader", e);
+            }
+        }
+    }
 }

amqp-impl/src/main/java/io/streamnative/pulsar/handlers/amqp/AmqpChannel.java

@@ -124,8 +124,9 @@ public class AmqpChannel implements ServerChannelMethodProcessor {
     /**
      * The delivery tag is unique per channel. This is pre-incremented before putting into the deliver frame so that
      * value of this represents the <b>last</b> tag sent out.
+     * AtomicLong is required because multi-bundle mode may deliver from multiple consumers concurrently.
      */
-    protected volatile long deliveryTag = 0;
+    protected final AtomicLong deliveryTag = new AtomicLong(0);
     protected final AmqpFlowCreditManager creditManager;
     protected final AtomicBoolean blockedOnCredit = new AtomicBoolean(false);
     public static final int DEFAULT_CONSUMER_PERMIT = 1000;
@@ -905,7 +906,7 @@ public void closeChannel(int cause, final String message) {
     }
 
     public long getNextDeliveryTag() {
-        return ++deliveryTag;
+        return deliveryTag.incrementAndGet();
     }
 
     public AmqpConnection getConnection() {

amqp-impl/src/main/java/io/streamnative/pulsar/handlers/amqp/AmqpProtocolHandler.java

@@ -24,6 +24,7 @@
 import io.streamnative.pulsar.handlers.amqp.utils.ConfigurationUtils;
 import java.net.InetSocketAddress;
 import java.util.Map;
+import java.util.Optional;
 import lombok.Getter;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.pulsar.broker.ServiceConfiguration;
@@ -46,6 +47,7 @@ public class AmqpProtocolHandler implements ProtocolHandler {
     public static final String SSL_PREFIX = "amqp+ssl://";
     public static final String PLAINTEXT_PREFIX = "amqp://";
     public static final String LISTENER_DEL = ",";
+    public static final String PROTOCOL_DATA_SEP = "|";
     public static final String LISTENER_PATTEN = "^(amqp)://[-a-zA-Z0-9+&@#/%?=~_|!:,.;]*[-0-9+]";
 
     @Getter
@@ -84,10 +86,13 @@ public void initialize(ServiceConfiguration conf) throws Exception {
     // This method is called after initialize
     @Override
     public String getProtocolDataToAdvertise() {
+        // Format: <amqp-listeners>|<admin-port>
+        // Admin port is required for multi-broker ownership redirects of AoP admin REST calls.
+        String protocolData = getAppliedAmqpListeners(amqpConfig) + PROTOCOL_DATA_SEP + amqpConfig.getAmqpAdminPort();
         if (log.isDebugEnabled()) {
-            log.debug("Get configured listeners: {}", getAppliedAmqpListeners(amqpConfig));
+            log.debug("Get protocol data to advertise: {}", protocolData);
         }
-        return getAppliedAmqpListeners(amqpConfig);
+        return protocolData;
     }
 
     @Override
@@ -182,10 +187,15 @@ public Map<InetSocketAddress, ChannelInitializer<SocketChannel>> newChannelIniti
     @Override
     public void close() {
         try {
-            webServer.stop();
+            if (webServer != null) {
+                webServer.stop();
+            }
         } catch (Exception e) {
             log.error("Failed to stop web server for aop", e);
         }
+        if (amqpBrokerService != null) {
+            amqpBrokerService.close();
+        }
     }
 
     public static int getListenerPort(String listener) {
@@ -207,4 +217,45 @@ public static String getAppliedAmqpListeners(AmqpServiceConfiguration configurat
     public static String amqpUrl(String host, int port) {
         return String.format("amqp://%s:%d", host, port);
     }
+
+    /**
+     * Extract AMQP listeners from protocol advertise data.
+     * Compatible with both legacy format (`amqp://host:port`) and
+     * new format (`amqp://host:port|adminPort`).
+     */
+    public static String extractAmqpListeners(String protocolData) {
+        if (protocolData == null) {
+            return null;
+        }
+        int sep = protocolData.lastIndexOf(PROTOCOL_DATA_SEP);
+        if (sep < 0) {
+            return protocolData;
+        }
+        String maybeAdminPort = protocolData.substring(sep + 1);
+        try {
+            Integer.parseInt(maybeAdminPort);
+            return protocolData.substring(0, sep);
+        } catch (NumberFormatException e) {
+            return protocolData;
+        }
+    }
+
+    /**
+     * Extract AoP admin port from protocol advertise data.
+     * Returns empty if the data uses the legacy format without admin port.
+     */
+    public static Optional<Integer> extractAmqpAdminPort(String protocolData) {
+        if (protocolData == null) {
+            return Optional.empty();
+        }
+        int sep = protocolData.lastIndexOf(PROTOCOL_DATA_SEP);
+        if (sep < 0) {
+            return Optional.empty();
+        }
+        try {
+            return Optional.of(Integer.parseInt(protocolData.substring(sep + 1)));
+        } catch (NumberFormatException e) {
+            return Optional.empty();
+        }
+    }
 }

amqp-impl/src/main/java/io/streamnative/pulsar/handlers/amqp/AmqpPulsarConsumer.java

@@ -75,6 +75,12 @@ private void consume() {
 
             MessageIdImpl messageId = (MessageIdImpl) message.getMessageId();
             long deliveryIndex = this.amqpChannel.getNextDeliveryTag();
+            // Register unacked before deliver so a fast client ack cannot miss the tag.
+            if (!this.autoAck) {
+                this.amqpChannel.getUnacknowledgedMessageMap().add(
+                        deliveryIndex, PositionFactory.create(messageId.getLedgerId(), messageId.getEntryId()),
+                        AmqpPulsarConsumer.this, message.size());
+            }
             this.amqpChannel.getConnection().getAmqpOutputConverter().writeDeliver(
                     MessageConvertUtils.messageToAmqpBody(message),
                     this.amqpChannel.getChannelId(),
@@ -87,10 +93,6 @@ private void consume() {
                             messageId, consumer.getTopic(), t);
                     return null;
                 });
-            } else {
-                this.amqpChannel.getUnacknowledgedMessageMap().add(
-                        deliveryIndex, PositionFactory.create(messageId.getLedgerId(), messageId.getEntryId()),
-                        AmqpPulsarConsumer.this, message.size());
             }
             consumeBackoff.reset();
             this.consume();

amqp-impl/src/main/java/io/streamnative/pulsar/handlers/amqp/MessageFetchContext.java

@@ -71,13 +71,15 @@ public static void handleFetch(AmqpChannel channel, AmqpConsumer consumer, boole
                 long deliveryTag = channel.getNextDeliveryTag();
                 boolean isRedelivery = consumer.getRedeliveryTracker()
                         .getRedeliveryCount(pair.getLeft().getLedgerId(), pair.getLeft().getEntryId()) > 0;
+                // Register unacked before get-ok so a fast client ack cannot miss the tag.
+                if (!autoAck) {
+                    channel.getUnacknowledgedMessageMap().add(deliveryTag, pair.getLeft(), consumer, 0);
+                    channel.getCreditManager().useCreditForMessages(1, 0);
+                }
                 channel.getConnection().getAmqpOutputConverter().writeGetOk(pair.getRight(), channel.getChannelId(),
                         isRedelivery, deliveryTag, 0);
                 if (autoAck) {
                     consumer.messageAck(pair.getLeft());
-                } else {
-                    channel.getUnacknowledgedMessageMap().add(deliveryTag, pair.getLeft(), consumer, 0);
-                    channel.getCreditManager().useCreditForMessages(1, 0);
                 }
             } else {
                 if (pair != null && pair.getLeft() != null) {

amqp-impl/src/main/java/io/streamnative/pulsar/handlers/amqp/admin/impl/BaseResources.java

@@ -30,19 +30,25 @@
 import java.net.URI;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.tuple.Pair;
 import org.apache.pulsar.broker.PulsarService;
+import org.apache.pulsar.broker.lookup.LookupResult;
 import org.apache.pulsar.broker.namespace.LookupOptions;
 import org.apache.pulsar.broker.namespace.NamespaceService;
+import org.apache.pulsar.broker.resources.MetadataStoreCacheLoader;
 import org.apache.pulsar.broker.resources.NamespaceResources;
 import org.apache.pulsar.broker.service.BrokerServiceException;
 import org.apache.pulsar.broker.web.RestException;
 import org.apache.pulsar.client.admin.PulsarAdminException;
+import org.apache.pulsar.common.lookup.data.LookupData;
 import org.apache.pulsar.common.naming.TopicName;
 import org.apache.pulsar.common.util.FutureUtil;
 import org.apache.pulsar.metadata.api.MetadataStoreException;
+import org.apache.pulsar.policies.data.loadbalancer.LoadManagerReport;
 
 /**
  * Base resources.
@@ -161,18 +167,23 @@ protected CompletableFuture<Void> validateTopicOwnershipAsync(TopicName topicNam
                                 "Failed to find ownership for topic:" + topicName);
                     }
                     return lookupResult.get();
-                }).thenCompose(webUrl -> nsService.isServiceUnitOwnedAsync(topicName)
-                        .thenApply(isTopicOwned -> Pair.of(webUrl, isTopicOwned))
+                }).thenCompose(lookupResult -> nsService.isServiceUnitOwnedAsync(topicName)
+                        .thenApply(isTopicOwned -> Pair.of(lookupResult, isTopicOwned))
                 ).thenAccept(pair -> {
-                    URI webUri = pair.getLeft().toLookupRedirectUri(uri.getRequestUri());
+                    LookupResult lookupResult = pair.getLeft();
+                    URI webUri = lookupResult.toLookupRedirectUri(uri.getRequestUri());
                     boolean isTopicOwned = pair.getRight();
 
                     if (!isTopicOwned) {
                         boolean newAuthoritative = isLeaderBroker(pulsar());
-                        // Replace the host and port of the current request and redirect
+                        int adminPort = resolveOwnerAmqpAdminPort(lookupResult)
+                                .orElseThrow(() -> new RestException(Response.Status.PRECONDITION_FAILED,
+                                        "Failed to resolve amqp admin port for topic:" + topicName));
+                        // Redirect to the owner broker's AoP admin endpoint.
+                        // Host comes from lookup; port must be the owner admin port (not local).
                         URI redirect = UriBuilder.fromUri(uri.getRequestUri())
                                 .host(webUri.getHost())
-                                .port(aop().getAmqpConfig().getAmqpAdminPort())
+                                .port(adminPort)
                                 .replaceQueryParam("authoritative", newAuthoritative)
                                 .build();
                         // Redirect
@@ -197,6 +208,41 @@ protected CompletableFuture<Void> validateTopicOwnershipAsync(TopicName topicNam
                 });
     }
 
+    private Optional<Integer> resolveOwnerAmqpAdminPort(LookupResult lookupResult) {
+        LookupData lookupData = lookupResult.getLookupData();
+        MetadataStoreCacheLoader cacheLoader = aop().getAmqpBrokerService().getMetadataStoreCacheLoader();
+        if (cacheLoader == null) {
+            log.warn("MetadataStoreCacheLoader is unavailable, cannot resolve owner amqp admin port");
+            return Optional.empty();
+        }
+        List<LoadManagerReport> brokers = cacheLoader.getAvailableBrokers();
+        Optional<LoadManagerReport> owner = brokers.stream()
+                .filter(report -> matchesOwnerBroker(report, lookupData))
+                .findFirst();
+        if (owner.isEmpty()) {
+            log.warn("Unable to locate load report for owner broker. httpUrl={}, brokerUrl={}, available={}",
+                    lookupData.getHttpUrl(), lookupData.getBrokerUrl(), brokers.size());
+            return Optional.empty();
+        }
+        Optional<String> protocolData = owner.get().getProtocol(AmqpProtocolHandler.PROTOCOL_NAME);
+        if (protocolData.isEmpty()) {
+            log.warn("Owner broker has no amqp protocol data. webServiceUrl={}", owner.get().getWebServiceUrl());
+            return Optional.empty();
+        }
+        Optional<Integer> adminPort = AmqpProtocolHandler.extractAmqpAdminPort(protocolData.get());
+        if (adminPort.isEmpty()) {
+            log.warn("Owner broker amqp protocol data has no admin port: {}", protocolData.get());
+        }
+        return adminPort;
+    }
+
+    private static boolean matchesOwnerBroker(LoadManagerReport report, LookupData lookupData) {
+        return StringUtils.equals(report.getWebServiceUrl(), lookupData.getHttpUrl())
+                || StringUtils.equals(report.getWebServiceUrlTls(), lookupData.getHttpUrlTls())
+                || StringUtils.equals(report.getPulsarServiceUrl(), lookupData.getBrokerUrl())
+                || StringUtils.equals(report.getPulsarServiceUrlTls(), lookupData.getBrokerUrlTls());
+    }
+
     protected static boolean isLeaderBroker(PulsarService pulsar) {
         return  pulsar.getLeaderElectionService().isLeader();
     }

amqp-impl/src/main/java/io/streamnative/pulsar/handlers/amqp/proxy/PulsarServiceLookupHandler.java

@@ -84,7 +84,7 @@ public CompletableFuture<Pair<String, Integer>> findBroker(TopicName topicName,
                 return;
             }
 
-            String amqpBrokerAddress = protocolData.get();
+            String amqpBrokerAddress = AmqpProtocolHandler.extractAmqpListeners(protocolData.get());
             if (!StringUtils.startsWith(amqpBrokerAddress, AmqpProtocolHandler.PLAINTEXT_PREFIX)
                 && !StringUtils.startsWith(amqpBrokerAddress, AmqpProtocolHandler.SSL_PREFIX)) {
                 amqpBrokerAddress = AmqpProtocolHandler.PLAINTEXT_PREFIX + amqpBrokerAddress;

amqp-impl/src/test/java/io/streamnative/pulsar/handlers/amqp/test/AmqpProtocolDataTest.java

@@ -0,0 +1,50 @@
+/**
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.streamnative.pulsar.handlers.amqp.test;
+
+import io.streamnative.pulsar.handlers.amqp.AmqpProtocolHandler;
+import io.streamnative.pulsar.handlers.amqp.AmqpServiceConfiguration;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+/**
+ * Protocol advertise data format tests.
+ */
+public class AmqpProtocolDataTest {
+
+    @Test
+    public void testAdvertiseDataIncludesAdminPort() {
+        AmqpServiceConfiguration conf = new AmqpServiceConfiguration();
+        conf.setAmqpListeners("amqp://127.0.0.1:5672");
+        conf.setAmqpAdminPort(15673);
+
+        AmqpProtocolHandler handler = new AmqpProtocolHandler();
+        try {
+            handler.initialize(conf);
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+        String protocolData = handler.getProtocolDataToAdvertise();
+        Assert.assertEquals(protocolData, "amqp://127.0.0.1:5672|15673");
+        Assert.assertEquals(AmqpProtocolHandler.extractAmqpListeners(protocolData), "amqp://127.0.0.1:5672");
+        Assert.assertEquals(AmqpProtocolHandler.extractAmqpAdminPort(protocolData).orElse(-1).intValue(), 15673);
+    }
+
+    @Test
+    public void testLegacyProtocolDataCompatible() {
+        String legacy = "amqp://127.0.0.1:5672";
+        Assert.assertEquals(AmqpProtocolHandler.extractAmqpListeners(legacy), legacy);
+        Assert.assertFalse(AmqpProtocolHandler.extractAmqpAdminPort(legacy).isPresent());
+    }
+}