You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pulsar.apache.org by lh...@apache.org on 2025/11/08 11:24:04 UTC
(pulsar) branch master updated: [feat] PIP-442: Add memory limits for CommandGetTopicsOfNamespace (#24833)
This is an automated email from the ASF dual-hosted git repository.
lhotari pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pulsar.git
The following commit(s) were added to refs/heads/master by this push:
new 3aaf34576f1 [feat] PIP-442: Add memory limits for CommandGetTopicsOfNamespace (#24833)
3aaf34576f1 is described below
commit 3aaf34576f15378669ff9b192ef4378908151509
Author: Lari Hotari <lh...@users.noreply.github.com>
AuthorDate: Sat Nov 8 13:23:58 2025 +0200
[feat] PIP-442: Add memory limits for CommandGetTopicsOfNamespace (#24833)
---
conf/broker.conf | 24 +
conf/proxy.conf | 24 +
conf/standalone.conf | 24 +
pip/pip-442.md | 600 +++++++++++------
.../apache/pulsar/broker/ServiceConfiguration.java | 36 +
.../topiclistlimit/TopicListMemoryLimiter.java | 312 +++++++++
.../topiclistlimit/TopicListSizeResultCache.java | 122 ++++
.../pulsar/broker/topiclistlimit/package-info.java | 22 +
.../stats/prometheus/PrometheusMetricsClient.java | 2 +
.../TopicListSizeResultCacheTest.java | 321 +++++++++
.../pulsar/broker/admin/impl/NamespacesBase.java | 68 +-
.../apache/pulsar/broker/admin/v1/Namespaces.java | 2 +-
.../apache/pulsar/broker/admin/v2/Namespaces.java | 2 +-
.../pulsar/broker/service/BrokerService.java | 19 +
.../pulsar/broker/service/PulsarCommandSender.java | 7 +-
.../broker/service/PulsarCommandSenderImpl.java | 18 +-
.../apache/pulsar/broker/service/ServerCnx.java | 138 ++--
.../broker/service/MessageCumulativeAckTest.java | 4 +-
.../PersistentDispatcherFailoverConsumerTest.java | 6 +-
.../pulsar/broker/service/PersistentTopicTest.java | 9 +-
...ackPressureMultipleConsumersHttpLookupTest.java | 50 ++
...nConsumerBackPressureMultipleConsumersTest.java | 228 +++++++
.../api/PatternConsumerBackPressureTest.java | 2 +
.../apache/pulsar/common/protocol/Commands.java | 17 +-
.../common/semaphore/AsyncDualMemoryLimiter.java | 114 ++++
.../semaphore/AsyncDualMemoryLimiterImpl.java | 177 +++++
.../semaphore/AsyncDualMemoryLimiterUtil.java | 123 ++++
.../pulsar/common/semaphore/AsyncSemaphore.java | 129 ++++
.../common/semaphore/AsyncSemaphoreImpl.java | 369 +++++++++++
.../pulsar/common/semaphore/package-info.java | 23 +
.../semaphore/AsyncDualMemoryLimiterImplTest.java | 721 +++++++++++++++++++++
.../semaphore/AsyncDualMemoryLimiterUtilTest.java | 581 +++++++++++++++++
.../common/semaphore/AsyncSemaphoreImplTest.java | 669 +++++++++++++++++++
pulsar-proxy/pom.xml | 12 +
.../pulsar/proxy/server/LookupProxyHandler.java | 94 ++-
.../pulsar/proxy/server/ProxyConfiguration.java | 36 +
.../apache/pulsar/proxy/server/ProxyService.java | 27 +-
.../InvalidProxyConfigForAuthorizationTest.java | 1 +
...nConsumerBackPressureMultipleConsumersTest.java | 92 +++
39 files changed, 4944 insertions(+), 281 deletions(-)
diff --git a/conf/broker.conf b/conf/broker.conf
index 29e060f4a38..997556c1740 100644
--- a/conf/broker.conf
+++ b/conf/broker.conf
@@ -567,6 +567,30 @@ allowOverrideEntryFilters=false
# Max number of concurrent lookup request broker allows to throttle heavy incoming lookup traffic
maxConcurrentLookupRequest=50000
+# Maximum heap memory for inflight topic list operations (MB).
+# Default: 100 MB (supports ~1M topic names assuming 100 bytes each)
+maxTopicListInFlightHeapMemSizeMB=100
+
+# Maximum direct memory for inflight topic list responses (MB).
+# Default: 100 MB (network buffers for serialized responses)
+maxTopicListInFlightDirectMemSizeMB=100
+
+# Timeout for acquiring heap memory permits (milliseconds).
+# Default: 25000 (25 seconds)
+maxTopicListInFlightHeapMemSizePermitsAcquireTimeoutMillis=25000
+
+# Maximum queue size for heap memory permit requests.
+# Default: 10000 (prevent unbounded queueing)
+maxTopicListInFlightHeapMemSizePermitsAcquireQueueSize=10000
+
+# Timeout for acquiring direct memory permits (milliseconds).
+# Default: 25000 (25 seconds)
+maxTopicListInFlightDirectMemSizePermitsAcquireTimeoutMillis=25000
+
+# Maximum queue size for direct memory permit requests.
+# Default: 10000 (prevent unbounded queueing)
+maxTopicListInFlightDirectMemSizePermitsAcquireQueueSize=10000
+
# Max number of concurrent topic loading request broker allows to control number of zk-operations
maxConcurrentTopicLoadRequest=5000
diff --git a/conf/proxy.conf b/conf/proxy.conf
index dc1fc5f002f..6cfe8213fab 100644
--- a/conf/proxy.conf
+++ b/conf/proxy.conf
@@ -228,6 +228,30 @@ maxConcurrentInboundConnectionsPerIp=0
# Max concurrent outbound connections. The proxy will error out requests beyond that.
maxConcurrentLookupRequests=50000
+# Maximum heap memory for inflight topic list operations (MB).
+# Default: 100 MB (supports ~1M topic names assuming 100 bytes each)
+maxTopicListInFlightHeapMemSizeMB=100
+
+# Maximum direct memory for inflight topic list responses (MB).
+# Default: 100 MB (network buffers for serialized responses)
+maxTopicListInFlightDirectMemSizeMB=100
+
+# Timeout for acquiring heap memory permits (milliseconds).
+# Default: 25000 (25 seconds)
+maxTopicListInFlightHeapMemSizePermitsAcquireTimeoutMillis=25000
+
+# Maximum queue size for heap memory permit requests.
+# Default: 10000 (prevent unbounded queueing)
+maxTopicListInFlightHeapMemSizePermitsAcquireQueueSize=10000
+
+# Timeout for acquiring direct memory permits (milliseconds).
+# Default: 25000 (25 seconds)
+maxTopicListInFlightDirectMemSizePermitsAcquireTimeoutMillis=25000
+
+# Maximum queue size for direct memory permit requests.
+# Default: 10000 (prevent unbounded queueing)
+maxTopicListInFlightDirectMemSizePermitsAcquireQueueSize=10000
+
##### --- TLS --- #####
# Deprecated - use servicePortTls and webServicePortTls instead
diff --git a/conf/standalone.conf b/conf/standalone.conf
index 5aa2b7cb73d..7f640a94f3d 100644
--- a/conf/standalone.conf
+++ b/conf/standalone.conf
@@ -320,6 +320,30 @@ preciseDispatcherFlowControl=false
# Max number of concurrent lookup request broker allows to throttle heavy incoming lookup traffic
maxConcurrentLookupRequest=50000
+# Maximum heap memory for inflight topic list operations (MB).
+# Default: 100 MB (supports ~1M topic names assuming 100 bytes each)
+maxTopicListInFlightHeapMemSizeMB=100
+
+# Maximum direct memory for inflight topic list responses (MB).
+# Default: 100 MB (network buffers for serialized responses)
+maxTopicListInFlightDirectMemSizeMB=100
+
+# Timeout for acquiring heap memory permits (milliseconds).
+# Default: 25000 (25 seconds)
+maxTopicListInFlightHeapMemSizePermitsAcquireTimeoutMillis=25000
+
+# Maximum queue size for heap memory permit requests.
+# Default: 10000 (prevent unbounded queueing)
+maxTopicListInFlightHeapMemSizePermitsAcquireQueueSize=10000
+
+# Timeout for acquiring direct memory permits (milliseconds).
+# Default: 25000 (25 seconds)
+maxTopicListInFlightDirectMemSizePermitsAcquireTimeoutMillis=25000
+
+# Maximum queue size for direct memory permit requests.
+# Default: 10000 (prevent unbounded queueing)
+maxTopicListInFlightDirectMemSizePermitsAcquireQueueSize=10000
+
# Max number of concurrent topic loading request broker allows to control number of zk-operations
maxConcurrentTopicLoadRequest=5000
diff --git a/pip/pip-442.md b/pip/pip-442.md
index 19e0500d2de..59a0cdebc04 100644
--- a/pip/pip-442.md
+++ b/pip/pip-442.md
@@ -14,8 +14,8 @@ Pulsar brokers already implement comprehensive memory management for most operat
- `maxMessagePublishBufferSizeInMB` (default: 50% direct memory): Limits memory used for buffering messages during publishing, providing backpressure when producers exceed broker capacity
**Managed Ledger Memory Limits:**
-- `managedLedgerMaxReadsInFlightSizeInMB` (default: 0, disabled): Controls memory allocation for concurrent read operations from BookKeeper, preventing excessive memory usage during high read loads. This limit extends to cover buffers that were read from BookKeeper and are waiting in channel outbound buffers to be written to client sockets.
-- `managedLedgerCacheSizeMB` (default: 20% of direct memory): Limits cache memory for recently read ledger entries, ensuring predictable memory usage for read caching. This limit extends to cover buffers that were read from the cache and are waiting in channel outbound buffers to be written to client sockets.
+- `managedLedgerMaxReadsInFlightSizeInMB` (default: 0, disabled): Controls memory allocation for concurrent read operations from BookKeeper, preventing excessive memory usage during high read loads. This limit extends to cover buffers that were read from BookKeeper and are waiting in channel outbound buffers to be written to client sockets.
+- `managedLedgerCacheSizeMB` (default: 20% of direct memory): Limits cache memory for recently read ledger entries, ensuring predictable memory usage for read caching. This limit extends to cover buffers that were read from the cache and are waiting in channel outbound buffers to be written to client sockets.
**Additional Memory Controls:**
- `maxConcurrentLookupRequest` (default: 50000): Limits concurrent topic lookup requests. The unit of this limit is the number of requests; it is not expressed in memory size.
@@ -46,14 +46,14 @@ These operations can allocate arbitrary amounts of memory based on namespace siz
The `getTopicsOfNamespace` request follows this path:
1. **Client Request**: Sends `CommandGetTopicsOfNamespace` via binary protocol
-2. **Request Handling**:
- - Broker: `ServerCnx.handleGetTopicsOfNamespace()`
- - Proxy: `ProxyConnection.handleGetTopicsOfNamespace()`
+2. **Request Handling**:
+ - Broker: `ServerCnx.handleGetTopicsOfNamespace()`
+ - Proxy: `LookupProxyHandler.handleGetTopicsOfNamespace()`
3. **Topic Retrieval**: `NamespaceService.getListOfUserTopics()` orchestrates:
- - Fetches persistent topics from `TopicResources`
- - Retrieves non-persistent topics from local cache or peer clusters
- - Filters system topics using `TopicList.filterSystemTopic()`
- - Implements caching via `inProgressQueryUserTopics` to prevent duplicate queries
+ - Fetches persistent topics from `TopicResources`
+ - Retrieves non-persistent topics from local cache or peer clusters
+ - Filters system topics using `TopicList.filterSystemTopic()`
+ - Implements caching via `inProgressQueryUserTopics` to prevent duplicate queries
4. **Response Construction**: Packages results with hash calculation and filtering metadata
5. **Response Transmission**: Sends complete response back to client
@@ -107,49 +107,75 @@ The lack of memory limits for topic listing commands creates the final significa
The solution introduces an `AsyncDualMemoryLimiter` that acts as a memory-aware semaphore for topic listing operations, completing Pulsar's memory management framework:
-1. **Memory Tracking**: Before processing requests or sending responses, the system estimates memory requirements and acquires permits from the limiter. When the permit cannot be estimated and allocated before the operation, an initial permit is acquired and updated before continuing with handling. Although not optimal, this will effectively limit memory usage across the broker.
+1. **Memory Tracking**: Before processing requests or sending responses, the system estimates memory requirements and acquires permits from the limiter. An initial permit is acquired with a fixed estimate (1KB), then updated to reflect the actual memory usage after the topic list is retrieved.
2. **Dual Memory Pools**: Separate tracking for heap memory (topic list assembly) and direct memory (network buffers) with independent limits, since topic listing operations use both types of memory.
-3. **Asynchronous Backpressure**: When memory limits are reached, requests queue with configurable timeouts rather than failing immediately, providing graceful degradation similar to `managedLedgerMaxReadsInFlightSizeInMB` behavior. This type of solution is helpful since rejecting requests and requiring clients to retry can cause more load on the system and would cause unfair queueing. When the queue is completely full, requests are rejected.
+3. **Asynchronous Backpressure**: When memory limits are reached, requests queue with configurable timeouts rather than failing immediately, providing graceful degradation similar to `managedLedgerMaxReadsInFlightSizeInMB` behavior. When the queue is completely full, requests are rejected.
4. **Graceful Degradation**: The system continues processing within memory limits rather than crashing, with clear metrics indicating when memory-based throttling occurs.
5. **Release Guarantees**: Memory permits are released after response transmission completes or on request failure, preventing memory leaks and ensuring accurate memory tracking.
+6. **Cancellation Support**: The implementation supports cancellation of permit requests when the client connection is closed, preventing unnecessary queueing and resource allocation.
+
## Detailed Design
### Design & Implementation Details
#### AsyncSemaphore Interface
-This is an abstraction for a generic asynchronous semaphore. The memory limiter implementation will use this abstraction to implement separate limiters for heap and direct memory.
+This is an abstraction for a generic asynchronous semaphore. The memory limiter implementation uses this abstraction to implement separate limiters for heap and direct memory.
```java
public interface AsyncSemaphore {
/**
* Acquire permits from the semaphore.
* Returned future completes when permits are available.
- * It will complete exceptionally with AsyncSemaphorePermitAcquireTimeoutException on timeout
- * and exceptionally with AsyncSemaphorePermitAcquireQueueFullException when queue full
+ * It will complete exceptionally with AsyncSemaphore.PermitAcquireTimeoutException on timeout
+ * and exceptionally with AsyncSemaphore.PermitAcquireQueueFullException when queue full
+ *
+ * @param permits number of permits to acquire
+ * @param isCancelled supplier that returns true if acquisition should be cancelled
* @return CompletableFuture that completes with permit when available
*/
- CompletableFuture<AsyncSemaphorePermitResult> acquire(long permits);
+ CompletableFuture<AsyncSemaphorePermit> acquire(long permits, BooleanSupplier isCancelled);
/**
* Acquire or release permits for previously acquired permits by updating the permits.
* Returns a future that completes when permits are available.
- * It will complete exceptionally with AsyncSemaphorePermitAcquireTimeoutException on timeout
- * and exceptionally with AsyncSemaphorePermitAcquireQueueFullException when queue full
+ * It will complete exceptionally with AsyncSemaphore.PermitAcquireTimeoutException on timeout
+ * and exceptionally with AsyncSemaphore.PermitAcquireQueueFullException when queue full
+ *
+ * @param permit previously acquired permit to update
+ * @param newPermits new number of permits to update to
+ * @param isCancelled supplier that returns true if update should be cancelled
* @return CompletableFuture that completes with permit when available
*/
- CompletableFuture<AsyncSemaphorePermit> update(AsyncSemaphorePermit permit, long newPermits);
-
+ CompletableFuture<AsyncSemaphorePermit> update(AsyncSemaphorePermit permit, long newPermits,
+ BooleanSupplier isCancelled);
+
/**
* Release previously acquired permit.
- * Must be called to prevent memory permit leaks.
+ * Must be called to prevent permit leaks.
+ *
+ * @param permit permit to release
*/
void release(AsyncSemaphorePermit permit);
+ /**
+ * Get the number of available permits.
+ */
+ long getAvailablePermits();
+
+ /**
+ * Get the number of acquired permits.
+ */
+ long getAcquiredPermits();
+
+ /**
+ * Get the current size of queued requests.
+ */
+ int getQueueSize();
}
```
@@ -163,227 +189,409 @@ public interface AsyncDualMemoryLimiter {
HEAP_MEMORY, // For heap memory allocation
DIRECT_MEMORY // For direct memory allocation
}
-
+
/**
* Acquire permits for the specified memory size.
* Returned future completes when memory permits are available.
- * It will complete exceptionally with AsyncSemaphorePermitAcquireTimeoutException on timeout
- * and exceptionally with AsyncSemaphorePermitAcquireQueueFullException when queue full
+ * It will complete exceptionally with AsyncSemaphore.PermitAcquireTimeoutException on timeout
+ * and exceptionally with AsyncSemaphore.PermitAcquireQueueFullException when queue full
+ *
+ * @param memorySize the size of memory to acquire permits for
+ * @param limitType the type of memory limit (HEAP_MEMORY or DIRECT_MEMORY)
+ * @param isCancelled supplier that returns true if acquisition should be cancelled
* @return CompletableFuture that completes with permit when available
*/
- CompletableFuture<AsyncDualMemoryLimiterPermit> acquire(long memorySize, LimitType limitType);
+ CompletableFuture<AsyncDualMemoryLimiterPermit> acquire(long memorySize, LimitType limitType,
+ BooleanSupplier isCancelled);
/**
* Acquire or release permits for previously acquired permits by updating the requested memory size.
* Returns a future that completes when permits are available.
- * It will complete exceptionally with AsyncSemaphorePermitAcquireTimeoutException on timeout
- * and exceptionally with AsyncSemaphorePermitAcquireQueueFullException when queue full
+ * It will complete exceptionally with AsyncSemaphore.PermitAcquireTimeoutException on timeout
+ * and exceptionally with AsyncSemaphore.PermitAcquireQueueFullException when queue full
+ * The provided permit is released when the permits are successfully acquired and the returned updated
+ * permit replaces the old instance.
+ *
+ * @param permit the previously acquired permit to update
+ * @param newMemorySize the new memory size to update to
+ * @param isCancelled supplier that returns true if update should be cancelled
* @return CompletableFuture that completes with permit when available
*/
- CompletableFuture<AsyncDualMemoryLimiterPermit> update(AsyncDualMemoryLimiterPermit permit, long newMemorySize);
-
+ CompletableFuture<AsyncDualMemoryLimiterPermit> update(AsyncDualMemoryLimiterPermit permit, long newMemorySize,
+ BooleanSupplier isCancelled);
+
/**
* Release previously acquired permit.
* Must be called to prevent memory permit leaks.
+ *
+ * @param permit the permit to release
*/
void release(AsyncDualMemoryLimiterPermit permit);
+ /**
+ * Execute the specified function with acquired permits and release the permits after the returned future completes.
+ * @param memorySize memory size to acquire permits for
+ * @param limitType memory limit type to acquire permits for
+ * @param function function to execute with acquired permits
+ * @return result of the function
+ * @param <T> type of the CompletableFuture returned by the function
+ */
+ default <T> CompletableFuture<T> withAcquiredPermits(long memorySize, LimitType limitType,
+ BooleanSupplier isCancelled,
+ Function<AsyncDualMemoryLimiterPermit,
+ CompletableFuture<T>> function,
+ Function<Throwable, CompletableFuture<T>>
+ permitAcquireErrorHandler) {
+ return AsyncDualMemoryLimiterUtil.withPermitsFuture(acquire(memorySize, limitType, isCancelled), function,
+ permitAcquireErrorHandler, this::release);
+ }
+
+ /**
+ * Executed the specified function with updated permits and release the permits after the returned future completes.
+ * @param initialPermit initial permit to update
+ * @param newMemorySize new memory size to update to
+ * @param function function to execute with updated permits
+ * @return result of the function
+ * @param <T> type of the CompletableFuture returned by the function
+ */
+ default <T> CompletableFuture<T> withUpdatedPermits(AsyncDualMemoryLimiterPermit initialPermit, long newMemorySize,
+ BooleanSupplier isCancelled,
+ Function<AsyncDualMemoryLimiterPermit,
+ CompletableFuture<T>> function,
+ Function<Throwable, CompletableFuture<T>>
+ permitAcquireErrorHandler) {
+ return AsyncDualMemoryLimiterUtil.withPermitsFuture(update(initialPermit, newMemorySize, isCancelled), function,
+ permitAcquireErrorHandler, this::release);
+ }
+}
+```
+
+#### AsyncDualMemoryLimiterUtil Helper
+
+A utility class provides helper methods for common patterns:
+
+```java
+public class AsyncDualMemoryLimiterUtil {
+ /**
+ * Execute a function with acquired permits and ensure permits are released after completion.
+ * This method handles the lifecycle of permits - acquisition, usage, and release, including error cases.
+ *
+ * @param permitsFuture Future that will complete with the required permits
+ * @param function Function to execute once permits are acquired that returns a CompletableFuture
+ * @param permitAcquireErrorHandler Handler for permit acquisition errors that returns a CompletableFuture
+ * @param releaser Consumer that handles releasing the permits
+ * @param <T> The type of result returned by the function
+ * @return CompletableFuture that completes with the result of the function execution
+ */
+ public static <T> CompletableFuture<T> withPermitsFuture(
+ CompletableFuture<AsyncDualMemoryLimiterPermit>
+ permitsFuture,
+ Function<AsyncDualMemoryLimiterPermit,
+ CompletableFuture<T>> function,
+ Function<Throwable, CompletableFuture<T>>
+ permitAcquireErrorHandler,
+ Consumer<AsyncDualMemoryLimiterPermit> releaser) {
+ // implementation omitted from PIP document
+ }
+
+ /**
+ * Acquires permits and writes the command as a response to the channel.
+ * Releases the permits after the response has been written to the socket or if the write fails.
+ *
+ * @param ctx the channel handler context used for writing the response
+ * @param dualMemoryLimiter the memory limiter used to acquire and release memory permits
+ * @param isCancelled supplier that indicates if the permit acquisition should be cancelled
+ * @param command the base command to serialize and write to the channel
+ * @param permitAcquireErrorHandler handler for errors that occur during permit acquisition
+ * @return a future that completes when the command has been written to the channel's outbound buffer
+ */
+ public static CompletableFuture<Void> acquireDirectMemoryPermitsAndWriteAndFlush(ChannelHandlerContext ctx,
+ AsyncDualMemoryLimiter
+ dualMemoryLimiter,
+ BooleanSupplier isCancelled,
+ BaseCommand command,
+ Consumer<Throwable>
+ permitAcquireErrorHandler
+ ) {
+ // implementation omitted from PIP document
+ }
}
```
#### Integration Points
-**1. Heap Memory Limiting (Post-Retrieval)**
+**1. Heap Memory Limiting (Post-Retrieval) - Broker**
-In `ServerCnx.handleGetTopicsOfNamespace`:
+In `ServerCnx.handleGetTopicsOfNamespace`, the implementation uses the helper methods:
```java
-// Acquire a fixed amount of permits initially since it's not known how much memory will be used
-// This will ensure that the operation continues only after it has the initial permits
-// It would be possible to use statistics for initial estimate, but this is simpler and sufficient
-maxTopicListInFlightLimiter.acquire(INITIAL_SIZE, AsyncDualMemoryLimiter.LimitType.HEAP_MEMORY)
- .thenCompose(initialPermit -> {
- getBrokerService().pulsar().getNamespaceService().getListOfUserTopics(namespaceName, mode)
- .thenCompose(topics -> {
- // Estimate memory after retrieval and update the permits to reflect the actual size
- long estimatedSize = topics.stream().mapToInt(String::length).sum();
- return maxTopicListInFlightLimiter
- .update(initialPermit, estimatedSize)
- .thenApply(permit -> Pair.of(topics, permit));
- })
- .thenAccept(topicsAndPermit -> {
- try {
- // Process and send response
- ...
- } finally {
- maxTopicListInFlightLimiter.release(topicsAndPermit.getRight());
- }
+ private void internalHandleGetTopicsOfNamespace(String namespace, NamespaceName namespaceName, long requestId,
+ CommandGetTopicsOfNamespace.Mode mode,
+ Optional<String> topicsPattern, Optional<String> topicsHash,
+ Semaphore lookupSemaphore) {
+ BooleanSupplier isPermitRequestCancelled = () -> !ctx().channel().isActive();
+ maxTopicListInFlightLimiter.withAcquiredPermits(INITIAL_TOPIC_LIST_HEAP_PERMITS_SIZE,
+ AsyncDualMemoryLimiter.LimitType.HEAP_MEMORY, isPermitRequestCancelled, initialPermits -> {
+ return getBrokerService().pulsar().getNamespaceService().getListOfUserTopics(namespaceName, mode)
+ .thenAccept(topics -> {
+ long actualSize = topics.stream().mapToInt(String::length).sum();
+ maxTopicListInFlightLimiter.withUpdatedPermits(initialPermits, actualSize,
+ isPermitRequestCancelled, permits -> {
+ boolean filterTopics = false;
+ // filter system topic
+ List<String> filteredTopics = topics;
+
+ if (enableSubscriptionPatternEvaluation && topicsPattern.isPresent()) {
+ if (topicsPattern.get().length() <= maxSubscriptionPatternLength) {
+ filterTopics = true;
+ filteredTopics = TopicList.filterTopics(filteredTopics, topicsPattern.get(),
+ topicsPatternImplementation);
+ } else {
+ log.info("[{}] Subscription pattern provided [{}] was longer than "
+ + "maximum {}.", remoteAddress, topicsPattern.get(),
+ maxSubscriptionPatternLength);
+ }
+ }
+ String hash = TopicList.calculateHash(filteredTopics);
+ boolean hashUnchanged = topicsHash.isPresent() && topicsHash.get().equals(hash);
+ if (hashUnchanged) {
+ filteredTopics = Collections.emptyList();
+ }
+ if (log.isDebugEnabled()) {
+ log.debug("[{}] Received CommandGetTopicsOfNamespace for namespace "
+ + "[//{}] by {}, size:{}", remoteAddress, namespace,
+ requestId,
+ topics.size());
+ }
+ commandSender.sendGetTopicsOfNamespaceResponse(filteredTopics, hash, filterTopics,
+ !hashUnchanged, requestId, ex -> {
+ log.warn("[{}] Failed to acquire direct memory permits for "
+ + "GetTopicsOfNamespace: {}", remoteAddress, ex.getMessage());
+ commandSender.sendErrorResponse(requestId, ServerError.TooManyRequests,
+ "Cannot acquire permits for direct memory");
+ });
+ return CompletableFuture.completedFuture(null);
+ }, t -> {
+ log.warn("[{}] Failed to acquire heap memory permits for "
+ + "GetTopicsOfNamespace: {}", remoteAddress, t.getMessage());
+ writeAndFlush(Commands.newError(requestId, ServerError.TooManyRequests,
+ "Failed due to heap memory limit exceeded"));
+ return CompletableFuture.completedFuture(null);
+ });
+ }).whenComplete((__, ___) -> {
+ lookupSemaphore.release();
+ }).exceptionally(ex -> {
+ log.warn("[{}] Error GetTopicsOfNamespace for namespace [//{}] by {}", remoteAddress,
+ namespace, requestId);
+ commandSender.sendErrorResponse(requestId,
+ BrokerServiceException.getClientErrorCode(new ServerMetadataException(ex)),
+ ex.getMessage());
+ return null;
+ });
+ }, t -> {
+ log.warn("[{}] Failed to acquire initial heap memory permits for GetTopicsOfNamespace: {}",
+ remoteAddress, t.getMessage());
+ writeAndFlush(Commands.newError(requestId, ServerError.TooManyRequests,
+ "Failed due to heap memory limit exceeded"));
+ lookupSemaphore.release();
+ return CompletableFuture.completedFuture(null);
});
- ...
- // For exceptional paths, initialPermit would need to be released
+}
```
-**2. Direct Memory Limiting (Pre-Serialization)**
+**2. Direct Memory Limiting (Pre-Serialization) - Broker**
-Modified `CommandSender` implementation:
+Modified `PulsarCommandSenderImpl`:
```java
@Override
public void sendGetTopicsOfNamespaceResponse(List<String> topics, String topicsHash,
- boolean filtered, boolean changed, long requestId) {
+ boolean filtered, boolean changed, long requestId,
+ Consumer<Throwable> permitAcquireErrorHandler) {
BaseCommand command = Commands.newGetTopicsOfNamespaceResponseCommand(topics, topicsHash,
filtered, changed, requestId);
safeIntercept(command, cnx);
- acquireMaxTopicListInFlightPermitsAndWriteAndFlush(command);
+ acquireDirectMemoryPermitsAndWriteAndFlush(cnx.ctx(), maxTopicListInFlightLimiter, () -> !cnx.isActive(),
+ command, permitAcquireErrorHandler);
}
+```
+
+The utility method implementation:
-private void acquireMaxTopicListInFlightPermitsAndWriteAndFlush(BaseCommand command) {
+```java
+public static CompletableFuture<Void> acquireDirectMemoryPermitsAndWriteAndFlush(ChannelHandlerContext ctx,
+ AsyncDualMemoryLimiter
+ dualMemoryLimiter,
+ BooleanSupplier isCancelled,
+ BaseCommand command,
+ Consumer<Throwable>
+ permitAcquireErrorHandler
+) {
// Calculate serialized size before acquiring permits
int serializedSize = command.getSerializedSize();
// Acquire permits
- maxTopicListInFlightLimiter.acquire(serializedSize, AsyncDualMemoryLimiter.LimitType.DIRECT_MEMORY)
- .thenAcceptAsync(permits -> {
+ return dualMemoryLimiter.acquire(serializedSize, AsyncDualMemoryLimiter.LimitType.DIRECT_MEMORY, isCancelled)
+ .whenComplete((permits, t) -> {
+ if (t != null) {
+ permitAcquireErrorHandler.accept(t);
+ return;
+ }
try {
// Serialize the response
ByteBuf outBuf = Commands.serializeWithPrecalculatedSerializedSize(command, serializedSize);
// Write the response
- cnx.ctx().writeAndFlush(outBuf).addListener(future -> {
+ ctx.writeAndFlush(outBuf).addListener(future -> {
// Release permits after the response has been written to the socket
- maxTopicListInFlightLimiter.release(permits);
+ dualMemoryLimiter.release(permits);
});
- } catch (Exception e) {
+ } catch (Throwable e) {
// Return permits if an exception occurs before writeAndFlush is called successfully
- maxTopicListInFlightLimiter.release(permits);
+ dualMemoryLimiter.release(permits);
throw e;
}
- }, cnx.ctx().executor());
+ }).thenApply(__ -> null);
}
```
-**3. Watch Command Memory Control**
+**3. Watch Command Memory Control - Broker**
-Similar memory limiting patterns apply to watch commands:
+Similar memory limiting patterns apply to watch commands in `TopicListService`:
```java
-public void sendWatchTopicListSuccess(long requestId, long watcherId, String topicsHash, List<String> topics) {
- BaseCommand command = Commands.newWatchTopicListSuccess(requestId, watcherId, topicsHash, topics);
- acquireMaxTopicListInFlightPermitsAndWriteAndFlush(command);
-}
-
-public void sendWatchTopicListUpdate(long watcherId, List<String> newTopics, List<String> deletedTopics, String topicsHash) {
- BaseCommand command = Commands.newWatchTopicUpdate(watcherId, newTopics, deletedTopics, topicsHash);
- acquireMaxTopicListInFlightPermitsAndWriteAndFlush(command);
+public void sendTopicListUpdate(long watcherId, String topicsHash,
+ List<String> deletedTopics, List<String> newTopics) {
+ connection.getCommandSender().sendWatchTopicListUpdate(
+ watcherId, newTopics, deletedTopics, topicsHash,
+ t -> {
+ log.warn("[{}] Cannot acquire direct memory tokens for sending topic list update",
+ connection.toString(), t);
+ });
}
```
-**4. Proxy Reading Memory Control**
-
-On the Pulsar Proxy side, the problem is slightly different. The problem occurs when the proxy receives a `CommandGetTopicsOfNamespace` command, forwards it to a broker, and receives a response. The proxy needs to deserialize and serialize the response before sending it to the client.
-Memory is allocated for both deserialization and serialization.
-
-Solving this requires a slight modification to PulsarDecoder.
+**4. Proxy Memory Control**
-In `PulsarDecoder.channelRead`, it would be necessary to record the size of the incoming message:
+On the Pulsar Proxy side in `LookupProxyHandler`:
```java
- // Get a buffer that contains the full frame
- ByteBuf buffer = (ByteBuf) msg;
- try {
- // De-serialize the command
- int cmdSize = (int) buffer.readUnsignedInt();
- cmd.parseFrom(buffer, cmdSize);
-```
+private void internalPerformGetTopicsOfNamespace(long clientRequestId, String namespaceName, ClientCnx clientCnx,
+ ByteBuf command, long requestId) {
+ BooleanSupplier isPermitRequestCancelled = () -> !proxyConnection.ctx().channel().isActive();
+ maxTopicListInFlightLimiter.withAcquiredPermits(INITIAL_TOPIC_LIST_HEAP_PERMITS_SIZE,
+ AsyncDualMemoryLimiter.LimitType.HEAP_MEMORY, isPermitRequestCancelled, initialPermits -> {
+ return clientCnx.newGetTopicsOfNamespace(command, requestId).whenComplete((r, t) -> {
+ if (t != null) {
+ log.warn("[{}] Failed to get TopicsOfNamespace {}: {}", clientAddress, namespaceName,
+ t.getMessage());
+ writeAndFlush(Commands.newError(clientRequestId, getServerError(t), t.getMessage()));
+ } else {
+ long actualSize =
+ r.getNonPartitionedOrPartitionTopics().stream().mapToInt(String::length).sum();
+ maxTopicListInFlightLimiter.withUpdatedPermits(initialPermits, actualSize,
+ isPermitRequestCancelled, permits -> {
+ return handleWritingGetTopicsResponse(clientRequestId, r, isPermitRequestCancelled);
+ }, t2 -> {
+ log.warn("[{}] Failed to acquire actual heap memory permits for "
+ + "GetTopicsOfNamespace: {}", clientAddress, t2.getMessage());
+ writeAndFlush(Commands.newError(clientRequestId, ServerError.TooManyRequests,
+ "Failed due to heap memory limit exceeded"));
+
+ return CompletableFuture.completedFuture(null);
+ });
+ }
+ }).thenApply(__ -> null);
+ }, t -> {
+ log.warn("[{}] Failed to acquire initial heap memory permits for GetTopicsOfNamespace: {}",
+ clientAddress, t.getMessage());
+ writeAndFlush(Commands.newError(clientRequestId, ServerError.TooManyRequests,
+ "Failed due to heap memory limit exceeded"));
+
+ return CompletableFuture.completedFuture(null);
+ }).exceptionally(ex -> {
+ writeAndFlush(Commands.newError(clientRequestId, getServerError(ex), ex.getMessage()));
+ return null;
+ });
+}
-It could be modified to store the `cmdSize` in a field instead of a local variable:
-```java
- protected int cmdSize;
-...
- // Get a buffer that contains the full frame
- ByteBuf buffer = (ByteBuf) msg;
- try {
- // De-serialize the command
- cmdSize = (int) buffer.readUnsignedInt();
- cmd.parseFrom(buffer, cmdSize);
+private CompletableFuture<Void> handleWritingGetTopicsResponse(long clientRequestId, GetTopicsResult r,
+ BooleanSupplier isCancelled) {
+ BaseCommand responseCommand = Commands.newGetTopicsOfNamespaceResponseCommand(
+ r.getNonPartitionedOrPartitionTopics(), r.getTopicsHash(), r.isFiltered(),
+ r.isChanged(), clientRequestId);
+ return acquireDirectMemoryPermitsAndWriteAndFlush(proxyConnection.ctx(), maxTopicListInFlightLimiter,
+ isCancelled, responseCommand, t -> {
+ log.warn("[{}] Failed to acquire actual direct memory permits for GetTopicsOfNamespace: {}",
+ clientAddress, t.getMessage());
+ writeAndFlush(Commands.newError(clientRequestId, ServerError.TooManyRequests,
+ "Failed due to direct memory limit exceeded"));
+ });
+}
```
-Changes would be needed to be able to use this serialized size so that it doesn't need to be re-calculated.
-`cmdSize` would be added as a field to `GetTopicsResult`:
-
-```java
-@Override
-protected void handleGetTopicsOfNamespaceSuccess(CommandGetTopicsOfNamespaceResponse success) {
- checkArgument(state == State.Ready);
-
- long requestId = success.getRequestId();
- List<String> topics = success.getTopicsList();
+#### TopicListMemoryLimiter - Metrics Integration
+The `TopicListMemoryLimiter` class extends `AsyncDualMemoryLimiterImpl` and adds Prometheus and OpenTelemetry metrics:
- if (log.isDebugEnabled()) {
- log.debug("{} Received get topics of namespace success response from server: {} - topics.size: {}",
- ctx.channel(), success.getRequestId(), topics.size());
+```java
+public class TopicListMemoryLimiter extends AsyncDualMemoryLimiterImpl {
+ private final CollectorRegistry collectorRegistry;
+ private final Gauge heapMemoryUsedBytes;
+ private final Gauge heapMemoryLimitBytes;
+ // ... other Prometheus metrics
+ private final ObservableDoubleGauge otelHeapMemoryUsedGauge;
+ // ... other OpenTelemetry metrics
+
+ public TopicListMemoryLimiter(CollectorRegistry collectorRegistry,
+ String prometheusPrefix,
+ Meter openTelemetryMeter,
+ long maxHeapMemory, int maxHeapQueueSize,
+ long heapTimeoutMillis, long maxDirectMemory,
+ int maxDirectQueueSize, long directTimeoutMillis) {
+ super(maxHeapMemory, maxHeapQueueSize, heapTimeoutMillis,
+ maxDirectMemory, maxDirectQueueSize, directTimeoutMillis);
+
+ // Initialize Prometheus metrics
+ this.heapMemoryUsedBytes = Gauge.build(
+ prometheusPrefix + "topic_list_heap_memory_used_bytes",
+ "Current heap memory used by topic listings")
+ .create()
+ .setChild(new Gauge.Child() {
+ @Override
+ public double get() {
+ return getLimiter(LimitType.HEAP_MEMORY).getAcquiredPermits();
+ }
+ })
+ .register(collectorRegistry);
+
+ // Initialize OpenTelemetry metrics
+ this.otelHeapMemoryUsedGauge = openTelemetryMeter
+ .gaugeBuilder("topic.list.heap.memory.used")
+ .setUnit("By")
+ .setDescription("Current heap memory used by topic listings")
+ .buildWithCallback(measurement -> {
+ measurement.record(getLimiter(LimitType.HEAP_MEMORY).getAcquiredPermits());
+ });
+
+ // ... initialize other metrics
}
-
- CompletableFuture<GetTopicsResult> requestFuture =
- (CompletableFuture<GetTopicsResult>) pendingRequests.remove(requestId);
- if (requestFuture != null) {
- requestFuture.complete(new GetTopicsResult(topics,
- success.hasTopicsHash() ? success.getTopicsHash() : null,
- success.isFiltered(),
- success.isChanged(),
- // Store cmdSize in the GetTopicsResult <----
- cmdSize));
- } else {
- duplicatedResponseCounter.incrementAndGet();
- log.warn("{} Received unknown request id from server: {}", ctx.channel(), success.getRequestId());
+
+ @Override
+ protected void recordHeapWaitTime(long waitTimeNanos) {
+ if (waitTimeNanos == Long.MAX_VALUE) {
+ heapTimeoutTotal.inc();
+ otelHeapTimeoutTotal.add(1);
+ } else {
+ heapWaitTimeMs.observe(TimeUnit.NANOSECONDS.toMillis(waitTimeNanos));
+ otelHeapWaitTime.record(waitTimeNanos / 1_000_000_000.0d);
+ }
+ }
+
+ @Override
+ protected void recordDirectWaitTime(long waitTimeNanos) {
+ // Similar implementation for direct memory
}
}
```
-The limiter would be integrated into `LookupProxyHandler`'s `performGetTopicsOfNamespace` in this way:
-
-```java
- proxyConnection.getConnectionPool().getConnection(addr).thenAccept(clientCnx -> {
- // Connected to backend broker
- long requestId = proxyConnection.newRequestId();
- ByteBuf command;
- command = Commands.newGetTopicsOfNamespaceRequest(namespaceName, requestId, mode,
- topicsPattern, topicsHash);
- // Acquire a fixed amount of permits initially since it's not known how much memory will be used
- // This will ensure that the operation continues only after it has the initial permits
- // It would be possible to use statistics for initial estimate, but this is simpler and sufficient
- maxTopicListInFlightLimiter.acquire(INITIAL_SIZE, AsyncDualMemoryLimiter.LimitType.HEAP_MEMORY)
- .thenCompose(initialPermit -> {
- clientCnx.newGetTopicsOfNamespace(command, requestId).whenComplete((r, t) -> {
- if (t != null) {
- maxTopicListInFlightLimiter.release(initialPermit);
- log.warn("[{}] Failed to get TopicsOfNamespace {}: {}",
- clientAddress, namespaceName, t.getMessage());
- writeAndFlush(
- Commands.newError(clientRequestId, getServerError(t), t.getMessage()));
- } else {
- // Update the initial permits to reflect the actual size of the response
- maxTopicListInFlightLimiter.update(initialPermit, r.getSerializedSize())
- .thenCompose(heapPermit -> {
- // Acquire a direct memory permit for serialization
- maxTopicListInFlightLimiter.acquire(r.getSerializedSize(), AsyncDualMemoryLimiter.LimitType.DIRECT_MEMORY)
- .thenAccept(directPermit -> {
- proxyConnection.ctx().writeAndFlush(
- Commands.newGetTopicsOfNamespaceResponse(r.getNonPartitionedOrPartitionTopics(),
- r.getTopicsHash(), r.isFiltered(),
- r.isChanged(), clientRequestId)
- ).addListener(future -> {
- // Release permits after the response has been written to the socket
- maxTopicListInFlightLimiter.release(heapPermit);
- maxTopicListInFlightLimiter.release(directPermit);
- });
- }) // Add exception handling for releasing directPermit
- }); // Add exception handling for releasing heapPermit
- }
- });
- });
- proxyConnection.getConnectionPool().releaseConnection(clientCnx);
- }).exceptionally(ex -> {
-```
-
### Public-facing Changes
#### Configuration
@@ -404,21 +612,21 @@ maxTopicListInFlightDirectMemSizeMB=100
maxTopicListInFlightHeapMemSizePermitsAcquireTimeoutMillis=25000
# Maximum queue size for heap memory permit requests
-# Default: 1000 (prevent unbounded queueing)
-maxTopicListInFlightHeapMemSizePermitsAcquireQueueSize=1000
+# Default: 10000 (prevent unbounded queueing)
+maxTopicListInFlightHeapMemSizePermitsAcquireQueueSize=10000
# Timeout for acquiring direct memory permits (milliseconds)
# Default: 25000 (25 seconds)
maxTopicListInFlightDirectMemSizePermitsAcquireTimeoutMillis=25000
# Maximum queue size for direct memory permit requests
-# Default: 1000 (prevent unbounded queueing)
-maxTopicListInFlightDirectMemSizePermitsAcquireQueueSize=1000
+# Default: 10000 (prevent unbounded queueing)
+maxTopicListInFlightDirectMemSizePermitsAcquireQueueSize=10000
```
#### Metrics
-New metrics under `pulsar_broker_topic_list_`/`pulsar_proxy_topic_list` prefix, complementing existing memory metrics:
+New metrics under `pulsar_broker_topic_list_`/`pulsar_proxy_topic_list_` prefix, complementing existing memory metrics:
| Metric Name | Type | Description | Labels |
|------------|------|-------------|--------|
@@ -430,11 +638,13 @@ New metrics under `pulsar_broker_topic_list_`/`pulsar_proxy_topic_list` prefix,
| `heap_queue_max_size` | Gauge | Maximum heap memory limiter queue size | `cluster` |
| `direct_queue_size` | Gauge | Current direct memory limiter queue size | `cluster` |
| `direct_queue_max_size` | Gauge | Maximum direct memory limiter queue size | `cluster` |
-| `heap_wait_time_ms` | Histogram | Wait time for heap memory permits | `cluster` |
-| `direct_wait_time_ms` | Histogram | Wait time for direct memory permits | `cluster` |
+| `heap_wait_time_ms` | Summary | Wait time for heap memory permits (quantiles: 0.5, 0.95, 0.99, 1.0) | `cluster` |
+| `direct_wait_time_ms` | Summary | Wait time for direct memory permits (quantiles: 0.5, 0.95, 0.99, 1.0) | `cluster` |
| `heap_timeout_total` | Counter | Total heap memory permit timeouts | `cluster` |
| `direct_timeout_total` | Counter | Total direct memory permit timeouts | `cluster` |
+OpenTelemetry equivalents are also provided with similar naming under the `topic.list.*` namespace.
+
#### Public API
No changes to REST API.
@@ -448,20 +658,20 @@ No protocol changes. Existing commands continue to work with added server-side m
Operators should monitor the following metrics alongside existing memory management metrics and set up alerts:
1. **Memory Utilization Alert**:
- - Alert when `heap_memory_used_bytes / heap_memory_limit_bytes > 0.8`
- - Indicates the need to increase limits or investigate namespace growth
+ - Alert when `heap_memory_used_bytes / heap_memory_limit_bytes > 0.8`
+ - Indicates the need to increase limits or investigate namespace growth
2. **Queue Saturation Alert**:
- - Alert when `heap_queue_size / heap_queue_max_size > 0.9`
- - Indicates sustained memory pressure requiring capacity adjustment
+ - Alert when `heap_queue_size / heap_queue_max_size > 0.9`
+ - Indicates sustained memory pressure requiring capacity adjustment
3. **Timeout Rate Alert**:
- - Alert when `rate(heap_timeout_total[5m]) > 1`
- - Indicates clients experiencing failures due to memory-based flow control
+ - Alert when `rate(heap_timeout_total[5m]) > 1`
+ - Indicates clients experiencing failures due to memory-based flow control
4. **P99 Wait Time Alert**:
- - Alert when `heap_wait_time_ms{quantile="0.99"} > 10000`
- - Indicates degraded client experience due to memory pressure
+ - Alert when `heap_wait_time_ms{quantile="0.99"} > 10000`
+ - Indicates degraded client experience due to memory pressure
These alerts should be configured alongside existing memory alerts for `managedLedgerCacheSizeMB`, `maxMessagePublishBufferSizeInMB`, and other memory limits to provide comprehensive memory management visibility.
@@ -473,30 +683,22 @@ The memory limiting mechanism introduces new denial-of-service protections:
2. **Fair Queueing**: The queue size limits prevent bad clients from monopolizing memory permits and blocking legitimate requests.
-3. **Multi-tenancy Isolation**: Consider per-tenant memory limits in future iterations to prevent one tenant from consuming all available topic listing memory permits, similar to how other memory limits could benefit from tenant isolation.
+3. **Cancellation Support**: The implementation includes cancellation support to prevent resource waste when clients disconnect, preventing malicious clients from queueing many requests and then disconnecting.
+
+4. **Multi-tenancy Isolation**: Consider per-tenant memory limits in future iterations to prevent one tenant from consuming all available topic listing memory permits, similar to how other memory limits could benefit from tenant isolation.
## Backward & Forward Compatibility
### Upgrade
-1. The feature can be disabled setting the limits set to 0 initially to maintain full compatibility
-2. After upgrade, gradually enable memory limits:
- ```bash
- # Start with high limits to understand current usage
- pulsar-admin brokers update-dynamic-config --config maxTopicListInFlightHeapMemSizeMB --value 512
- pulsar-admin brokers update-dynamic-config --config maxTopicListInFlightDirectMemSizeMB --value 512
-
- # Monitor metrics and adjust downward based on actual usage patterns
- pulsar-admin brokers update-dynamic-config --config maxTopicListInFlightHeapMemSizeMB --value 200
- pulsar-admin brokers update-dynamic-config --config maxTopicListInFlightDirectMemSizeMB --value 200
- pulsar-admin brokers update-dynamic-config --config maxTopicListInFlightHeapMemSizeMB --value 100
- pulsar-admin brokers update-dynamic-config --config maxTopicListInFlightDirectMemSizeMB --value 100
- ```
+1. The feature is enabled by default with reasonable limits (100MB each for heap and direct memory)
+2. After upgrade monitor metrics
3. No client changes are required
### Downgrade / Rollback
- No changes required
+- If downgrading, the memory limiting behavior will be lost but no functional issues will occur
### Pulsar Geo-Replication Upgrade & Downgrade/Rollback Considerations
diff --git a/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/ServiceConfiguration.java b/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/ServiceConfiguration.java
index e63c7d2b286..a3183316ccf 100644
--- a/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/ServiceConfiguration.java
+++ b/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/ServiceConfiguration.java
@@ -1391,6 +1391,42 @@ public class ServiceConfiguration implements PulsarConfiguration {
doc = "Max number of concurrent lookup request broker allows to throttle heavy incoming lookup traffic")
private int maxConcurrentLookupRequest = 50000;
+ @FieldContext(
+ category = CATEGORY_SERVER,
+ doc = "Maximum heap memory for inflight topic list operations (MB).\n"
+ + "Default: 100 MB (supports ~1M topic names assuming 100 bytes each)")
+ private int maxTopicListInFlightHeapMemSizeMB = 100;
+
+ @FieldContext(
+ category = CATEGORY_SERVER,
+ doc = "Maximum direct memory for inflight topic list responses (MB).\n"
+ + "Default: 100 MB (network buffers for serialized responses)")
+ private int maxTopicListInFlightDirectMemSizeMB = 100;
+
+ @FieldContext(
+ category = CATEGORY_SERVER,
+ doc = "Timeout for acquiring heap memory permits (milliseconds).\n"
+ + "Default: 25000 (25 seconds)")
+ private int maxTopicListInFlightHeapMemSizePermitsAcquireTimeoutMillis = 25000;
+
+ @FieldContext(
+ category = CATEGORY_SERVER,
+ doc = "Maximum queue size for heap memory permit requests.\n"
+ + "Default: 10000 (prevent unbounded queueing)")
+ private int maxTopicListInFlightHeapMemSizePermitsAcquireQueueSize = 10000;
+
+ @FieldContext(
+ category = CATEGORY_SERVER,
+ doc = "Timeout for acquiring direct memory permits (milliseconds).\n"
+ + "Default: 25000 (25 seconds)")
+ private int maxTopicListInFlightDirectMemSizePermitsAcquireTimeoutMillis = 25000;
+
+ @FieldContext(
+ category = CATEGORY_SERVER,
+ doc = "Maximum queue size for direct memory permit requests.\n"
+ + "Default: 10000 (prevent unbounded queueing)")
+ private int maxTopicListInFlightDirectMemSizePermitsAcquireQueueSize = 10000;
+
@FieldContext(
dynamic = true,
category = CATEGORY_SERVER,
diff --git a/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/topiclistlimit/TopicListMemoryLimiter.java b/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/topiclistlimit/TopicListMemoryLimiter.java
new file mode 100644
index 00000000000..42f8d473776
--- /dev/null
+++ b/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/topiclistlimit/TopicListMemoryLimiter.java
@@ -0,0 +1,312 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.pulsar.broker.topiclistlimit;
+
+import io.netty.buffer.ByteBufUtil;
+import io.opentelemetry.api.metrics.DoubleGauge;
+import io.opentelemetry.api.metrics.DoubleHistogram;
+import io.opentelemetry.api.metrics.LongCounter;
+import io.opentelemetry.api.metrics.LongGauge;
+import io.opentelemetry.api.metrics.Meter;
+import io.opentelemetry.api.metrics.ObservableDoubleGauge;
+import io.opentelemetry.api.metrics.ObservableLongUpDownCounter;
+import io.prometheus.client.Collector;
+import io.prometheus.client.CollectorRegistry;
+import io.prometheus.client.Counter;
+import io.prometheus.client.Gauge;
+import io.prometheus.client.Summary;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.pulsar.common.semaphore.AsyncDualMemoryLimiterImpl;
+import org.apache.pulsar.common.semaphore.AsyncSemaphore;
+
+/**
+ * Topic list memory limiter that exposes both Prometheus metrics and OpenTelemetry metrics.
+ */
+@Slf4j
+public class TopicListMemoryLimiter extends AsyncDualMemoryLimiterImpl {
+ private final CollectorRegistry collectorRegistry;
+ private final Gauge heapMemoryUsedBytes;
+ private final Gauge heapMemoryLimitBytes;
+ private final Gauge directMemoryUsedBytes;
+ private final Gauge directMemoryLimitBytes;
+ private final Gauge heapQueueSize;
+ private final Gauge heapQueueMaxSize;
+ private final Gauge directQueueSize;
+ private final Gauge directQueueMaxSize;
+ private final Summary heapWaitTimeMs;
+ private final Summary directWaitTimeMs;
+ private final Counter heapTimeoutTotal;
+ private final Counter directTimeoutTotal;
+ private final ObservableDoubleGauge otelHeapMemoryUsedGauge;
+ private final DoubleGauge otelHeapMemoryLimitGauge;
+ private final ObservableDoubleGauge otelDirectMemoryUsedGauge;
+ private final DoubleGauge otelDirectMemoryLimitGauge;
+ private final ObservableLongUpDownCounter otelHeapQueueSize;
+ private final ObservableLongUpDownCounter otelDirectQueueSize;
+ private final DoubleHistogram otelHeapWaitTime;
+ private final DoubleHistogram otelDirectWaitTime;
+ private final LongCounter otelHeapTimeoutTotal;
+ private final LongCounter otelDirectTimeoutTotal;
+
+ public TopicListMemoryLimiter(CollectorRegistry collectorRegistry, String prometheusPrefix,
+ Meter openTelemetryMeter,
+ long maxHeapMemory, int maxHeapQueueSize,
+ long heapTimeoutMillis, long maxDirectMemory, int maxDirectQueueSize,
+ long directTimeoutMillis) {
+ super(maxHeapMemory, maxHeapQueueSize, heapTimeoutMillis, maxDirectMemory, maxDirectQueueSize,
+ directTimeoutMillis);
+ this.collectorRegistry = collectorRegistry;
+
+ AsyncSemaphore heapMemoryLimiter = getLimiter(LimitType.HEAP_MEMORY);
+ AsyncSemaphore directMemoryLimiter = getLimiter(LimitType.DIRECT_MEMORY);
+
+ this.heapMemoryUsedBytes = register(Gauge.build(prometheusPrefix + "topic_list_heap_memory_used_bytes",
+ "Current heap memory used by topic listings")
+ .create()
+ .setChild(new Gauge.Child() {
+ @Override
+ public double get() {
+ return heapMemoryLimiter.getAcquiredPermits();
+ }
+ }));
+ this.otelHeapMemoryUsedGauge = openTelemetryMeter.gaugeBuilder("topic.list.heap.memory.used")
+ .setUnit("By")
+ .setDescription("Current heap memory used by topic listings")
+ .buildWithCallback(observableDoubleMeasurement -> {
+ observableDoubleMeasurement.record(heapMemoryLimiter.getAcquiredPermits());
+ });
+
+ this.heapMemoryLimitBytes = register(Gauge.build(prometheusPrefix + "topic_list_heap_memory_limit_bytes",
+ "Configured heap memory limit")
+ .create());
+ this.heapMemoryLimitBytes.set(maxHeapMemory);
+ this.otelHeapMemoryLimitGauge = openTelemetryMeter.gaugeBuilder("topic.list.heap.memory.limit")
+ .setUnit("By")
+ .setDescription("Configured heap memory limit")
+ .build();
+ this.otelHeapMemoryLimitGauge.set(maxHeapMemory);
+
+ this.directMemoryUsedBytes = register(Gauge.build(prometheusPrefix + "topic_list_direct_memory_used_bytes",
+ "Current direct memory used by topic listings")
+ .create()
+ .setChild(new Gauge.Child() {
+ @Override
+ public double get() {
+ return directMemoryLimiter.getAcquiredPermits();
+ }
+ }));
+ this.otelDirectMemoryUsedGauge = openTelemetryMeter.gaugeBuilder("topic.list.direct.memory.used")
+ .setUnit("By")
+ .setDescription("Current direct memory used by topic listings")
+ .buildWithCallback(observableDoubleMeasurement -> {
+ observableDoubleMeasurement.record(directMemoryLimiter.getAcquiredPermits());
+ });
+
+ this.directMemoryLimitBytes = register(Gauge.build(prometheusPrefix + "topic_list_direct_memory_limit_bytes",
+ "Configured direct memory limit")
+ .create());
+ this.directMemoryLimitBytes.set(maxDirectMemory);
+ this.otelDirectMemoryLimitGauge = openTelemetryMeter.gaugeBuilder("topic.list.direct.memory.limit")
+ .setUnit("By")
+ .setDescription("Configured direct memory limit")
+ .build();
+ this.otelDirectMemoryLimitGauge.set(maxDirectMemory);
+
+ this.heapQueueSize = register(Gauge.build(prometheusPrefix + "topic_list_heap_queue_size",
+ "Current heap memory limiter queue size")
+ .create()
+ .setChild(new Gauge.Child() {
+ @Override
+ public double get() {
+ return heapMemoryLimiter.getQueueSize();
+ }
+ }));
+ this.otelHeapQueueSize = openTelemetryMeter
+ .upDownCounterBuilder("topic.list.heap.queue.size")
+ .setDescription("Current heap memory limiter queue size")
+ .setUnit("1")
+ .buildWithCallback(observableLongMeasurement -> {
+ observableLongMeasurement.record(heapMemoryLimiter.getQueueSize());
+ });
+
+ this.heapQueueMaxSize = register(Gauge.build(prometheusPrefix + "topic_list_heap_queue_max_size",
+ "Maximum heap memory limiter queue size")
+ .create());
+ this.heapQueueMaxSize.set(maxHeapQueueSize);
+ LongGauge otelHeapQueueMaxSize = openTelemetryMeter
+ .gaugeBuilder("topic.list.heap.queue.max.size")
+ .setDescription("Maximum heap memory limiter queue size")
+ .setUnit("1")
+ .ofLongs()
+ .build();
+ otelHeapQueueMaxSize.set(maxHeapQueueSize);
+
+ this.directQueueSize = register(Gauge.build(prometheusPrefix + "topic_list_direct_queue_size",
+ "Current direct memory limiter queue size")
+ .create()
+ .setChild(new Gauge.Child() {
+ @Override
+ public double get() {
+ return directMemoryLimiter.getQueueSize();
+ }
+ }));
+ this.otelDirectQueueSize = openTelemetryMeter
+ .upDownCounterBuilder("topic.list.direct.queue.size")
+ .setDescription("Current direct memory limiter queue size")
+ .setUnit("1")
+ .buildWithCallback(observableLongMeasurement -> {
+ observableLongMeasurement.record(directMemoryLimiter.getQueueSize());
+ });
+
+ this.directQueueMaxSize = register(Gauge.build(prometheusPrefix + "topic_list_direct_queue_max_size",
+ "Maximum direct memory limiter queue size")
+ .create());
+ this.directQueueMaxSize.set(maxDirectQueueSize);
+ LongGauge otelDirectQueueMaxSize = openTelemetryMeter
+ .gaugeBuilder("topic.list.direct.queue.max.size")
+ .setDescription("Maximum direct memory limiter queue size")
+ .setUnit("1")
+ .ofLongs()
+ .build();
+ otelDirectQueueMaxSize.set(maxDirectQueueSize);
+
+ this.heapWaitTimeMs = register(Summary.build(prometheusPrefix + "topic_list_heap_wait_time_ms",
+ "Wait time for heap memory permits")
+ .quantile(0.50, 0.01)
+ .quantile(0.95, 0.01)
+ .quantile(0.99, 0.01)
+ .quantile(1, 0.01)
+ .create());
+ this.otelHeapWaitTime = openTelemetryMeter.histogramBuilder("topic.list.heap.wait.time")
+ .setUnit("s")
+ .setDescription("Wait time for heap memory permits")
+ .build();
+
+ this.directWaitTimeMs = register(Summary.build(prometheusPrefix + "topic_list_direct_wait_time_ms",
+ "Wait time for direct memory permits")
+ .quantile(0.50, 0.01)
+ .quantile(0.95, 0.01)
+ .quantile(0.99, 0.01)
+ .quantile(1, 0.01)
+ .create());
+ this.otelDirectWaitTime = openTelemetryMeter.histogramBuilder("topic.list.direct.wait.time")
+ .setUnit("s")
+ .setDescription("Wait time for direct memory permits")
+ .build();
+
+ this.heapTimeoutTotal = register(Counter.build(prometheusPrefix + "topic_list_heap_timeout_total",
+ "Total heap memory permit timeouts")
+ .create());
+ this.otelHeapTimeoutTotal = openTelemetryMeter.counterBuilder("topic.list.heap.timeout.total")
+ .setDescription("Total heap memory permit timeouts")
+ .setUnit("1")
+ .build();
+
+ this.directTimeoutTotal = register(Counter.build(prometheusPrefix + "topic_list_direct_timeout_total",
+ "Total direct memory permit timeouts")
+ .create());
+ this.otelDirectTimeoutTotal = openTelemetryMeter.counterBuilder("topic.list.direct.timeout.total")
+ .setDescription("Total direct memory permit timeouts")
+ .setUnit("1")
+ .build();
+ }
+
+ private <T extends Collector> T register(T collector) {
+ try {
+ collectorRegistry.register(collector);
+ } catch (Exception e) {
+ // ignore exception when registering a collector that is already registered
+ if (log.isDebugEnabled()) {
+ log.debug("Failed to register Prometheus collector {}", collector, e);
+ }
+ }
+ return collector;
+ }
+
+ private void unregister(Collector collector) {
+ try {
+ collectorRegistry.unregister(collector);
+ } catch (Exception e) {
+ // ignore exception when unregistering a collector that is not registered
+ if (log.isDebugEnabled()) {
+ log.debug("Failed to unregister Prometheus collector {}", collector, e);
+ }
+ }
+ }
+
+
+ @Override
+ protected void recordHeapWaitTime(long waitTimeNanos) {
+ if (waitTimeNanos == Long.MAX_VALUE) {
+ heapTimeoutTotal.inc();
+ otelHeapTimeoutTotal.add(1);
+ } else {
+ heapWaitTimeMs.observe(TimeUnit.NANOSECONDS.toMillis(waitTimeNanos));
+ otelHeapWaitTime.record(waitTimeNanos / 1_000_000_000.0d);
+ }
+ }
+
+ @Override
+ protected void recordDirectWaitTime(long waitTimeNanos) {
+ if (waitTimeNanos == Long.MAX_VALUE) {
+ directTimeoutTotal.inc();
+ otelDirectTimeoutTotal.add(1);
+ } else {
+ directWaitTimeMs.observe(TimeUnit.NANOSECONDS.toMillis(waitTimeNanos));
+ otelDirectWaitTime.record(waitTimeNanos / 1_000_000_000.0d);
+ }
+ }
+
+ @Override
+ public void close() {
+ super.close();
+ unregister(heapMemoryUsedBytes);
+ unregister(heapMemoryLimitBytes);
+ unregister(directMemoryUsedBytes);
+ unregister(directMemoryLimitBytes);
+ unregister(heapQueueSize);
+ unregister(heapQueueMaxSize);
+ unregister(directQueueSize);
+ unregister(directQueueMaxSize);
+ unregister(heapWaitTimeMs);
+ unregister(directWaitTimeMs);
+ unregister(heapTimeoutTotal);
+ unregister(directTimeoutTotal);
+ otelHeapMemoryUsedGauge.close();
+ otelDirectMemoryUsedGauge.close();
+ otelHeapQueueSize.close();
+ otelDirectQueueSize.close();
+ }
+
+ /**
+ * Estimate the heap memory size of a topic list.
+ * @param topicList the topic list to estimate
+ * @return the estimated heap memory size in bytes
+ */
+ public static long estimateTopicListSize(List<String> topicList) {
+ return topicList.stream()
+ .mapToLong(ByteBufUtil::utf8Bytes) // convert character count to bytes
+ // add 32 bytes overhead for each entry
+ // 16 bytes for object header, 16 bytes for java.lang.String fields
+ .map(n -> n + 32)
+ .sum();
+ }
+}
diff --git a/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/topiclistlimit/TopicListSizeResultCache.java b/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/topiclistlimit/TopicListSizeResultCache.java
new file mode 100644
index 00000000000..90e34844790
--- /dev/null
+++ b/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/topiclistlimit/TopicListSizeResultCache.java
@@ -0,0 +1,122 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.pulsar.broker.topiclistlimit;
+
+import com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.Caffeine;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicReference;
+import org.apache.pulsar.common.api.proto.CommandGetTopicsOfNamespace;
+
+/**
+ * A cache for storing and managing topic list size estimates in namespaces.
+ * This class provides functionality to:
+ * - Cache and retrieve topic list size estimates for different namespaces and modes
+ * - Handle concurrent requests for topic list sizes efficiently
+ * - Maintain and update size estimates based on actual topic list sizes
+ * - Prevent thundering herd problems when multiple concurrent requests for a namespace are made without
+ * a previous size estimate
+ * The cache uses namespace name and topic list mode (PERSISTENT/NON_PERSISTENT/ALL) as keys
+ * and maintains size estimates that are refined with actual usage.
+ */
+public class TopicListSizeResultCache {
+ // 10kB initial estimate for topic list heap size
+ private static final long INITIAL_TOPIC_LIST_HEAP_SIZE = 10 * 1024;
+
+ private Cache<CacheKey, ResultHolder> topicListSizeCache = Caffeine.newBuilder()
+ .maximumSize(10000)
+ .build();
+
+ record CacheKey(String namespaceName, CommandGetTopicsOfNamespace.Mode mode) {
+ }
+
+ /**
+ * Holds the topic list size estimate and future for the topic list size.
+ * The size is returned by calling {@link #getSizeAsync()} method which is asynchronous.
+ * The estimate is updated by calling {@link #updateSize(long)} method.
+ */
+ public static class ResultHolder {
+ private final AtomicReference<CompletableFuture<Long>> topicListSizeFuture =
+ new AtomicReference<>(null);
+
+ /**
+ * Get the topic list size estimate. The first request will return the initial estimate
+ * and update the estimate based on the returned size of the topic list. Other concurrent requests
+ * will wait for the first request to complete and use the estimate of the first request.
+ * Subsequent requests will use estimate which gets updated based on the returned size of the topic list
+ * of each request.
+ * @return a future that will return the topic list size estimate
+ */
+ public CompletableFuture<Long> getSizeAsync() {
+ if (topicListSizeFuture.compareAndSet(null, new CompletableFuture<>())) {
+ // let the first request proceed with the initial estimate
+ return CompletableFuture.completedFuture(INITIAL_TOPIC_LIST_HEAP_SIZE);
+ } else {
+ // all other requests wait for the first one to complete
+ return topicListSizeFuture.get();
+ }
+ }
+
+ /**
+ * Update the topic list size estimate. The last changed value will be used.
+ *
+ * @param actualSize the actual size of the topic list
+ */
+ public void updateSize(long actualSize) {
+ CompletableFuture<Long> currentFuture = topicListSizeFuture.get();
+ if (currentFuture != null && !currentFuture.isDone()) {
+ // complete the future if it's not done yet
+ currentFuture.complete(actualSize);
+ } else if (currentFuture == null || currentFuture.getNow(0L).longValue() != actualSize) {
+ // only update the future if the current value is different from the existing value
+ topicListSizeFuture.compareAndSet(currentFuture, CompletableFuture.completedFuture(actualSize));
+ }
+ }
+
+ /**
+ * After errors, it's necessary to call this method to ensure that the instance isn't left in a state
+ * where concurrent requests are waiting for the first request to complete the future by calling updateSize.
+ */
+ public void resetIfInitializing() {
+ CompletableFuture<Long> currentFuture = topicListSizeFuture.getAndUpdate(value -> {
+ if (value != null && !value.isDone()) {
+ return null;
+ } else {
+ return value;
+ }
+ });
+ // let all current requests proceed with the initial estimate if this were to happen
+ if (currentFuture != null && !currentFuture.isDone()) {
+ currentFuture.complete(INITIAL_TOPIC_LIST_HEAP_SIZE);
+ }
+ }
+ }
+
+ /**
+ * Get the topic list size result holder for the given namespace and mode.
+ * @param namespaceName the namespace name in the format of "tenant/namespace"
+ * @param mode the mode of the topic list request (PERSISTENT, NON_PERSISTENT, ALL)
+ * @return the topic list size result holder
+ */
+ public ResultHolder getTopicListSize(String namespaceName,
+ CommandGetTopicsOfNamespace.Mode mode) {
+ return topicListSizeCache.asMap()
+ .computeIfAbsent(new CacheKey(namespaceName, mode), __ -> new ResultHolder());
+ }
+}
diff --git a/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/topiclistlimit/package-info.java b/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/topiclistlimit/package-info.java
new file mode 100644
index 00000000000..42911acf708
--- /dev/null
+++ b/pulsar-broker-common/src/main/java/org/apache/pulsar/broker/topiclistlimit/package-info.java
@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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.
+ */
+/**
+ * Topic Listing memory limiter, PIP-442
+ */
+package org.apache.pulsar.broker.topiclistlimit;
diff --git a/pulsar-broker-common/src/test/java/org/apache/pulsar/broker/stats/prometheus/PrometheusMetricsClient.java b/pulsar-broker-common/src/test/java/org/apache/pulsar/broker/stats/prometheus/PrometheusMetricsClient.java
index cf6d63e96ce..298cbc2f658 100644
--- a/pulsar-broker-common/src/test/java/org/apache/pulsar/broker/stats/prometheus/PrometheusMetricsClient.java
+++ b/pulsar-broker-common/src/test/java/org/apache/pulsar/broker/stats/prometheus/PrometheusMetricsClient.java
@@ -31,6 +31,7 @@ import java.util.Map;
import java.util.TreeMap;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
+import lombok.Getter;
import org.apache.commons.lang3.tuple.Pair;
public class PrometheusMetricsClient {
@@ -110,6 +111,7 @@ public class PrometheusMetricsClient {
}
public static class Metrics {
+ @Getter
final Multimap<String, Metric> nameToDataPoints;
public Metrics(Multimap<String, Metric> nameToDataPoints) {
diff --git a/pulsar-broker-common/src/test/java/org/apache/pulsar/broker/topiclistlimit/TopicListSizeResultCacheTest.java b/pulsar-broker-common/src/test/java/org/apache/pulsar/broker/topiclistlimit/TopicListSizeResultCacheTest.java
new file mode 100644
index 00000000000..cea075886ad
--- /dev/null
+++ b/pulsar-broker-common/src/test/java/org/apache/pulsar/broker/topiclistlimit/TopicListSizeResultCacheTest.java
@@ -0,0 +1,321 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.pulsar.broker.topiclistlimit;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotEquals;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertNotSame;
+import static org.testng.Assert.assertSame;
+import static org.testng.Assert.assertTrue;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import org.apache.pulsar.common.api.proto.CommandGetTopicsOfNamespace;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+public class TopicListSizeResultCacheTest {
+
+ private TopicListSizeResultCache cache;
+ private ExecutorService executorService;
+
+ @BeforeMethod
+ public void setup() {
+ cache = new TopicListSizeResultCache();
+ executorService = Executors.newFixedThreadPool(10);
+ }
+
+ @AfterMethod(alwaysRun = true)
+ public void teardown() {
+ if (executorService != null) {
+ executorService.shutdownNow();
+ }
+ }
+
+ @Test
+ public void testGetTopicListSize_returnsSameInstanceForSameKey() {
+ String namespace = "tenant/namespace";
+ CommandGetTopicsOfNamespace.Mode mode = CommandGetTopicsOfNamespace.Mode.ALL;
+
+ TopicListSizeResultCache.ResultHolder holder1 = cache.getTopicListSize(namespace, mode);
+ TopicListSizeResultCache.ResultHolder holder2 = cache.getTopicListSize(namespace, mode);
+
+ assertSame(holder1, holder2, "Should return the same ResultHolder instance for the same key");
+ }
+
+ @Test
+ public void testGetTopicListSize_returnsDifferentInstancesForDifferentKeys() {
+ String namespace1 = "tenant/namespace1";
+ String namespace2 = "tenant/namespace2";
+ CommandGetTopicsOfNamespace.Mode mode = CommandGetTopicsOfNamespace.Mode.ALL;
+
+ TopicListSizeResultCache.ResultHolder holder1 = cache.getTopicListSize(namespace1, mode);
+ TopicListSizeResultCache.ResultHolder holder2 = cache.getTopicListSize(namespace2, mode);
+
+ assertNotSame(holder1, holder2, "Should return different ResultHolder instances for different namespaces");
+ }
+
+ @Test
+ public void testGetTopicListSize_returnsDifferentInstancesForDifferentModes() {
+ String namespace = "tenant/namespace";
+
+ TopicListSizeResultCache.ResultHolder holder1 = cache.getTopicListSize(namespace,
+ CommandGetTopicsOfNamespace.Mode.ALL);
+ TopicListSizeResultCache.ResultHolder holder2 = cache.getTopicListSize(namespace,
+ CommandGetTopicsOfNamespace.Mode.PERSISTENT);
+
+ assertNotSame(holder1, holder2, "Should return different ResultHolder instances for different modes");
+ }
+
+ @Test
+ public void testResultHolder_firstRequestReturnsInitialEstimate() throws Exception {
+ TopicListSizeResultCache.ResultHolder holder = new TopicListSizeResultCache.ResultHolder();
+
+ CompletableFuture<Long> sizeFuture = holder.getSizeAsync();
+
+ assertTrue(sizeFuture.isDone(), "First request should complete immediately");
+ assertEquals(sizeFuture.get().longValue(), 10 * 1024L,
+ "First request should return initial estimate of 10KB");
+ }
+
+ @Test
+ public void testResultHolder_concurrentRequestsWaitForFirstToComplete() throws Exception {
+ TopicListSizeResultCache.ResultHolder holder = new TopicListSizeResultCache.ResultHolder();
+ int numConcurrentRequests = 5;
+ CountDownLatch startLatch = new CountDownLatch(1);
+ CountDownLatch completeLatch = new CountDownLatch(numConcurrentRequests);
+
+ List<CompletableFuture<Long>> futures = new ArrayList<>();
+
+ // Start concurrent requests
+ for (int i = 0; i < numConcurrentRequests; i++) {
+ executorService.submit(() -> {
+ try {
+ startLatch.await();
+ CompletableFuture<Long> future = holder.getSizeAsync();
+ futures.add(future);
+ completeLatch.countDown();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ });
+ }
+
+ // Trigger all requests
+ startLatch.countDown();
+ completeLatch.await(5, TimeUnit.SECONDS);
+
+ // First request completes immediately with initial estimate
+ assertTrue(futures.get(0).isDone());
+ assertEquals(futures.get(0).get().longValue(), 10 * 1024L);
+
+ // Other requests should be waiting
+ for (int i = 1; i < numConcurrentRequests; i++) {
+ assertFalse(futures.get(i).isDone(), "Concurrent request " + i + " should be waiting");
+ }
+
+ // Update size to complete waiting requests
+ long actualSize = 20 * 1024L;
+ holder.updateSize(actualSize);
+
+ // All waiting requests should now complete with the actual size
+ for (int i = 1; i < numConcurrentRequests; i++) {
+ assertTrue(futures.get(i).isDone(), "Request " + i + " should complete after updateSize");
+ assertEquals(futures.get(i).get().longValue(), actualSize);
+ }
+ }
+
+ @Test
+ public void testResultHolder_updateSize_firstUpdate() {
+ TopicListSizeResultCache.ResultHolder holder = new TopicListSizeResultCache.ResultHolder();
+
+ // Get initial size
+ holder.getSizeAsync();
+
+ // Update with actual size
+ long actualSize = 20 * 1024L;
+ holder.updateSize(actualSize);
+
+ // Next request should return the actual size
+ CompletableFuture<Long> nextFuture = holder.getSizeAsync();
+ assertEquals(nextFuture.join().longValue(), actualSize);
+ }
+
+ @Test
+ public void testResultHolder_updateSize_completesWaitingFuture() throws Exception {
+ TopicListSizeResultCache.ResultHolder holder = new TopicListSizeResultCache.ResultHolder();
+
+ // First request gets initial estimate
+ CompletableFuture<Long> first = holder.getSizeAsync();
+ assertTrue(first.isDone());
+
+ // Second request should wait
+ CompletableFuture<Long> second = holder.getSizeAsync();
+ assertFalse(second.isDone(), "Second request should be waiting");
+
+ // Update size should complete the waiting future
+ long actualSize = 15000L;
+ holder.updateSize(actualSize);
+
+ assertTrue(second.isDone(), "Update should complete waiting future");
+ assertEquals(second.get().longValue(), actualSize);
+ }
+
+ @Test
+ public void testResultHolder_resetIfInitializing_resetsUncompletedFuture() throws Exception {
+ TopicListSizeResultCache.ResultHolder holder = new TopicListSizeResultCache.ResultHolder();
+
+ // First request gets initial estimate
+ holder.getSizeAsync();
+
+ // Second request should wait
+ CompletableFuture<Long> waitingFuture = holder.getSizeAsync();
+ assertFalse(waitingFuture.isDone(), "Second request should be waiting");
+
+ // Reset should complete the waiting future with initial estimate
+ holder.resetIfInitializing();
+
+ assertTrue(waitingFuture.isDone(), "Reset should complete waiting future");
+ assertEquals(waitingFuture.get().longValue(), 10 * 1024L,
+ "Reset should complete with initial estimate");
+
+ // Next request should start fresh
+ CompletableFuture<Long> afterReset = holder.getSizeAsync();
+ assertTrue(afterReset.isDone(), "Request after reset should complete immediately");
+ assertEquals(afterReset.get().longValue(), 10 * 1024L);
+ }
+
+ @Test
+ public void testResultHolder_resetIfInitializing_noEffectOnCompletedFuture() throws Exception {
+ TopicListSizeResultCache.ResultHolder holder = new TopicListSizeResultCache.ResultHolder();
+
+ // Get initial size and update
+ holder.getSizeAsync();
+ holder.updateSize(20000L);
+
+ CompletableFuture<Long> future = holder.getSizeAsync();
+ assertTrue(future.isDone());
+ long valueBefore = future.get();
+
+ // Reset should have no effect on completed future
+ holder.resetIfInitializing();
+
+ CompletableFuture<Long> futureAfter = holder.getSizeAsync();
+ assertEquals(futureAfter.join().longValue(), valueBefore,
+ "Reset should not affect already completed future");
+ }
+
+ @Test
+ public void testResultHolder_multipleSequentialRequests() throws Exception {
+ TopicListSizeResultCache.ResultHolder holder = new TopicListSizeResultCache.ResultHolder();
+
+ // First request
+ CompletableFuture<Long> future1 = holder.getSizeAsync();
+ assertEquals(future1.get().longValue(), 10 * 1024L);
+
+ // Update
+ holder.updateSize(15000L);
+
+ // Second request
+ CompletableFuture<Long> future2 = holder.getSizeAsync();
+ assertEquals(future2.get().longValue(), 15000L);
+
+ // Update again
+ holder.updateSize(25000L);
+
+ // Third request
+ CompletableFuture<Long> future3 = holder.getSizeAsync();
+ assertEquals(future3.get().longValue(), 25000L, "Should be the last value");
+ }
+
+ @Test
+ public void testCacheKey_equalityForSameValues() {
+ TopicListSizeResultCache.CacheKey key1 = new TopicListSizeResultCache.CacheKey("tenant/ns",
+ CommandGetTopicsOfNamespace.Mode.ALL);
+ TopicListSizeResultCache.CacheKey key2 = new TopicListSizeResultCache.CacheKey("tenant/ns",
+ CommandGetTopicsOfNamespace.Mode.ALL);
+
+ assertEquals(key1, key2, "Keys with same values should be equal");
+ assertEquals(key1.hashCode(), key2.hashCode(), "Equal keys should have same hash code");
+ }
+
+ @Test
+ public void testCacheKey_inequalityForDifferentNamespaces() {
+ TopicListSizeResultCache.CacheKey key1 = new TopicListSizeResultCache.CacheKey("tenant/ns1",
+ CommandGetTopicsOfNamespace.Mode.ALL);
+ TopicListSizeResultCache.CacheKey key2 = new TopicListSizeResultCache.CacheKey("tenant/ns2",
+ CommandGetTopicsOfNamespace.Mode.ALL);
+
+ assertNotEquals(key1, key2, "Keys with different namespaces should not be equal");
+ }
+
+ @Test
+ public void testCacheKey_inequalityForDifferentModes() {
+ TopicListSizeResultCache.CacheKey key1 = new TopicListSizeResultCache.CacheKey("tenant/ns",
+ CommandGetTopicsOfNamespace.Mode.ALL);
+ TopicListSizeResultCache.CacheKey key2 = new TopicListSizeResultCache.CacheKey("tenant/ns",
+ CommandGetTopicsOfNamespace.Mode.PERSISTENT);
+
+ assertNotEquals(key1, key2, "Keys with different modes should not be equal");
+ }
+
+ @Test
+ public void testHighConcurrency() throws Exception {
+ TopicListSizeResultCache.ResultHolder holder = new TopicListSizeResultCache.ResultHolder();
+ int numThreads = 20;
+ CountDownLatch startLatch = new CountDownLatch(1);
+ CountDownLatch completeLatch = new CountDownLatch(numThreads);
+
+ List<CompletableFuture<Long>> futures = new ArrayList<>();
+
+ // Launch many concurrent requests
+ for (int i = 0; i < numThreads; i++) {
+ executorService.submit(() -> {
+ try {
+ startLatch.await();
+ CompletableFuture<Long> future = holder.getSizeAsync();
+ synchronized (futures) {
+ futures.add(future);
+ }
+ completeLatch.countDown();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ });
+ }
+
+ startLatch.countDown();
+ assertTrue(completeLatch.await(5, TimeUnit.SECONDS), "All requests should complete");
+
+ // Update to complete all waiting requests
+ holder.updateSize(30000L);
+
+ // Verify all futures complete successfully
+ for (CompletableFuture<Long> future : futures) {
+ assertNotNull(future.get(5, TimeUnit.SECONDS), "All futures should complete");
+ }
+ }
+}
\ No newline at end of file
diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/admin/impl/NamespacesBase.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/admin/impl/NamespacesBase.java
index 01659400e08..39718c3902e 100644
--- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/admin/impl/NamespacesBase.java
+++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/admin/impl/NamespacesBase.java
@@ -36,13 +36,17 @@ import java.util.Set;
import java.util.SortedSet;
import java.util.TreeSet;
import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.BooleanSupplier;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import javax.ws.rs.WebApplicationException;
import javax.ws.rs.container.AsyncResponse;
+import javax.ws.rs.container.CompletionCallback;
import javax.ws.rs.core.Response;
import javax.ws.rs.core.Response.Status;
import javax.ws.rs.core.UriBuilder;
@@ -62,6 +66,8 @@ import org.apache.pulsar.broker.service.Subscription;
import org.apache.pulsar.broker.service.Topic;
import org.apache.pulsar.broker.service.persistent.PersistentReplicator;
import org.apache.pulsar.broker.service.persistent.PersistentTopic;
+import org.apache.pulsar.broker.topiclistlimit.TopicListMemoryLimiter;
+import org.apache.pulsar.broker.topiclistlimit.TopicListSizeResultCache;
import org.apache.pulsar.broker.web.RestException;
import org.apache.pulsar.client.admin.GrantTopicPermissionOptions;
import org.apache.pulsar.client.admin.PulsarAdmin;
@@ -109,6 +115,8 @@ import org.apache.pulsar.common.policies.data.TopicType;
import org.apache.pulsar.common.policies.data.ValidateResult;
import org.apache.pulsar.common.policies.data.impl.AutoTopicCreationOverrideImpl;
import org.apache.pulsar.common.policies.data.impl.DispatchRateImpl;
+import org.apache.pulsar.common.semaphore.AsyncDualMemoryLimiter;
+import org.apache.pulsar.common.semaphore.AsyncDualMemoryLimiterImpl;
import org.apache.pulsar.common.util.Codec;
import org.apache.pulsar.common.util.FutureUtil;
import org.apache.pulsar.metadata.api.MetadataStoreException;
@@ -163,8 +171,66 @@ public abstract class NamespacesBase extends AdminResource {
.thenAccept(__ -> log.info("[{}] Created namespace {}", clientAppId(), namespaceName));
}
- protected CompletableFuture<List<String>> internalGetListOfTopics(Policies policies,
+ protected CompletableFuture<List<String>> internalGetListOfTopics(AsyncResponse response, Policies policies,
CommandGetTopicsOfNamespace.Mode mode) {
+ // Use maxTopicListInFlightLimiter to limit inflight get topic listing responses
+ // to avoid OOME caused by a lot of clients using HTTP service lookups to list topics
+ AsyncDualMemoryLimiterImpl maxTopicListInFlightLimiter =
+ pulsar().getBrokerService().getMaxTopicListInFlightLimiter();
+ TopicListSizeResultCache.ResultHolder listSizeHolder = pulsar().getBrokerService().getTopicListSizeResultCache()
+ .getTopicListSize(namespaceName.toString(), mode);
+ // setup the permit cancellation function
+ AtomicBoolean permitRequestCancelled = new AtomicBoolean(false);
+ BooleanSupplier isPermitRequestCancelled = permitRequestCancelled::get;
+ // add callback that releases permits when the response completes
+ AtomicReference<AsyncDualMemoryLimiter.AsyncDualMemoryLimiterPermit> initialPermitsRef =
+ new AtomicReference<>();
+ AtomicReference<AsyncDualMemoryLimiter.AsyncDualMemoryLimiterPermit> permitsRef = new AtomicReference<>();
+ response.register(new CompletionCallback() {
+ @Override
+ public void onComplete(Throwable throwable) {
+ if (throwable != null) {
+ // for failed request
+ // handle resetting the TopicListSizeResultCache.ResultHolder
+ listSizeHolder.resetIfInitializing();
+ // cancel any pending permit request
+ permitRequestCancelled.set(true);
+ }
+ AsyncDualMemoryLimiter.AsyncDualMemoryLimiterPermit initialPermit = initialPermitsRef.get();
+ if (initialPermit != null) {
+ maxTopicListInFlightLimiter.release(initialPermit);
+ }
+ AsyncDualMemoryLimiter.AsyncDualMemoryLimiterPermit permits = permitsRef.get();
+ if (permits != null) {
+ maxTopicListInFlightLimiter.release(permits);
+ }
+ }
+ });
+ return listSizeHolder.getSizeAsync().thenCompose(initialSize -> maxTopicListInFlightLimiter.acquire(initialSize,
+ AsyncDualMemoryLimiter.LimitType.HEAP_MEMORY, isPermitRequestCancelled).exceptionally(t -> {
+ throw new CompletionException(
+ new RestException(Status.TOO_MANY_REQUESTS, "Failed due to heap memory limit exceeded"));
+ }).thenCompose(initialPermits -> {
+ initialPermitsRef.set(initialPermits);
+ // perform the actual get list of topics operation
+ return doInternalGetListOfTopics(policies, mode).thenCompose(topicList -> {
+ long actualSize = TopicListMemoryLimiter.estimateTopicListSize(topicList);
+ listSizeHolder.updateSize(actualSize);
+ return maxTopicListInFlightLimiter.update(initialPermits, actualSize, isPermitRequestCancelled)
+ .exceptionally(t -> {
+ throw new CompletionException(new RestException(Status.TOO_MANY_REQUESTS,
+ "Failed due to heap memory limit exceeded"));
+ }).thenApply(permits -> {
+ permitsRef.set(permits);
+ initialPermitsRef.set(null);
+ return topicList;
+ });
+ });
+ }));
+ }
+
+ private CompletableFuture<List<String>> doInternalGetListOfTopics(Policies policies,
+ CommandGetTopicsOfNamespace.Mode mode) {
switch (mode) {
case ALL:
return pulsar().getNamespaceService().getListOfPersistentTopics(namespaceName)
diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/admin/v1/Namespaces.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/admin/v1/Namespaces.java
index 4424ceded81..ad1b223e8c2 100644
--- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/admin/v1/Namespaces.java
+++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/admin/v1/Namespaces.java
@@ -146,7 +146,7 @@ public class Namespaces extends NamespacesBase {
validateNamespaceOperationAsync(NamespaceName.get(property, namespace), NamespaceOperation.GET_TOPICS)
// Validate that namespace exists, throws 404 if it doesn't exist
.thenCompose(__ -> getNamespacePoliciesAsync(namespaceName))
- .thenCompose(policies -> internalGetListOfTopics(policies, mode))
+ .thenCompose(policies -> internalGetListOfTopics(response, policies, mode))
.thenApply(topics -> filterSystemTopic(topics, includeSystemTopic))
.thenAccept(response::resume)
.exceptionally(ex -> {
diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/admin/v2/Namespaces.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/admin/v2/Namespaces.java
index 3b045259f9a..5632de07f47 100644
--- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/admin/v2/Namespaces.java
+++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/admin/v2/Namespaces.java
@@ -131,7 +131,7 @@ public class Namespaces extends NamespacesBase {
validateNamespaceOperationAsync(NamespaceName.get(tenant, namespace), NamespaceOperation.GET_TOPICS)
// Validate that namespace exists, throws 404 if it doesn't exist
.thenCompose(__ -> getNamespacePoliciesAsync(namespaceName))
- .thenCompose(policies -> internalGetListOfTopics(policies, mode))
+ .thenCompose(policies -> internalGetListOfTopics(response, policies, mode))
.thenApply(topics -> filterSystemTopic(topics, includeSystemTopic))
.thenAccept(response::resume)
.exceptionally(ex -> {
diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/BrokerService.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/BrokerService.java
index 3323b9db713..e4b2f335023 100644
--- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/BrokerService.java
+++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/BrokerService.java
@@ -42,6 +42,7 @@ import io.netty.handler.ssl.SslContext;
import io.netty.util.concurrent.DefaultThreadFactory;
import io.opentelemetry.api.metrics.LongCounter;
import io.opentelemetry.api.metrics.ObservableLongUpDownCounter;
+import io.prometheus.client.CollectorRegistry;
import io.prometheus.client.Gauge;
import io.prometheus.client.Histogram;
import java.io.Closeable;
@@ -139,6 +140,8 @@ import org.apache.pulsar.broker.stats.prometheus.metrics.ObserverGauge;
import org.apache.pulsar.broker.stats.prometheus.metrics.Summary;
import org.apache.pulsar.broker.storage.ManagedLedgerStorage;
import org.apache.pulsar.broker.storage.ManagedLedgerStorageClass;
+import org.apache.pulsar.broker.topiclistlimit.TopicListMemoryLimiter;
+import org.apache.pulsar.broker.topiclistlimit.TopicListSizeResultCache;
import org.apache.pulsar.broker.validator.BindAddressValidator;
import org.apache.pulsar.client.admin.PulsarAdmin;
import org.apache.pulsar.client.admin.PulsarAdminBuilder;
@@ -180,6 +183,7 @@ import org.apache.pulsar.common.policies.data.TopicType;
import org.apache.pulsar.common.policies.data.impl.AutoSubscriptionCreationOverrideImpl;
import org.apache.pulsar.common.policies.data.stats.TopicStatsImpl;
import org.apache.pulsar.common.protocol.schema.SchemaVersion;
+import org.apache.pulsar.common.semaphore.AsyncDualMemoryLimiterImpl;
import org.apache.pulsar.common.stats.Metrics;
import org.apache.pulsar.common.util.FieldParser;
import org.apache.pulsar.common.util.FutureUtil;
@@ -250,6 +254,10 @@ public class BrokerService implements Closeable {
private final SingleThreadNonConcurrentFixedRateScheduler backlogQuotaChecker;
protected final AtomicReference<Semaphore> lookupRequestSemaphore;
+ @Getter
+ private final AsyncDualMemoryLimiterImpl maxTopicListInFlightLimiter;
+ @Getter
+ private final TopicListSizeResultCache topicListSizeResultCache = new TopicListSizeResultCache();
protected final AtomicReference<Semaphore> topicLoadRequestSemaphore;
public static final String TOPIC_LOOKUP_USAGE_METRIC_NAME = "pulsar.broker.request.topic.lookup.concurrent.usage";
@@ -387,6 +395,15 @@ public class BrokerService implements Closeable {
updateConfigurationAndRegisterListeners();
this.lookupRequestSemaphore = new AtomicReference<>(
new Semaphore(pulsar.getConfiguration().getMaxConcurrentLookupRequest(), false));
+ // Initialize topic list memory limiter
+ this.maxTopicListInFlightLimiter = new TopicListMemoryLimiter(
+ CollectorRegistry.defaultRegistry, "pulsar_broker_", pulsar.getOpenTelemetry().getMeter(),
+ pulsar.getConfiguration().getMaxTopicListInFlightHeapMemSizeMB() * 1024L * 1024L,
+ pulsar.getConfiguration().getMaxTopicListInFlightHeapMemSizePermitsAcquireQueueSize(),
+ pulsar.getConfiguration().getMaxTopicListInFlightHeapMemSizePermitsAcquireTimeoutMillis(),
+ pulsar.getConfiguration().getMaxTopicListInFlightDirectMemSizeMB() * 1024L * 1024L,
+ pulsar.getConfiguration().getMaxTopicListInFlightDirectMemSizePermitsAcquireQueueSize(),
+ pulsar.getConfiguration().getMaxTopicListInFlightDirectMemSizePermitsAcquireTimeoutMillis());
this.topicLoadRequestSemaphore = new AtomicReference<>(
new Semaphore(pulsar.getConfiguration().getMaxConcurrentTopicLoadRequest(), false));
if (pulsar.getConfiguration().getMaxUnackedMessagesPerBroker() > 0
@@ -890,6 +907,8 @@ public class BrokerService implements Closeable {
}
});
+ maxTopicListInFlightLimiter.close();
+
if (interceptor != null) {
interceptor.close();
interceptor = null;
diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/PulsarCommandSender.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/PulsarCommandSender.java
index f25703757d3..8e2deff31d0 100644
--- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/PulsarCommandSender.java
+++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/PulsarCommandSender.java
@@ -22,6 +22,8 @@ package org.apache.pulsar.broker.service;
import io.netty.util.concurrent.Future;
import java.util.List;
import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+import java.util.function.Consumer;
import org.apache.bookkeeper.mledger.Entry;
import org.apache.pulsar.client.api.transaction.TxnID;
import org.apache.pulsar.common.api.proto.CommandLookupTopicResponse;
@@ -51,8 +53,9 @@ public interface PulsarCommandSender {
void sendSendError(long producerId, long sequenceId, ServerError error, String errorMsg);
- void sendGetTopicsOfNamespaceResponse(List<String> topics, String topicsHash, boolean filtered,
- boolean changed, long requestId);
+ CompletableFuture<Void> sendGetTopicsOfNamespaceResponse(List<String> topics, String topicsHash, boolean filtered,
+ boolean changed, long requestId,
+ Consumer<Throwable> permitAcquireErrorHandler);
void sendGetSchemaResponse(long requestId, SchemaInfo schema, SchemaVersion version);
diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/PulsarCommandSenderImpl.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/PulsarCommandSenderImpl.java
index 105650caaaf..275e049255c 100644
--- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/PulsarCommandSenderImpl.java
+++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/PulsarCommandSenderImpl.java
@@ -18,6 +18,7 @@
*/
package org.apache.pulsar.broker.service;
+import static org.apache.pulsar.common.semaphore.AsyncDualMemoryLimiterUtil.acquireDirectMemoryPermitsAndWriteAndFlush;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.channel.ChannelHandlerContext;
@@ -25,6 +26,8 @@ import io.netty.channel.ChannelPromise;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
+import java.util.function.Consumer;
import lombok.extern.slf4j.Slf4j;
import org.apache.bookkeeper.mledger.Entry;
import org.apache.pulsar.broker.intercept.BrokerInterceptor;
@@ -37,6 +40,7 @@ import org.apache.pulsar.common.api.proto.ServerError;
import org.apache.pulsar.common.protocol.Commands;
import org.apache.pulsar.common.protocol.schema.SchemaVersion;
import org.apache.pulsar.common.schema.SchemaInfo;
+import org.apache.pulsar.common.semaphore.AsyncDualMemoryLimiter;
import org.apache.pulsar.common.util.netty.NettyChannelUtil;
@Slf4j
@@ -44,10 +48,13 @@ public class PulsarCommandSenderImpl implements PulsarCommandSender {
private final BrokerInterceptor interceptor;
private final ServerCnx cnx;
+ private final AsyncDualMemoryLimiter maxTopicListInFlightLimiter;
- public PulsarCommandSenderImpl(BrokerInterceptor interceptor, ServerCnx cnx) {
+ public PulsarCommandSenderImpl(BrokerInterceptor interceptor, ServerCnx cnx,
+ AsyncDualMemoryLimiter maxTopicListInFlightLimiter) {
this.interceptor = interceptor;
this.cnx = cnx;
+ this.maxTopicListInFlightLimiter = maxTopicListInFlightLimiter;
}
@Override
@@ -120,13 +127,14 @@ public class PulsarCommandSenderImpl implements PulsarCommandSender {
}
@Override
- public void sendGetTopicsOfNamespaceResponse(List<String> topics, String topicsHash,
- boolean filtered, boolean changed, long requestId) {
+ public CompletableFuture<Void> sendGetTopicsOfNamespaceResponse(List<String> topics, String topicsHash,
+ boolean filtered, boolean changed, long requestId,
+ Consumer<Throwable> permitAcquireErrorHandler) {
BaseCommand command = Commands.newGetTopicsOfNamespaceResponseCommand(topics, topicsHash,
filtered, changed, requestId);
safeIntercept(command, cnx);
- ByteBuf outBuf = Commands.serializeWithSize(command);
- writeAndFlush(outBuf);
+ return acquireDirectMemoryPermitsAndWriteAndFlush(cnx.ctx(), maxTopicListInFlightLimiter, () -> !cnx.isActive(),
+ command, permitAcquireErrorHandler);
}
@Override
diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ServerCnx.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ServerCnx.java
index d7010e3cf8c..55cd0fe7e3d 100644
--- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ServerCnx.java
+++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ServerCnx.java
@@ -61,6 +61,7 @@ import java.util.concurrent.CompletionStage;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
+import java.util.function.BooleanSupplier;
import java.util.stream.Collectors;
import javax.naming.AuthenticationException;
import javax.net.ssl.SSLSession;
@@ -99,6 +100,8 @@ import org.apache.pulsar.broker.service.persistent.PersistentSubscription;
import org.apache.pulsar.broker.service.persistent.PersistentTopic;
import org.apache.pulsar.broker.service.schema.SchemaRegistryService;
import org.apache.pulsar.broker.service.schema.exceptions.IncompatibleSchemaException;
+import org.apache.pulsar.broker.topiclistlimit.TopicListMemoryLimiter;
+import org.apache.pulsar.broker.topiclistlimit.TopicListSizeResultCache;
import org.apache.pulsar.broker.web.RestException;
import org.apache.pulsar.client.api.PulsarClientException;
import org.apache.pulsar.client.api.transaction.TxnID;
@@ -172,6 +175,8 @@ import org.apache.pulsar.common.protocol.PulsarHandler;
import org.apache.pulsar.common.protocol.schema.SchemaData;
import org.apache.pulsar.common.protocol.schema.SchemaVersion;
import org.apache.pulsar.common.schema.SchemaType;
+import org.apache.pulsar.common.semaphore.AsyncDualMemoryLimiter;
+import org.apache.pulsar.common.semaphore.AsyncDualMemoryLimiterImpl;
import org.apache.pulsar.common.topics.TopicList;
import org.apache.pulsar.common.topics.TopicsPattern;
import org.apache.pulsar.common.util.FutureUtil;
@@ -260,6 +265,7 @@ public class ServerCnx extends PulsarHandler implements TransportCnx {
private final TimedSingleThreadRateLimiter requestRateLimiter;
private final int pauseReceivingCooldownMilliSeconds;
private boolean pausedDueToRateLimitation = false;
+ private AsyncDualMemoryLimiterImpl maxTopicListInFlightLimiter;
// Tracks and limits number of bytes pending to be published from a single specific IO thread.
static final class PendingBytesPerThreadTracker {
@@ -367,6 +373,7 @@ public class ServerCnx extends PulsarHandler implements TransportCnx {
this.connectionController = new ConnectionController.DefaultConnectionController(
conf.getBrokerMaxConnections(),
conf.getBrokerMaxConnectionsPerIp());
+ this.maxTopicListInFlightLimiter = pulsar.getBrokerService().getMaxTopicListInFlightLimiter();
this.enableSubscriptionPatternEvaluation = conf.isEnableBrokerSideSubscriptionPatternEvaluation();
this.maxSubscriptionPatternLength = conf.getSubscriptionPatternMaxLength();
this.topicListService = new TopicListService(pulsar, this,
@@ -395,7 +402,8 @@ public class ServerCnx extends PulsarHandler implements TransportCnx {
log.debug("New connection from {}", remoteAddress);
}
this.ctx = ctx;
- this.commandSender = new PulsarCommandSenderImpl(brokerInterceptor, this);
+ this.commandSender =
+ new PulsarCommandSenderImpl(brokerInterceptor, this, this.service.getMaxTopicListInFlightLimiter());
this.service.getPulsarStats().recordConnectionCreate();
cnxsPerThread.get().add(this);
service.getPulsar().runWhenReadyForIncomingRequests(() -> {
@@ -2546,45 +2554,8 @@ public class ServerCnx extends PulsarHandler implements TransportCnx {
if (lookupSemaphore.tryAcquire()) {
isNamespaceOperationAllowed(namespaceName, NamespaceOperation.GET_TOPICS).thenApply(isAuthorized -> {
if (isAuthorized) {
- getBrokerService().pulsar().getNamespaceService().getListOfUserTopics(namespaceName, mode)
- .thenAccept(topics -> {
- boolean filterTopics = false;
- // filter system topic
- List<String> filteredTopics = topics;
-
- if (enableSubscriptionPatternEvaluation && topicsPattern.isPresent()) {
- if (topicsPattern.get().length() <= maxSubscriptionPatternLength) {
- filterTopics = true;
- filteredTopics = TopicList.filterTopics(filteredTopics, topicsPattern.get(),
- topicsPatternImplementation);
- } else {
- log.info("[{}] Subscription pattern provided [{}] was longer than maximum {}.",
- remoteAddress, topicsPattern.get(), maxSubscriptionPatternLength);
- }
- }
- String hash = TopicList.calculateHash(filteredTopics);
- boolean hashUnchanged = topicsHash.isPresent() && topicsHash.get().equals(hash);
- if (hashUnchanged) {
- filteredTopics = Collections.emptyList();
- }
- if (log.isDebugEnabled()) {
- log.debug(
- "[{}] Received CommandGetTopicsOfNamespace for namespace [//{}] by {}, size:{}",
- remoteAddress, namespace, requestId, topics.size());
- }
- commandSender.sendGetTopicsOfNamespaceResponse(filteredTopics, hash, filterTopics,
- !hashUnchanged, requestId);
- lookupSemaphore.release();
- })
- .exceptionally(ex -> {
- log.warn("[{}] Error GetTopicsOfNamespace for namespace [//{}] by {}",
- remoteAddress, namespace, requestId);
- commandSender.sendErrorResponse(requestId,
- BrokerServiceException.getClientErrorCode(new ServerMetadataException(ex)),
- ex.getMessage());
- lookupSemaphore.release();
- return null;
- });
+ internalHandleGetTopicsOfNamespace(namespace, namespaceName, requestId, mode, topicsPattern,
+ topicsHash, lookupSemaphore);
} else {
final String msg = "Client is not authorized to GetTopicsOfNamespace";
log.warn("[{}] {} with role {} on namespace {}", remoteAddress, msg, getPrincipal(), namespaceName);
@@ -2610,6 +2581,93 @@ public class ServerCnx extends PulsarHandler implements TransportCnx {
}
}
+ private void internalHandleGetTopicsOfNamespace(String namespace, NamespaceName namespaceName, long requestId,
+ CommandGetTopicsOfNamespace.Mode mode,
+ Optional<String> topicsPattern, Optional<String> topicsHash,
+ Semaphore lookupSemaphore) {
+ BooleanSupplier isPermitRequestCancelled = () -> !ctx().channel().isActive();
+ TopicListSizeResultCache.ResultHolder
+ listSizeHolder = service.getTopicListSizeResultCache().getTopicListSize(namespaceName.toString(), mode);
+ listSizeHolder.getSizeAsync().thenAccept(initialSize -> {
+ maxTopicListInFlightLimiter.withAcquiredPermits(initialSize,
+ AsyncDualMemoryLimiter.LimitType.HEAP_MEMORY, isPermitRequestCancelled, initialPermits -> {
+ return getBrokerService().pulsar().getNamespaceService()
+ .getListOfUserTopics(namespaceName, mode)
+ .thenAccept(topics -> {
+ long actualSize = TopicListMemoryLimiter.estimateTopicListSize(topics);
+ listSizeHolder.updateSize(actualSize);
+ maxTopicListInFlightLimiter.withUpdatedPermits(initialPermits, actualSize,
+ isPermitRequestCancelled, permits -> {
+ boolean filterTopics = false;
+ // filter system topic
+ List<String> filteredTopics = topics;
+
+ if (enableSubscriptionPatternEvaluation && topicsPattern.isPresent()) {
+ if (topicsPattern.get().length() <= maxSubscriptionPatternLength) {
+ filterTopics = true;
+ filteredTopics = TopicList.filterTopics(filteredTopics,
+ topicsPattern.get(),
+ topicsPatternImplementation);
+ } else {
+ log.info(
+ "[{}] Subscription pattern provided [{}] was longer "
+ + "than maximum {}.", remoteAddress,
+ topicsPattern.get(),
+ maxSubscriptionPatternLength);
+ }
+ }
+ String hash = TopicList.calculateHash(filteredTopics);
+ boolean hashUnchanged =
+ topicsHash.isPresent() && topicsHash.get().equals(hash);
+ if (hashUnchanged) {
+ filteredTopics = Collections.emptyList();
+ }
+ if (log.isDebugEnabled()) {
+ log.debug("[{}] Received CommandGetTopicsOfNamespace for namespace "
+ + "[//{}] by {}, size:{}", remoteAddress, namespace,
+ requestId,
+ topics.size());
+ }
+ return commandSender.sendGetTopicsOfNamespaceResponse(filteredTopics,
+ hash,
+ filterTopics, !hashUnchanged, requestId, ex -> {
+ log.warn("[{}] Failed to acquire direct memory permits for "
+ + "GetTopicsOfNamespace: {}", remoteAddress,
+ ex.getMessage());
+ commandSender.sendErrorResponse(requestId,
+ ServerError.TooManyRequests,
+ "Cannot acquire permits for direct memory");
+ });
+ }, t -> {
+ log.warn("[{}] Failed to acquire heap memory permits for "
+ + "GetTopicsOfNamespace: {}", remoteAddress, t.getMessage());
+ writeAndFlush(Commands.newError(requestId, ServerError.TooManyRequests,
+ "Failed due to heap memory limit exceeded"));
+ return CompletableFuture.completedFuture(null);
+ });
+ }).whenComplete((__, ___) -> {
+ lookupSemaphore.release();
+ }).exceptionally(ex -> {
+ log.warn("[{}] Error GetTopicsOfNamespace for namespace [//{}] by {}",
+ remoteAddress,
+ namespace, requestId);
+ listSizeHolder.resetIfInitializing();
+ commandSender.sendErrorResponse(requestId,
+ BrokerServiceException.getClientErrorCode(new ServerMetadataException(ex)),
+ ex.getMessage());
+ return null;
+ });
+ }, t -> {
+ log.warn("[{}] Failed to acquire initial heap memory permits for GetTopicsOfNamespace: {}",
+ remoteAddress, t.getMessage());
+ listSizeHolder.resetIfInitializing();
+ writeAndFlush(Commands.newError(requestId, ServerError.TooManyRequests,
+ "Failed due to heap memory limit exceeded"));
+ lookupSemaphore.release();
+ return CompletableFuture.completedFuture(null);
+ });
+ });
+ }
@Override
diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/MessageCumulativeAckTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/MessageCumulativeAckTest.java
index 4a8067b2433..2ca25c89420 100644
--- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/MessageCumulativeAckTest.java
+++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/MessageCumulativeAckTest.java
@@ -47,6 +47,7 @@ import org.apache.pulsar.common.api.proto.CommandAck;
import org.apache.pulsar.common.api.proto.CommandSubscribe;
import org.apache.pulsar.common.api.proto.ProtocolVersion;
import org.apache.pulsar.common.naming.TopicName;
+import org.apache.pulsar.common.semaphore.AsyncDualMemoryLimiter;
import org.apache.pulsar.common.util.Codec;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
@@ -72,7 +73,8 @@ public class MessageCumulativeAckTest {
doReturn(new InetSocketAddress("localhost", 1234)).when(serverCnx).clientAddress();
when(serverCnx.getRemoteEndpointProtocolVersion()).thenReturn(ProtocolVersion.v12.getValue());
when(serverCnx.ctx()).thenReturn(mock(ChannelHandlerContext.class));
- doReturn(new PulsarCommandSenderImpl(null, serverCnx))
+ AsyncDualMemoryLimiter maxTopicListInFlightLimiter = mock(AsyncDualMemoryLimiter.class);
+ doReturn(new PulsarCommandSenderImpl(null, serverCnx, maxTopicListInFlightLimiter))
.when(serverCnx).getCommandSender();
String topicName = TopicName.get("MessageCumulativeAckTest").toString();
diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/PersistentDispatcherFailoverConsumerTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/PersistentDispatcherFailoverConsumerTest.java
index 37cf75d84ca..9ac922ce469 100644
--- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/PersistentDispatcherFailoverConsumerTest.java
+++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/PersistentDispatcherFailoverConsumerTest.java
@@ -79,6 +79,7 @@ import org.apache.pulsar.common.api.proto.CommandSubscribe.InitialPosition;
import org.apache.pulsar.common.api.proto.CommandSubscribe.SubType;
import org.apache.pulsar.common.api.proto.ProtocolVersion;
import org.apache.pulsar.common.naming.NamespaceBundle;
+import org.apache.pulsar.common.semaphore.AsyncDualMemoryLimiter;
import org.apache.pulsar.common.util.netty.EventLoopUtil;
import org.awaitility.Awaitility;
import org.slf4j.Logger;
@@ -145,7 +146,8 @@ public class PersistentDispatcherFailoverConsumerTest {
doReturn(new InetSocketAddress("localhost", 1234)).when(serverCnx).clientAddress();
when(serverCnx.getRemoteEndpointProtocolVersion()).thenReturn(ProtocolVersion.v12.getValue());
when(serverCnx.ctx()).thenReturn(channelCtx);
- doReturn(new PulsarCommandSenderImpl(null, serverCnx))
+ AsyncDualMemoryLimiter maxTopicListInFlightLimiter = mock(AsyncDualMemoryLimiter.class);
+ doReturn(new PulsarCommandSenderImpl(null, serverCnx, maxTopicListInFlightLimiter))
.when(serverCnx).getCommandSender();
serverCnxWithOldVersion = pulsarTestContext.createServerCnxSpy();
@@ -156,7 +158,7 @@ public class PersistentDispatcherFailoverConsumerTest {
when(serverCnxWithOldVersion.getRemoteEndpointProtocolVersion())
.thenReturn(ProtocolVersion.v11.getValue());
when(serverCnxWithOldVersion.ctx()).thenReturn(channelCtx);
- doReturn(new PulsarCommandSenderImpl(null, serverCnxWithOldVersion))
+ doReturn(new PulsarCommandSenderImpl(null, serverCnxWithOldVersion, maxTopicListInFlightLimiter))
.when(serverCnxWithOldVersion).getCommandSender();
NamespaceService nsSvc = pulsarTestContext.getPulsarService().getNamespaceService();
diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/PersistentTopicTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/PersistentTopicTest.java
index 37e1e058cc2..097a4f6de05 100644
--- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/PersistentTopicTest.java
+++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/PersistentTopicTest.java
@@ -135,6 +135,7 @@ import org.apache.pulsar.common.policies.data.TopicPolicies;
import org.apache.pulsar.common.policies.data.stats.SubscriptionStatsImpl;
import org.apache.pulsar.common.protocol.ByteBufPair;
import org.apache.pulsar.common.protocol.schema.SchemaVersion;
+import org.apache.pulsar.common.semaphore.AsyncDualMemoryLimiter;
import org.apache.pulsar.common.util.Codec;
import org.apache.pulsar.common.util.FutureUtil;
import org.apache.pulsar.compaction.CompactedTopic;
@@ -175,6 +176,7 @@ public class PersistentTopicTest extends MockedBookKeeperTestCase {
private EventLoopGroup eventLoopGroup;
private ManagedLedgerFactory managedLedgerFactory;
private ChannelHandlerContext ctx;
+ private AsyncDualMemoryLimiter maxTopicListInFlightLimiter;
@BeforeMethod(alwaysRun = true)
public void setup() throws Exception {
@@ -210,7 +212,8 @@ public class PersistentTopicTest extends MockedBookKeeperTestCase {
doReturn(true).when(serverCnx).isActive();
doReturn(true).when(serverCnx).isWritable();
doReturn(new InetSocketAddress("localhost", 1234)).when(serverCnx).clientAddress();
- doReturn(new PulsarCommandSenderImpl(null, serverCnx))
+ maxTopicListInFlightLimiter = mock(AsyncDualMemoryLimiter.class);
+ doReturn(new PulsarCommandSenderImpl(null, serverCnx, maxTopicListInFlightLimiter))
.when(serverCnx).getCommandSender();
ctx = mock(ChannelHandlerContext.class);
Channel channel = mock(Channel.class);
@@ -531,7 +534,7 @@ public class PersistentTopicTest extends MockedBookKeeperTestCase {
doReturn(true).when(cnx).isWritable();
doReturn(new InetSocketAddress(address, 1234)).when(cnx).clientAddress();
doReturn(address.getHostAddress()).when(cnx).clientSourceAddress();
- doReturn(new PulsarCommandSenderImpl(null, cnx)).when(cnx).getCommandSender();
+ doReturn(new PulsarCommandSenderImpl(null, cnx, maxTopicListInFlightLimiter)).when(cnx).getCommandSender();
return new Producer(topic, cnx, producerId, producerNameBase + producerId, role, false, null,
SchemaVersion.Latest, 0, false, ProducerAccessMode.Shared, Optional.empty(), true);
@@ -967,7 +970,7 @@ public class PersistentTopicTest extends MockedBookKeeperTestCase {
doReturn(true).when(cnx).isWritable();
doReturn(new InetSocketAddress(address, 1234)).when(cnx).clientAddress();
doReturn(address.getHostAddress()).when(cnx).clientSourceAddress();
- doReturn(new PulsarCommandSenderImpl(null, cnx)).when(cnx).getCommandSender();
+ doReturn(new PulsarCommandSenderImpl(null, cnx, maxTopicListInFlightLimiter)).when(cnx).getCommandSender();
return new Consumer(sub, SubType.Shared, topic.getName(), consumerId, 0, consumerNameBase + consumerId, true,
cnx, role, Collections.emptyMap(), false, null, MessageId.latest, DEFAULT_CONSUMER_EPOCH);
diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/client/api/PatternConsumerBackPressureMultipleConsumersHttpLookupTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/client/api/PatternConsumerBackPressureMultipleConsumersHttpLookupTest.java
new file mode 100644
index 00000000000..11916926a0a
--- /dev/null
+++ b/pulsar-broker/src/test/java/org/apache/pulsar/client/api/PatternConsumerBackPressureMultipleConsumersHttpLookupTest.java
@@ -0,0 +1,50 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.pulsar.client.api;
+
+import lombok.extern.slf4j.Slf4j;
+import org.testng.annotations.Test;
+
+@Slf4j
+@Test(groups = "broker-impl")
+public class PatternConsumerBackPressureMultipleConsumersHttpLookupTest
+ extends PatternConsumerBackPressureMultipleConsumersTest {
+ @Override
+ protected boolean useTcpLookup() {
+ return false;
+ }
+
+ @Override
+ protected int getMaxRequestsInFlight() {
+ // AsyncHttpClient is poorly implemented and will use a single byte[] arrays for the response
+ // this is why higher concurrency would lead to OOME due to heap fragmentation
+ return 50;
+ }
+
+ @Override
+ protected int getNumberOfRequests() {
+ return getNumberOfClients();
+ }
+
+ @Override
+ protected int getDirectMemoryRequiredMB() {
+ // don't do any allocations to limit direct memory
+ return Integer.MAX_VALUE;
+ }
+}
diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/client/api/PatternConsumerBackPressureMultipleConsumersTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/client/api/PatternConsumerBackPressureMultipleConsumersTest.java
new file mode 100644
index 00000000000..d91106ab683
--- /dev/null
+++ b/pulsar-broker/src/test/java/org/apache/pulsar/client/api/PatternConsumerBackPressureMultipleConsumersTest.java
@@ -0,0 +1,228 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.pulsar.client.api;
+
+import static org.testng.Assert.assertEquals;
+import io.netty.buffer.ByteBuf;
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.UUID;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.atomic.AtomicInteger;
+import lombok.Cleanup;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.lang3.StringUtils;
+import org.apache.commons.lang3.mutable.MutableInt;
+import org.apache.pulsar.broker.auth.MockedPulsarServiceBaseTest;
+import org.apache.pulsar.broker.stats.prometheus.PrometheusMetricsClient;
+import org.apache.pulsar.client.admin.PulsarAdminException;
+import org.apache.pulsar.client.impl.PulsarClientImpl;
+import org.apache.pulsar.common.allocator.PulsarByteBufAllocator;
+import org.apache.pulsar.common.api.proto.CommandGetTopicsOfNamespace;
+import org.apache.pulsar.common.naming.NamespaceName;
+import org.apache.pulsar.common.stats.JvmMetrics;
+import org.apache.pulsar.common.util.DirectMemoryUtils;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+@Slf4j
+@Test(groups = "broker-impl")
+public class PatternConsumerBackPressureMultipleConsumersTest extends MockedPulsarServiceBaseTest {
+
+ @Override
+ @BeforeMethod
+ protected void setup() throws Exception {
+ isTcpLookup = useTcpLookup();
+ super.internalSetup();
+ setupDefaultTenantAndNamespace();
+ }
+
+ protected boolean useTcpLookup() {
+ return true;
+ }
+
+ @Override
+ @AfterMethod(alwaysRun = true)
+ protected void cleanup() throws Exception {
+ super.internalCleanup();
+ }
+
+ @Test(timeOut = 60 * 1000)
+ public void testGetTopicsWithLargeAmountOfConcurrentClientConnections()
+ throws PulsarAdminException, InterruptedException, IOException {
+ // number of requests to send to the broker
+ final int requests = getNumberOfRequests();
+ // use multiple clients so that each client has a separate connection to the broker
+ final int numberOfClients = getNumberOfClients();
+ // create a long topic name to consume more memory per topic
+ final String topicName = StringUtils.repeat('a', 512) + UUID.randomUUID();
+ // number of topics to create
+ final int topicCount = 8192;
+ // maximum number of requests in flight at any given time
+ final int maxRequestsInFlight = getMaxRequestsInFlight();
+
+ // create a single topic with multiple partitions
+ admin.topics().createPartitionedTopic(topicName, topicCount);
+
+ // reduce available direct memory to reproduce issues with less concurrency
+ long directMemoryRequired = getDirectMemoryRequiredMB() * 1024 * 1024;
+ List<ByteBuf> buffers = allocateDirectMemory(directMemoryRequired);
+ @Cleanup
+ Closeable releaseBuffers = () -> {
+ for (ByteBuf byteBuf : buffers) {
+ byteBuf.release();
+ }
+ };
+
+ @Cleanup("shutdownNow")
+ final ExecutorService executorService = Executors.newFixedThreadPool(Runtime.getRuntime()
+ .availableProcessors());
+
+ @Cleanup
+ PulsarClientSharedResources sharedResources =
+ PulsarClientSharedResources.builder().build();
+ List<PulsarClientImpl> clients = new ArrayList<>(numberOfClients);
+ @Cleanup
+ Closeable closeClients = () -> {
+ for (PulsarClient client : clients) {
+ try {
+ client.close();
+ } catch (PulsarClientException e) {
+ log.error("Failed to close client {}", client, e);
+ }
+ }
+ };
+ for (int i = 0; i < numberOfClients; i++) {
+ PulsarClientImpl client = (PulsarClientImpl) PulsarClient.builder()
+ .serviceUrl(getClientServiceUrl())
+ .sharedResources(sharedResources)
+ .build();
+ clients.add(client);
+ }
+
+ final AtomicInteger success = new AtomicInteger(0);
+ final CountDownLatch latch = new CountDownLatch(requests);
+ final Semaphore semaphore = new Semaphore(maxRequestsInFlight);
+ for (int i = 0; i < requests; i++) {
+ PulsarClientImpl pulsarClientImpl = clients.get(i % numberOfClients);
+ executorService.execute(() -> {
+ semaphore.acquireUninterruptibly();
+ try {
+ pulsarClientImpl.getLookup()
+ .getTopicsUnderNamespace(NamespaceName.get("public", "default"),
+ CommandGetTopicsOfNamespace.Mode.PERSISTENT, ".*", "")
+ .whenComplete((result, ex) -> {
+ semaphore.release();
+ if (ex == null) {
+ success.incrementAndGet();
+ } else {
+ log.error("Failed to get topic list.", ex);
+ }
+ log.info(
+ "latch-count: {}, succeed: {}, available direct mem: {} MB, free heap mem: {}"
+ + " MB",
+ latch.getCount(), success.get(),
+ (DirectMemoryUtils.jvmMaxDirectMemory() - JvmMetrics.getJvmDirectMemoryUsed())
+ / (1024 * 1024), Runtime.getRuntime().freeMemory() / (1024 * 1024));
+ latch.countDown();
+ });
+ } catch (Exception e) {
+ semaphore.release();
+ latch.countDown();
+ log.error("Failed to execute getTopicsUnderNamespace request.", e);
+ }
+ });
+ }
+ latch.await();
+ assertEquals(success.get(), requests);
+
+ validateTopiclistPrometheusMetrics();
+ }
+
+ protected int getNumberOfClients() {
+ return 200;
+ }
+
+ protected int getNumberOfRequests() {
+ return 500;
+ }
+
+ protected int getMaxRequestsInFlight() {
+ return 500;
+ }
+
+ private void validateTopiclistPrometheusMetrics() {
+ MutableInt metricsCount = new MutableInt(0);
+ PrometheusMetricsClient metricsClient =
+ new PrometheusMetricsClient("localhost", pulsar.getListenPortHTTP().get());
+ PrometheusMetricsClient.Metrics metrics = metricsClient.getMetrics();
+ metrics.getNameToDataPoints().entries().forEach(entry -> {
+ if (entry.getKey().startsWith(getMetricsPrefix())) {
+ metricsCount.increment();
+ System.out.println(entry.getKey() + " -> " + entry.getValue());
+ }
+ });
+ assertEquals(metricsCount.intValue(), 26);
+ }
+
+ protected String getMetricsPrefix() {
+ return "pulsar_broker_topic_list_";
+ }
+
+ protected int getDirectMemoryRequiredMB() {
+ return 175;
+ }
+
+ protected String getClientServiceUrl() {
+ return lookupUrl.toString();
+ }
+
+ /**
+ * Allocate direct memory to reduce available direct memory to the given amount of required memory.
+ * @param directMemoryRequired required direct memory in bytes
+ * @return list of ByteBufs allocated to reduce available direct memory
+ */
+ private static List<ByteBuf> allocateDirectMemory(long directMemoryRequired) {
+ if (directMemoryRequired <= 0 || directMemoryRequired == Integer.MAX_VALUE) {
+ return List.of();
+ }
+ long usedMemory = JvmMetrics.getJvmDirectMemoryUsed();
+ long maxMemory = DirectMemoryUtils.jvmMaxDirectMemory();
+ long availableMemory = maxMemory - usedMemory;
+ List<ByteBuf> buffers = new ArrayList<>();
+ if (availableMemory > directMemoryRequired) {
+ long allocateRemaining = availableMemory - directMemoryRequired;
+ log.info("Making allocations for {} MB to reduce available direct memory",
+ allocateRemaining / (1024 * 1024));
+ int blockSize = 5 * 1024 * 1024;
+ while (allocateRemaining > 0) {
+ ByteBuf byteBuf = PulsarByteBufAllocator.DEFAULT.directBuffer(blockSize);
+ buffers.add(byteBuf);
+ allocateRemaining -= blockSize;
+ }
+ }
+ return buffers;
+ }
+}
diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/client/api/PatternConsumerBackPressureTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/client/api/PatternConsumerBackPressureTest.java
index f33de0a40b6..16a16ebcead 100644
--- a/pulsar-broker/src/test/java/org/apache/pulsar/client/api/PatternConsumerBackPressureTest.java
+++ b/pulsar-broker/src/test/java/org/apache/pulsar/client/api/PatternConsumerBackPressureTest.java
@@ -23,6 +23,7 @@ import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicInteger;
+import lombok.Cleanup;
import lombok.extern.slf4j.Slf4j;
import org.apache.pulsar.broker.auth.MockedPulsarServiceBaseTest;
import org.apache.pulsar.client.admin.PulsarAdminException;
@@ -70,6 +71,7 @@ public class PatternConsumerBackPressureTest extends MockedPulsarServiceBaseTest
final int requests = 2048;
final String topicName = UUID.randomUUID().toString();
admin.topics().createPartitionedTopic(topicName, topicCount);
+ @Cleanup("shutdownNow")
final ExecutorService executorService = Executors.newFixedThreadPool(Runtime.getRuntime()
.availableProcessors());
diff --git a/pulsar-common/src/main/java/org/apache/pulsar/common/protocol/Commands.java b/pulsar-common/src/main/java/org/apache/pulsar/common/protocol/Commands.java
index cab4dc8bcab..85e6a73f787 100644
--- a/pulsar-common/src/main/java/org/apache/pulsar/common/protocol/Commands.java
+++ b/pulsar-common/src/main/java/org/apache/pulsar/common/protocol/Commands.java
@@ -1236,7 +1236,7 @@ public class Commands {
public static BaseCommand newGetTopicsOfNamespaceResponseCommand(List<String> topics, String topicsHash,
boolean filtered, boolean changed,
long requestId) {
- BaseCommand cmd = localCmd(Type.GET_TOPICS_OF_NAMESPACE_RESPONSE);
+ BaseCommand cmd = new BaseCommand().setType(Type.GET_TOPICS_OF_NAMESPACE_RESPONSE);
CommandGetTopicsOfNamespaceResponse topicsResponse = cmd.setGetTopicsOfNamespaceResponse();
topicsResponse.setRequestId(requestId);
for (int i = 0; i < topics.size(); i++) {
@@ -1250,12 +1250,6 @@ public class Commands {
return cmd;
}
- public static ByteBuf newGetTopicsOfNamespaceResponse(List<String> topics, String topicsHash,
- boolean filtered, boolean changed, long requestId) {
- return serializeWithSize(newGetTopicsOfNamespaceResponseCommand(
- topics, topicsHash, filtered, changed, requestId));
- }
-
private static final ByteBuf cmdPing;
static {
@@ -1635,7 +1629,7 @@ public class Commands {
*/
public static BaseCommand newWatchTopicListSuccess(long requestId, long watcherId, String topicsHash,
List<String> topics) {
- BaseCommand cmd = localCmd(Type.WATCH_TOPIC_LIST_SUCCESS);
+ BaseCommand cmd = new BaseCommand().setType(Type.WATCH_TOPIC_LIST_SUCCESS);
cmd.setWatchTopicListSuccess()
.setRequestId(requestId)
.setWatcherId(watcherId);
@@ -1654,7 +1648,7 @@ public class Commands {
*/
public static BaseCommand newWatchTopicUpdate(long watcherId,
List<String> newTopics, List<String> deletedTopics, String topicsHash) {
- BaseCommand cmd = localCmd(Type.WATCH_TOPIC_UPDATE);
+ BaseCommand cmd = new BaseCommand().setType(Type.WATCH_TOPIC_UPDATE);
cmd.setWatchTopicUpdate()
.setWatcherId(watcherId)
.setTopicsHash(topicsHash)
@@ -1672,9 +1666,12 @@ public class Commands {
}
public static ByteBuf serializeWithSize(BaseCommand cmd) {
+ return serializeWithPrecalculatedSerializedSize(cmd, cmd.getSerializedSize());
+ }
+
+ public static ByteBuf serializeWithPrecalculatedSerializedSize(BaseCommand cmd, int cmdSize) {
// / Wire format
// [TOTAL_SIZE] [CMD_SIZE][CMD]
- int cmdSize = cmd.getSerializedSize();
int totalSize = cmdSize + 4;
int frameSize = totalSize + 4;
diff --git a/pulsar-common/src/main/java/org/apache/pulsar/common/semaphore/AsyncDualMemoryLimiter.java b/pulsar-common/src/main/java/org/apache/pulsar/common/semaphore/AsyncDualMemoryLimiter.java
new file mode 100644
index 00000000000..59be8db8923
--- /dev/null
+++ b/pulsar-common/src/main/java/org/apache/pulsar/common/semaphore/AsyncDualMemoryLimiter.java
@@ -0,0 +1,114 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.pulsar.common.semaphore;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.function.BooleanSupplier;
+import java.util.function.Function;
+
+/**
+ * An abstraction for an asynchronous memory semaphore that tracks separate limits for heap and direct memory.
+ */
+public interface AsyncDualMemoryLimiter {
+ enum LimitType {
+ HEAP_MEMORY, // For heap memory allocation
+ DIRECT_MEMORY // For direct memory allocation
+ }
+
+ /**
+ * Acquire permits for the specified memory size.
+ * Returned future completes when memory permits are available.
+ * It will complete exceptionally with AsyncSemaphore.PermitAcquireTimeoutException on timeout
+ * and exceptionally with AsyncSemaphore.PermitAcquireQueueFullException when queue full
+ *
+ * @param memorySize the size of memory to acquire permits for
+ * @param limitType the type of memory limit (HEAP_MEMORY or DIRECT_MEMORY)
+ * @param isCancelled supplier that returns true if acquisition should be cancelled
+ * @return CompletableFuture that completes with permit when available
+ */
+ CompletableFuture<AsyncDualMemoryLimiterPermit> acquire(long memorySize, LimitType limitType,
+ BooleanSupplier isCancelled);
+
+ /**
+ * Acquire or release permits for previously acquired permits by updating the requested memory size.
+ * Returns a future that completes when permits are available.
+ * It will complete exceptionally with AsyncSemaphore.PermitAcquireTimeoutException on timeout
+ * and exceptionally with AsyncSemaphore.PermitAcquireQueueFullException when queue full
+ * The provided permit is released when the permits are successfully acquired and the returned updated
+ * permit replaces the old instance.
+ *
+ * @param permit the previously acquired permit to update
+ * @param newMemorySize the new memory size to update to
+ * @param isCancelled supplier that returns true if update should be cancelled
+ * @return CompletableFuture that completes with permit when available
+ */
+ CompletableFuture<AsyncDualMemoryLimiterPermit> update(AsyncDualMemoryLimiterPermit permit, long newMemorySize,
+ BooleanSupplier isCancelled);
+
+ /**
+ * Release previously acquired permit.
+ * Must be called to prevent memory permit leaks.
+ *
+ * @param permit the permit to release
+ */
+ void release(AsyncDualMemoryLimiterPermit permit);
+ /**
+ * Execute the specified function with acquired permits and release the permits after the returned future completes.
+ * @param memorySize memory size to acquire permits for
+ * @param limitType memory limit type to acquire permits for
+ * @param function function to execute with acquired permits
+ * @return result of the function
+ * @param <T> type of the CompletableFuture returned by the function
+ */
+ default <T> CompletableFuture<T> withAcquiredPermits(long memorySize, LimitType limitType,
+ BooleanSupplier isCancelled,
+ Function<AsyncDualMemoryLimiterPermit,
+ CompletableFuture<T>> function,
+ Function<Throwable, CompletableFuture<T>>
+ permitAcquireErrorHandler) {
+ return AsyncDualMemoryLimiterUtil.withPermitsFuture(acquire(memorySize, limitType, isCancelled), function,
+ permitAcquireErrorHandler, this::release);
+ }
+
+ /**
+ * Executed the specified function with updated permits and release the permits after the returned future completes.
+ * @param initialPermit initial permit to update
+ * @param newMemorySize new memory size to update to
+ * @param function function to execute with updated permits
+ * @return result of the function
+ * @param <T> type of the CompletableFuture returned by the function
+ */
+ default <T> CompletableFuture<T> withUpdatedPermits(AsyncDualMemoryLimiterPermit initialPermit, long newMemorySize,
+ BooleanSupplier isCancelled,
+ Function<AsyncDualMemoryLimiterPermit,
+ CompletableFuture<T>> function,
+ Function<Throwable, CompletableFuture<T>>
+ permitAcquireErrorHandler) {
+ return AsyncDualMemoryLimiterUtil.withPermitsFuture(update(initialPermit, newMemorySize, isCancelled), function,
+ permitAcquireErrorHandler, this::release);
+ }
+
+ /**
+ * Represents an acquired permit for memory limiting that can be updated or released.
+ */
+ interface AsyncDualMemoryLimiterPermit {
+ long getPermits();
+ LimitType getLimitType();
+ }
+}
diff --git a/pulsar-common/src/main/java/org/apache/pulsar/common/semaphore/AsyncDualMemoryLimiterImpl.java b/pulsar-common/src/main/java/org/apache/pulsar/common/semaphore/AsyncDualMemoryLimiterImpl.java
new file mode 100644
index 00000000000..d6b31aa72ee
--- /dev/null
+++ b/pulsar-common/src/main/java/org/apache/pulsar/common/semaphore/AsyncDualMemoryLimiterImpl.java
@@ -0,0 +1,177 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.pulsar.common.semaphore;
+
+import io.netty.util.concurrent.DefaultThreadFactory;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.function.BooleanSupplier;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Implementation of AsyncDualMemoryLimiter with separate limits for heap and direct memory.
+ */
+public class AsyncDualMemoryLimiterImpl implements AsyncDualMemoryLimiter, AutoCloseable {
+ private static final Logger log = LoggerFactory.getLogger(AsyncDualMemoryLimiterImpl.class);
+
+ private final ScheduledExecutorService executor;
+ private final boolean shutdownExecutor;
+ private final AsyncSemaphoreImpl heapLimiter;
+ private final AsyncSemaphoreImpl directLimiter;
+
+ /**
+ * Creates an AsyncDualMemoryLimiterImpl with the given parameters.
+ * @param maxHeapMemory max heap memory available for allocation
+ * @param maxHeapQueueSize max number of requests that can be queued for heap memory allocation
+ * @param heapTimeoutMillis timeout in milliseconds for heap memory allocation
+ * @param maxDirectMemory max direct memory available for allocation
+ * @param maxDirectQueueSize max number of requests that can be queued for direct memory allocation
+ * @param directTimeoutMillis timeout in milliseconds for direct memory allocation
+ * @param executor executor service to use for scheduling timeouts, it is expected to be single threaded
+ */
+ public AsyncDualMemoryLimiterImpl(long maxHeapMemory, int maxHeapQueueSize, long heapTimeoutMillis,
+ long maxDirectMemory, int maxDirectQueueSize, long directTimeoutMillis,
+ ScheduledExecutorService executor) {
+ this(maxHeapMemory, maxHeapQueueSize, heapTimeoutMillis, maxDirectMemory, maxDirectQueueSize,
+ directTimeoutMillis, executor, false);
+ }
+
+ /**
+ * Creates an AsyncDualMemoryLimiterImpl with the given parameters.
+ * @param maxHeapMemory max heap memory available for allocation
+ * @param maxHeapQueueSize max number of requests that can be queued for heap memory allocation
+ * @param heapTimeoutMillis timeout in milliseconds for heap memory allocation
+ * @param maxDirectMemory max direct memory available for allocation
+ * @param maxDirectQueueSize max number of requests that can be queued for direct memory allocation
+ * @param directTimeoutMillis timeout in milliseconds for direct memory allocation
+ */
+ public AsyncDualMemoryLimiterImpl(long maxHeapMemory, int maxHeapQueueSize, long heapTimeoutMillis,
+ long maxDirectMemory, int maxDirectQueueSize, long directTimeoutMillis) {
+ this(maxHeapMemory, maxHeapQueueSize, heapTimeoutMillis, maxDirectMemory, maxDirectQueueSize,
+ directTimeoutMillis, createExecutor(), true);
+ }
+
+ AsyncDualMemoryLimiterImpl(long maxHeapMemory, int maxHeapQueueSize, long heapTimeoutMillis,
+ long maxDirectMemory, int maxDirectQueueSize, long directTimeoutMillis,
+ ScheduledExecutorService executor, boolean shutdownExecutor) {
+ this.executor = executor;
+ this.shutdownExecutor = shutdownExecutor;
+ this.heapLimiter = new AsyncSemaphoreImpl(maxHeapMemory, maxHeapQueueSize, heapTimeoutMillis, executor,
+ this::recordHeapWaitTime);
+ this.directLimiter = new AsyncSemaphoreImpl(maxDirectMemory, maxDirectQueueSize, directTimeoutMillis, executor,
+ this::recordDirectWaitTime);
+ }
+
+ private static ScheduledExecutorService createExecutor() {
+ return Executors.newSingleThreadScheduledExecutor(
+ new DefaultThreadFactory("async-dual-memory-limiter"));
+ }
+
+ @Override
+ public CompletableFuture<AsyncDualMemoryLimiterPermit> acquire(long memorySize, LimitType limitType,
+ BooleanSupplier isCancelled) {
+ AsyncSemaphore limiter = getLimiter(limitType);
+ return limiter.acquire(memorySize, isCancelled).thenApply(result ->
+ new DualMemoryLimiterPermit(limitType, result));
+ }
+
+ protected AsyncSemaphore getLimiter(LimitType limitType) {
+ switch (limitType) {
+ case HEAP_MEMORY:
+ return heapLimiter;
+ case DIRECT_MEMORY:
+ return directLimiter;
+ default:
+ throw new IllegalArgumentException("Unsupported limit type: " + limitType);
+ }
+ }
+
+ @Override
+ public CompletableFuture<AsyncDualMemoryLimiterPermit> update(AsyncDualMemoryLimiterPermit permit,
+ long newMemorySize, BooleanSupplier isCancelled) {
+ AsyncSemaphore limiter = getLimiter(permit.getLimitType());
+ return limiter.update(castToImplementation(permit).getUnderlyingPermit(), newMemorySize, isCancelled)
+ .thenApply(updatedPermit -> new DualMemoryLimiterPermit(permit.getLimitType(), updatedPermit));
+ }
+
+ @Override
+ public void release(AsyncDualMemoryLimiterPermit permit) {
+ AsyncSemaphore limiter = getLimiter(permit.getLimitType());
+ limiter.release(castToImplementation(permit).getUnderlyingPermit());
+ }
+
+ private DualMemoryLimiterPermit castToImplementation(AsyncDualMemoryLimiterPermit permit) {
+ if (permit instanceof DualMemoryLimiterPermit dualMemoryLimiterPermit) {
+ return dualMemoryLimiterPermit;
+ } else {
+ throw new IllegalArgumentException("Invalid permit type");
+ }
+ }
+
+ /**
+ * Record the wait time for a heap memory allocation permit.
+ * @param waitTimeNanos wait time in nanoseconds, or Long.MAX_VALUE if the allocation timed out
+ */
+ protected void recordHeapWaitTime(long waitTimeNanos) {
+
+ }
+
+ /**
+ * Record the wait time for a direct memory allocation permit.
+ * @param waitTimeNanos wait time in nanoseconds, or Long.MAX_VALUE if the allocation timed out
+ */
+ protected void recordDirectWaitTime(long waitTimeNanos) {
+
+ }
+
+ @Override
+ public void close() {
+ heapLimiter.close();
+ directLimiter.close();
+ if (shutdownExecutor) {
+ executor.shutdown();
+ }
+ }
+
+ private static class DualMemoryLimiterPermit implements AsyncDualMemoryLimiterPermit {
+ private final LimitType limitType;
+ private final AsyncSemaphore.AsyncSemaphorePermit underlyingPermit;
+
+ DualMemoryLimiterPermit(LimitType limitType, AsyncSemaphore.AsyncSemaphorePermit underlyingPermit) {
+ this.limitType = limitType;
+ this.underlyingPermit = underlyingPermit;
+ }
+
+ @Override
+ public long getPermits() {
+ return underlyingPermit.getPermits();
+ }
+
+ @Override
+ public LimitType getLimitType() {
+ return limitType;
+ }
+
+ public AsyncSemaphore.AsyncSemaphorePermit getUnderlyingPermit() {
+ return underlyingPermit;
+ }
+ }
+}
diff --git a/pulsar-common/src/main/java/org/apache/pulsar/common/semaphore/AsyncDualMemoryLimiterUtil.java b/pulsar-common/src/main/java/org/apache/pulsar/common/semaphore/AsyncDualMemoryLimiterUtil.java
new file mode 100644
index 00000000000..616681dc155
--- /dev/null
+++ b/pulsar-common/src/main/java/org/apache/pulsar/common/semaphore/AsyncDualMemoryLimiterUtil.java
@@ -0,0 +1,123 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.pulsar.common.semaphore;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.channel.ChannelHandlerContext;
+import java.util.concurrent.CompletableFuture;
+import java.util.function.BooleanSupplier;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import lombok.experimental.UtilityClass;
+import org.apache.commons.lang3.tuple.Pair;
+import org.apache.pulsar.common.api.proto.BaseCommand;
+import org.apache.pulsar.common.protocol.Commands;
+import org.apache.pulsar.common.semaphore.AsyncDualMemoryLimiter.AsyncDualMemoryLimiterPermit;
+
+@UtilityClass
+public class AsyncDualMemoryLimiterUtil {
+
+ /**
+ * Execute a function with acquired permits and ensure permits are released after completion.
+ * This method handles the lifecycle of permits - acquisition, usage, and release, including error cases.
+ *
+ * @param permitsFuture Future that will complete with the required permits
+ * @param function Function to execute once permits are acquired that returns a CompletableFuture
+ * @param permitAcquireErrorHandler Handler for permit acquisition errors that returns a CompletableFuture
+ * @param releaser Consumer that handles releasing the permits
+ * @param <T> The type of result returned by the function
+ * @return CompletableFuture that completes with the result of the function execution
+ */
+ public static <T> CompletableFuture<T> withPermitsFuture(
+ CompletableFuture<AsyncDualMemoryLimiterPermit>
+ permitsFuture,
+ Function<AsyncDualMemoryLimiterPermit,
+ CompletableFuture<T>> function,
+ Function<Throwable, CompletableFuture<T>>
+ permitAcquireErrorHandler,
+ Consumer<AsyncDualMemoryLimiterPermit> releaser) {
+ return permitsFuture
+ // combine the permits and error into a single pair so that it can be used in thenCompose
+ .handle((permits, permitAcquireError) ->
+ Pair.of(permits, permitAcquireError))
+ .thenCompose(permitsAndError -> {
+ if (permitsAndError.getRight() != null) {
+ // permits weren't acquired
+ return permitAcquireErrorHandler.apply(permitsAndError.getRight());
+ } else {
+ // permits were acquired
+ AsyncDualMemoryLimiterPermit permits = permitsAndError.getLeft();
+ try {
+ return function.apply(permits)
+ .whenComplete((__, ___) ->
+ // release the permits
+ releaser.accept(permits));
+ } catch (Throwable t) {
+ // release the permits if an exception occurs before the function returns
+ releaser.accept(permits);
+ throw t;
+ }
+ }
+ });
+ }
+
+ /**
+ * Acquires permits and writes the command as a response to the channel.
+ * Releases the permits after the response has been written to the socket or if the write fails.
+ *
+ * @param ctx the channel handler context used for writing the response
+ * @param dualMemoryLimiter the memory limiter used to acquire and release memory permits
+ * @param isCancelled supplier that indicates if the permit acquisition should be cancelled
+ * @param command the base command to serialize and write to the channel
+ * @param permitAcquireErrorHandler handler for errors that occur during permit acquisition
+ * @return a future that completes when the command has been written to the channel's outbound buffer
+ */
+ public static CompletableFuture<Void> acquireDirectMemoryPermitsAndWriteAndFlush(ChannelHandlerContext ctx,
+ AsyncDualMemoryLimiter
+ dualMemoryLimiter,
+ BooleanSupplier isCancelled,
+ BaseCommand command,
+ Consumer<Throwable>
+ permitAcquireErrorHandler
+ ) {
+ // Calculate serialized size before acquiring permits
+ int serializedSize = command.getSerializedSize();
+ // Acquire permits
+ return dualMemoryLimiter.acquire(serializedSize, AsyncDualMemoryLimiter.LimitType.DIRECT_MEMORY, isCancelled)
+ .whenComplete((permits, t) -> {
+ if (t != null) {
+ permitAcquireErrorHandler.accept(t);
+ return;
+ }
+ try {
+ // Serialize the response
+ ByteBuf outBuf = Commands.serializeWithPrecalculatedSerializedSize(command, serializedSize);
+ // Write the response
+ ctx.writeAndFlush(outBuf).addListener(future -> {
+ // Release permits after the response has been written to the socket
+ dualMemoryLimiter.release(permits);
+ });
+ } catch (Throwable e) {
+ // Return permits if an exception occurs before writeAndFlush is called successfully
+ dualMemoryLimiter.release(permits);
+ throw e;
+ }
+ }).thenApply(__ -> null);
+ }
+}
diff --git a/pulsar-common/src/main/java/org/apache/pulsar/common/semaphore/AsyncSemaphore.java b/pulsar-common/src/main/java/org/apache/pulsar/common/semaphore/AsyncSemaphore.java
new file mode 100644
index 00000000000..53ca0670ae6
--- /dev/null
+++ b/pulsar-common/src/main/java/org/apache/pulsar/common/semaphore/AsyncSemaphore.java
@@ -0,0 +1,129 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.pulsar.common.semaphore;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.function.BooleanSupplier;
+
+/**
+ * An abstraction for a generic asynchronous semaphore.
+ */
+public interface AsyncSemaphore {
+ /**
+ * Acquire permits from the semaphore.
+ * Returned future completes when permits are available.
+ * It will complete exceptionally with AsyncSemaphore.PermitAcquireTimeoutException on timeout
+ * and exceptionally with AsyncSemaphore.PermitAcquireQueueFullException when queue full
+ *
+ * @param permits number of permits to acquire
+ * @param isCancelled supplier that returns true if acquisition should be cancelled
+ * @return CompletableFuture that completes with permit when available
+ */
+ CompletableFuture<AsyncSemaphorePermit> acquire(long permits, BooleanSupplier isCancelled);
+
+ /**
+ * Acquire or release permits for previously acquired permits by updating the permits.
+ * Returns a future that completes when permits are available.
+ * The provided permit is released when the permits are successfully acquired and the returned updated
+ * permit replaces the old instance.
+ * It will complete exceptionally with AsyncSemaphore.PermitAcquireTimeoutException on timeout
+ * and exceptionally with AsyncSemaphore.PermitAcquireQueueFullException when queue full
+ *
+ * @param permit previously acquired permit to update
+ * @param newPermits new number of permits to update to
+ * @param isCancelled supplier that returns true if update should be cancelled
+ * @return CompletableFuture that completes with permit when available
+ */
+ CompletableFuture<AsyncSemaphorePermit> update(AsyncSemaphorePermit permit, long newPermits,
+ BooleanSupplier isCancelled);
+
+ /**
+ * Release previously acquired permit.
+ * Must be called to prevent permit leaks.
+ *
+ * @param permit permit to release
+ */
+ void release(AsyncSemaphorePermit permit);
+ /**
+ * Get the number of available permits.
+ */
+ long getAvailablePermits();
+
+ /**
+ * Get the number of acquired permits.
+ */
+ long getAcquiredPermits();
+
+ /**
+ * Get the current size of queued requests.
+ */
+ int getQueueSize();
+
+ /**
+ * Abstract base class for all exceptions thrown by acquire or update.
+ */
+ abstract class PermitAcquireException extends RuntimeException {
+ public PermitAcquireException(String message) {
+ super(message);
+ }
+ }
+
+ /**
+ * Exception thrown when permit acquisition times out.
+ */
+ class PermitAcquireTimeoutException extends PermitAcquireException {
+ public PermitAcquireTimeoutException(String message) {
+ super(message);
+ }
+ }
+
+ /**
+ * Exception thrown when permit acquisition queue is full.
+ */
+ class PermitAcquireQueueFullException extends PermitAcquireException {
+ public PermitAcquireQueueFullException(String message) {
+ super(message);
+ }
+ }
+
+ /**
+ * Exception thrown when permit acquisition is attempted on a closed semaphore.
+ */
+ class PermitAcquireAlreadyClosedException extends PermitAcquireException {
+ public PermitAcquireAlreadyClosedException(String message) {
+ super(message);
+ }
+ }
+
+ /**
+ * Exception thrown when permit acquisition is cancelled.
+ */
+ class PermitAcquireCancelledException extends PermitAcquireException {
+ public PermitAcquireCancelledException(String message) {
+ super(message);
+ }
+ }
+
+ /**
+ * Represents an acquired permit that can be updated or released.
+ */
+ interface AsyncSemaphorePermit {
+ long getPermits();
+ }
+}
diff --git a/pulsar-common/src/main/java/org/apache/pulsar/common/semaphore/AsyncSemaphoreImpl.java b/pulsar-common/src/main/java/org/apache/pulsar/common/semaphore/AsyncSemaphoreImpl.java
new file mode 100644
index 00000000000..7fc506753a8
--- /dev/null
+++ b/pulsar-common/src/main/java/org/apache/pulsar/common/semaphore/AsyncSemaphoreImpl.java
@@ -0,0 +1,369 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.pulsar.common.semaphore;
+
+import io.netty.util.concurrent.DefaultThreadFactory;
+import java.util.Queue;
+import java.util.concurrent.ArrayBlockingQueue;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executors;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicLongFieldUpdater;
+import java.util.function.BooleanSupplier;
+import java.util.function.LongConsumer;
+import org.apache.pulsar.common.util.FutureUtil;
+import org.apache.pulsar.common.util.Runnables;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Implementation of AsyncSemaphore with timeout and queue size limits.
+ */
+public class AsyncSemaphoreImpl implements AsyncSemaphore, AutoCloseable {
+ private static final Logger log = LoggerFactory.getLogger(AsyncSemaphoreImpl.class);
+
+ private final AtomicLong availablePermits;
+ private final Queue<PendingRequest> queue;
+ private final long maxPermits;
+ private final long timeoutMillis;
+ private final ScheduledExecutorService executor;
+ private final boolean shutdownExecutor;
+ private final LongConsumer queueLatencyRecorder;
+ private final AtomicBoolean closed = new AtomicBoolean(false);
+ private final Runnable processQueueRunnable = Runnables.catchingAndLoggingThrowables(this::internalProcessQueue);
+ private final ScheduledFuture<?> processQueueScheduledFuture;
+
+ /**
+ * Creates an AsyncSemaphoreImpl with the given parameters.
+ * @param maxPermits max number of permits available for acquisition, set to <= 0 for unbounded semaphore (not
+ * recommended)
+ * @param maxQueueSize max number of requests that can be queued, set to <= 0 for unbounded queue (not recommended)
+ * @param timeoutMillis timeout in milliseconds for acquiring permits
+ */
+ public AsyncSemaphoreImpl(long maxPermits, int maxQueueSize, long timeoutMillis) {
+ this(maxPermits, maxQueueSize, timeoutMillis, maxPermits > 0 ? createExecutor() : null, maxPermits > 0, null);
+ }
+
+ /**
+ * Creates an AsyncSemaphoreImpl with the given parameters.
+ * @param maxPermits max number of permits available for acquisition, set to <= 0 for unbounded semaphore (not
+ * recommended)
+ * @param maxQueueSize max number of requests that can be queued, set to <= 0 for unbounded queue (not recommended)
+ * @param timeoutMillis timeout in milliseconds for acquiring permits
+ * @param executor executor service to use for scheduling timeouts, it is expected to be single threaded
+ * @param queueLatencyRecorder consumer to record queue latency, Long.MAX_VALUE is used for requests that timed out
+ */
+ public AsyncSemaphoreImpl(long maxPermits, int maxQueueSize, long timeoutMillis,
+ ScheduledExecutorService executor, LongConsumer queueLatencyRecorder) {
+ this(maxPermits, maxQueueSize, timeoutMillis, executor, false, queueLatencyRecorder);
+ }
+
+ AsyncSemaphoreImpl(long maxPermits, int maxQueueSize, long timeoutMillis, ScheduledExecutorService executor,
+ boolean shutdownExecutor, LongConsumer queueLatencyRecorder) {
+ this.availablePermits = new AtomicLong(maxPermits);
+ this.maxPermits = maxPermits;
+ this.queue = maxQueueSize > 0 ? new ArrayBlockingQueue<>(maxQueueSize) : new LinkedBlockingQueue<>();
+ this.timeoutMillis = timeoutMillis;
+ this.executor = executor;
+ this.shutdownExecutor = shutdownExecutor;
+ this.queueLatencyRecorder = queueLatencyRecorder;
+ // scheduled task that runs the processQueue method every half of the timeout
+ // this is to support cancellation in cases where the head of the queue request is blocking others
+ // from proceeding and it happens to already be cancelled.
+ this.processQueueScheduledFuture = executor != null
+ ? executor.schedule(processQueueRunnable, timeoutMillis / 2, TimeUnit.MILLISECONDS) : null;
+ }
+
+ private static ScheduledExecutorService createExecutor() {
+ return Executors.newSingleThreadScheduledExecutor(
+ new DefaultThreadFactory("async-semaphore-executor"));
+ }
+
+ @Override
+ public CompletableFuture<AsyncSemaphorePermit> acquire(long permits, BooleanSupplier isCancelled) {
+ return internalAcquire(permits, permits, isCancelled);
+ }
+
+ private CompletableFuture<AsyncSemaphorePermit> internalAcquire(long permits, long acquirePermits,
+ BooleanSupplier isCancelled) {
+ validatePermits(permits);
+
+ // if maximum permits is <= 0, then the semaphore is unbounded
+ if (isUnbounded()) {
+ return CompletableFuture.completedFuture(new SemaphorePermit(permits));
+ }
+
+ CompletableFuture<AsyncSemaphorePermit> future = new CompletableFuture<>();
+
+ if (closed.get()) {
+ future.completeExceptionally(new PermitAcquireAlreadyClosedException("Semaphore is closed"));
+ return future;
+ }
+
+ PendingRequest request = new PendingRequest(permits, acquirePermits, future, isCancelled);
+ if (!queue.offer(request)) {
+ future.completeExceptionally(new PermitAcquireQueueFullException(
+ "Semaphore queue is full"));
+ return future;
+ }
+ // Schedule timeout
+ ScheduledFuture<?> timeoutTask = executor.schedule(() -> {
+ if (!request.future.isDone() && queue.remove(request)) {
+ // timeout is recorded with Long.MAX_VALUE as the age
+ recordQueueLatency(Long.MAX_VALUE);
+ // also record the time in the queue
+ recordQueueLatency(request.getAgeNanos());
+ future.completeExceptionally(new PermitAcquireTimeoutException(
+ "Permit acquisition timed out"));
+ // the next request might have smaller permits and that might be processed
+ processQueue();
+ }
+ }, timeoutMillis, TimeUnit.MILLISECONDS);
+ request.setTimeoutTask(timeoutTask);
+
+ processQueue();
+ return future;
+ }
+
+ private void validatePermits(long permits) {
+ if (permits < 0) {
+ throw new IllegalArgumentException("Invalid negative permits value: " + permits);
+ }
+ if (!isUnbounded() && permits > maxPermits) {
+ throw new IllegalArgumentException(
+ "Requested permits=" + permits + " is larger than maxPermits=" + maxPermits);
+ }
+ }
+
+ private boolean isUnbounded() {
+ return maxPermits <= 0;
+ }
+
+ private void recordQueueLatency(long ageNanos) {
+ if (queueLatencyRecorder != null) {
+ queueLatencyRecorder.accept(ageNanos);
+ }
+ }
+
+ @Override
+ public CompletableFuture<AsyncSemaphorePermit> update(AsyncSemaphorePermit permit, long newPermits,
+ BooleanSupplier isCancelled) {
+ validatePermits(newPermits);
+ if (isUnbounded()) {
+ return CompletableFuture.completedFuture(new SemaphorePermit(newPermits));
+ }
+ long oldPermits = permit.getPermits();
+ long additionalPermits = newPermits - oldPermits;
+ if (additionalPermits > 0) {
+ CompletableFuture<AsyncSemaphorePermit> acquireFuture =
+ internalAcquire(newPermits, additionalPermits, isCancelled);
+ // return a future that completes after original permits have been released when the acquisition
+ // has been successfully completed
+ CompletableFuture<AsyncSemaphorePermit> returnedFuture =
+ acquireFuture.thenApply(p -> {
+ // mark the old permits as released without adding the permits to availablePermits
+ castToImplementation(permit).releasePermits();
+ return p;
+ });
+ // add cancellation support for returned future, so that it cancels the acquireFuture if the returnedFuture
+ // is cancelled
+ returnedFuture.whenComplete((p, t) -> {
+ if (t != null && FutureUtil.unwrapCompletionException(t) instanceof CancellationException) {
+ acquireFuture.cancel(false);
+ }
+ });
+ return returnedFuture;
+ }
+ if (additionalPermits < 0) {
+ // new permits are less than the old ones, so we return the difference
+ availablePermits.addAndGet(-additionalPermits);
+ processQueue();
+ }
+ // mark the old permits as released without adding the permits to availablePermits
+ castToImplementation(permit).releasePermits();
+ // return the new permits immediately
+ return CompletableFuture.completedFuture(new SemaphorePermit(newPermits));
+ }
+
+ @Override
+ public void release(AsyncSemaphorePermit permit) {
+ if (isUnbounded()) {
+ return;
+ }
+ long releasedPermits = castToImplementation(permit).releasePermits();
+ if (releasedPermits > 0) {
+ availablePermits.addAndGet(releasedPermits);
+ processQueue();
+ }
+ }
+
+ @Override
+ public long getAvailablePermits() {
+ if (isUnbounded()) {
+ return Long.MAX_VALUE;
+ }
+ return availablePermits.get();
+ }
+
+ @Override
+ public long getAcquiredPermits() {
+ if (isUnbounded()) {
+ return 0;
+ }
+ return maxPermits - availablePermits.get();
+ }
+
+ @Override
+ public int getQueueSize() {
+ return queue.size();
+ }
+
+ private SemaphorePermit castToImplementation(AsyncSemaphorePermit permit) {
+ if (permit instanceof SemaphorePermit semaphorePermit) {
+ return semaphorePermit;
+ } else {
+ throw new IllegalArgumentException("Invalid permit type");
+ }
+ }
+
+ private void processQueue() {
+ if (closed.get()) {
+ return;
+ }
+ executor.execute(processQueueRunnable);
+ }
+
+ private synchronized void internalProcessQueue() {
+ while (!closed.get() && !queue.isEmpty()) {
+ long current = availablePermits.get();
+ if (current <= 0) {
+ break;
+ }
+
+ PendingRequest request = queue.peek();
+ if (request == null) {
+ break;
+ }
+
+ if (request.isCancelled.getAsBoolean()) {
+ request.cancelTimeoutTask();
+ queue.remove(request);
+ request.future.completeExceptionally(
+ new PermitAcquireCancelledException("Permit acquisition was cancelled"));
+ continue;
+ }
+
+ // request future has been completed by user code cancellation, remove it from the queue
+ if (request.future.isDone()) {
+ request.cancelTimeoutTask();
+ queue.remove(request);
+ continue;
+ }
+
+ if (request.acquirePermits <= current) {
+ availablePermits.addAndGet(-request.acquirePermits);
+ request.cancelTimeoutTask();
+ queue.remove(request);
+ SemaphorePermit permit = new SemaphorePermit(request.permits);
+ recordQueueLatency(request.getAgeNanos());
+ boolean futureCompleted = request.future.complete(permit);
+ if (!futureCompleted) {
+ // request was cancelled by user code, return permits
+ availablePermits.addAndGet(request.acquirePermits);
+ }
+ } else {
+ break;
+ }
+ }
+ }
+
+ @Override
+ public void close() {
+ if (closed.compareAndSet(false, true)) {
+ if (processQueueScheduledFuture != null) {
+ processQueueScheduledFuture.cancel(false);
+ }
+ while (!queue.isEmpty()) {
+ PendingRequest request = queue.poll();
+ request.cancelTimeoutTask();
+ request.future.completeExceptionally(new PermitAcquireAlreadyClosedException("Semaphore is closed"));
+ }
+ if (shutdownExecutor) {
+ executor.shutdownNow();
+ }
+ }
+ }
+
+ private static class PendingRequest {
+ final long permits;
+ private final long acquirePermits;
+ final CompletableFuture<AsyncSemaphorePermit> future;
+ private final BooleanSupplier isCancelled;
+ private volatile ScheduledFuture<?> timeoutTask;
+ private final long requestCreatedNanos = System.nanoTime();
+
+ PendingRequest(long permits, long acquirePermits, CompletableFuture<AsyncSemaphorePermit> future,
+ BooleanSupplier isCancelled) {
+ this.permits = permits;
+ this.acquirePermits = acquirePermits;
+ this.future = future;
+ this.isCancelled = isCancelled;
+ }
+
+ void setTimeoutTask(ScheduledFuture<?> timeoutTask) {
+ this.timeoutTask = timeoutTask;
+ }
+
+ void cancelTimeoutTask() {
+ if (timeoutTask != null) {
+ timeoutTask.cancel(false);
+ timeoutTask = null;
+ }
+ }
+
+ long getAgeNanos() {
+ return System.nanoTime() - requestCreatedNanos;
+ }
+ }
+
+ private static class SemaphorePermit implements AsyncSemaphorePermit {
+ private static final AtomicLongFieldUpdater<SemaphorePermit> PERMITS_UPDATER = AtomicLongFieldUpdater
+ .newUpdater(SemaphorePermit.class, "permits");
+ private volatile long permits;
+
+ SemaphorePermit(long permits) {
+ this.permits = permits;
+ }
+
+ @Override
+ public long getPermits() {
+ return permits;
+ }
+
+ public long releasePermits() {
+ return PERMITS_UPDATER.getAndSet(this, 0);
+ }
+ }
+}
diff --git a/pulsar-common/src/main/java/org/apache/pulsar/common/semaphore/package-info.java b/pulsar-common/src/main/java/org/apache/pulsar/common/semaphore/package-info.java
new file mode 100644
index 00000000000..6d798ac9d80
--- /dev/null
+++ b/pulsar-common/src/main/java/org/apache/pulsar/common/semaphore/package-info.java
@@ -0,0 +1,23 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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.
+ */
+
+/**
+ * Concurrent utilities for Pulsar.
+ */
+package org.apache.pulsar.common.semaphore;
diff --git a/pulsar-common/src/test/java/org/apache/pulsar/common/semaphore/AsyncDualMemoryLimiterImplTest.java b/pulsar-common/src/test/java/org/apache/pulsar/common/semaphore/AsyncDualMemoryLimiterImplTest.java
new file mode 100644
index 00000000000..fd23d95151c
--- /dev/null
+++ b/pulsar-common/src/test/java/org/apache/pulsar/common/semaphore/AsyncDualMemoryLimiterImplTest.java
@@ -0,0 +1,721 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.pulsar.common.semaphore;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.fail;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import org.apache.pulsar.common.semaphore.AsyncDualMemoryLimiter.AsyncDualMemoryLimiterPermit;
+import org.apache.pulsar.common.semaphore.AsyncDualMemoryLimiter.LimitType;
+import org.apache.pulsar.common.semaphore.AsyncSemaphore.PermitAcquireAlreadyClosedException;
+import org.apache.pulsar.common.semaphore.AsyncSemaphore.PermitAcquireCancelledException;
+import org.apache.pulsar.common.semaphore.AsyncSemaphore.PermitAcquireQueueFullException;
+import org.apache.pulsar.common.semaphore.AsyncSemaphore.PermitAcquireTimeoutException;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.Test;
+
+public class AsyncDualMemoryLimiterImplTest {
+
+ private AsyncDualMemoryLimiterImpl limiter;
+ private ScheduledExecutorService executor;
+
+ @AfterMethod(alwaysRun = true)
+ public void cleanup() {
+ if (limiter != null) {
+ limiter.close();
+ limiter = null;
+ }
+ if (executor != null) {
+ executor.shutdownNow();
+ executor = null;
+ }
+ }
+
+ @Test
+ public void testAcquireAndReleaseHeapMemory() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 10, 5000, 1000, 10, 5000);
+
+ CompletableFuture<AsyncDualMemoryLimiterPermit> future =
+ limiter.acquire(100, LimitType.HEAP_MEMORY, () -> false);
+ AsyncDualMemoryLimiterPermit permit = future.get(1, TimeUnit.SECONDS);
+
+ assertNotNull(permit);
+ assertEquals(permit.getPermits(), 100);
+ assertEquals(permit.getLimitType(), LimitType.HEAP_MEMORY);
+
+ limiter.release(permit);
+ }
+
+ @Test
+ public void testAcquireAndReleaseDirectMemory() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 10, 5000, 1000, 10, 5000);
+
+ CompletableFuture<AsyncDualMemoryLimiterPermit> future =
+ limiter.acquire(100, LimitType.DIRECT_MEMORY, () -> false);
+ AsyncDualMemoryLimiterPermit permit = future.get(1, TimeUnit.SECONDS);
+
+ assertNotNull(permit);
+ assertEquals(permit.getPermits(), 100);
+ assertEquals(permit.getLimitType(), LimitType.DIRECT_MEMORY);
+
+ limiter.release(permit);
+ }
+
+ @Test
+ public void testAcquireMultiplePermitsHeap() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 10, 5000, 1000, 10, 5000);
+
+ List<AsyncDualMemoryLimiterPermit> permits = new ArrayList<>();
+ for (int i = 0; i < 5; i++) {
+ CompletableFuture<AsyncDualMemoryLimiterPermit> future =
+ limiter.acquire(100, LimitType.HEAP_MEMORY, () -> false);
+ permits.add(future.get(1, TimeUnit.SECONDS));
+ }
+
+ assertEquals(permits.size(), 5);
+ for (AsyncDualMemoryLimiterPermit permit : permits) {
+ assertEquals(permit.getPermits(), 100);
+ assertEquals(permit.getLimitType(), LimitType.HEAP_MEMORY);
+ limiter.release(permit);
+ }
+ }
+
+ @Test
+ public void testAcquireMultiplePermitsDirect() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 10, 5000, 1000, 10, 5000);
+
+ List<AsyncDualMemoryLimiterPermit> permits = new ArrayList<>();
+ for (int i = 0; i < 5; i++) {
+ CompletableFuture<AsyncDualMemoryLimiterPermit> future =
+ limiter.acquire(100, LimitType.DIRECT_MEMORY, () -> false);
+ permits.add(future.get(1, TimeUnit.SECONDS));
+ }
+
+ assertEquals(permits.size(), 5);
+ for (AsyncDualMemoryLimiterPermit permit : permits) {
+ assertEquals(permit.getPermits(), 100);
+ assertEquals(permit.getLimitType(), LimitType.DIRECT_MEMORY);
+ limiter.release(permit);
+ }
+ }
+
+ @Test
+ public void testIndependentHeapAndDirectLimits() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 10, 5000, 1000, 10, 5000);
+
+ // Acquire all heap memory
+ AsyncDualMemoryLimiterPermit heapPermit =
+ limiter.acquire(1000, LimitType.HEAP_MEMORY, () -> false)
+ .get(1, TimeUnit.SECONDS);
+
+ // Should still be able to acquire direct memory
+ AsyncDualMemoryLimiterPermit directPermit =
+ limiter.acquire(500, LimitType.DIRECT_MEMORY, () -> false)
+ .get(1, TimeUnit.SECONDS);
+
+ assertNotNull(heapPermit);
+ assertNotNull(directPermit);
+ assertEquals(heapPermit.getLimitType(), LimitType.HEAP_MEMORY);
+ assertEquals(directPermit.getLimitType(), LimitType.DIRECT_MEMORY);
+
+ limiter.release(heapPermit);
+ limiter.release(directPermit);
+ }
+
+ @Test
+ public void testQueueingWhenHeapMemoryNotAvailable() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 10, 5000, 1000, 10, 5000);
+
+ // Acquire all heap memory
+ AsyncDualMemoryLimiterPermit permit1 =
+ limiter.acquire(1000, LimitType.HEAP_MEMORY, () -> false)
+ .get(1, TimeUnit.SECONDS);
+
+ // Try to acquire more - should be queued
+ CompletableFuture<AsyncDualMemoryLimiterPermit> future =
+ limiter.acquire(100, LimitType.HEAP_MEMORY, () -> false);
+ assertFalse(future.isDone());
+
+ // Release the first permit
+ limiter.release(permit1);
+
+ // The queued request should now complete
+ AsyncDualMemoryLimiterPermit permit2 = future.get(1, TimeUnit.SECONDS);
+ assertNotNull(permit2);
+ assertEquals(permit2.getPermits(), 100);
+
+ limiter.release(permit2);
+ }
+
+ @Test
+ public void testQueueingWhenDirectMemoryNotAvailable() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 10, 5000, 1000, 10, 5000);
+
+ // Acquire all direct memory
+ AsyncDualMemoryLimiterPermit permit1 =
+ limiter.acquire(1000, LimitType.DIRECT_MEMORY, () -> false)
+ .get(1, TimeUnit.SECONDS);
+
+ // Try to acquire more - should be queued
+ CompletableFuture<AsyncDualMemoryLimiterPermit> future =
+ limiter.acquire(100, LimitType.DIRECT_MEMORY, () -> false);
+ assertFalse(future.isDone());
+
+ // Release the first permit
+ limiter.release(permit1);
+
+ // The queued request should now complete
+ AsyncDualMemoryLimiterPermit permit2 = future.get(1, TimeUnit.SECONDS);
+ assertNotNull(permit2);
+ assertEquals(permit2.getPermits(), 100);
+
+ limiter.release(permit2);
+ }
+
+ @Test
+ public void testHeapQueueFullException() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 2, 5000, 1000, 10, 5000);
+
+ // Acquire all heap memory
+ limiter.acquire(1000, LimitType.HEAP_MEMORY, () -> false)
+ .get(1, TimeUnit.SECONDS);
+
+ // Fill the queue
+ limiter.acquire(100, LimitType.HEAP_MEMORY, () -> false);
+ limiter.acquire(100, LimitType.HEAP_MEMORY, () -> false);
+
+ // This should fail with queue full exception
+ CompletableFuture<AsyncDualMemoryLimiterPermit> future =
+ limiter.acquire(100, LimitType.HEAP_MEMORY, () -> false);
+
+ try {
+ future.get(1, TimeUnit.SECONDS);
+ fail("Expected exception");
+ } catch (ExecutionException e) {
+ assertTrue(e.getCause() instanceof PermitAcquireQueueFullException);
+ }
+ }
+
+ @Test
+ public void testDirectQueueFullException() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 10, 5000, 1000, 2, 5000);
+
+ // Acquire all direct memory
+ limiter.acquire(1000, LimitType.DIRECT_MEMORY, () -> false)
+ .get(1, TimeUnit.SECONDS);
+
+ // Fill the queue
+ limiter.acquire(100, LimitType.DIRECT_MEMORY, () -> false);
+ limiter.acquire(100, LimitType.DIRECT_MEMORY, () -> false);
+
+ // This should fail with queue full exception
+ CompletableFuture<AsyncDualMemoryLimiterPermit> future =
+ limiter.acquire(100, LimitType.DIRECT_MEMORY, () -> false);
+
+ try {
+ future.get(1, TimeUnit.SECONDS);
+ fail("Expected exception");
+ } catch (ExecutionException e) {
+ assertTrue(e.getCause() instanceof PermitAcquireQueueFullException);
+ }
+ }
+
+ @Test
+ public void testHeapTimeoutException() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 10, 100, 1000, 10, 5000);
+
+ // Acquire all heap memory
+ limiter.acquire(1000, LimitType.HEAP_MEMORY, () -> false)
+ .get(1, TimeUnit.SECONDS);
+
+ // Try to acquire more - should timeout
+ CompletableFuture<AsyncDualMemoryLimiterPermit> future =
+ limiter.acquire(100, LimitType.HEAP_MEMORY, () -> false);
+
+ try {
+ future.get(2, TimeUnit.SECONDS);
+ fail("Expected timeout exception");
+ } catch (ExecutionException e) {
+ assertTrue(e.getCause() instanceof PermitAcquireTimeoutException);
+ }
+ }
+
+ @Test
+ public void testDirectTimeoutException() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 10, 5000, 1000, 10, 100);
+
+ // Acquire all direct memory
+ limiter.acquire(1000, LimitType.DIRECT_MEMORY, () -> false)
+ .get(1, TimeUnit.SECONDS);
+
+ // Try to acquire more - should timeout
+ CompletableFuture<AsyncDualMemoryLimiterPermit> future =
+ limiter.acquire(100, LimitType.DIRECT_MEMORY, () -> false);
+
+ try {
+ future.get(2, TimeUnit.SECONDS);
+ fail("Expected timeout exception");
+ } catch (ExecutionException e) {
+ assertTrue(e.getCause() instanceof PermitAcquireTimeoutException);
+ }
+ }
+
+ @Test
+ public void testHeapCancellation() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 10, 500, 1000, 10, 500);
+
+ // Acquire all direct memory
+ CompletableFuture<AsyncDualMemoryLimiterPermit> future1 =
+ limiter.acquire(1000, LimitType.HEAP_MEMORY, () -> false);
+ AsyncDualMemoryLimiterPermit permit = future1.get(1, TimeUnit.SECONDS);
+
+ // Try to acquire more with cancellation
+ AtomicBoolean cancelled = new AtomicBoolean(false);
+ CompletableFuture<AsyncDualMemoryLimiterPermit> future2 =
+ limiter.acquire(100, LimitType.HEAP_MEMORY, cancelled::get);
+
+ assertFalse(future2.isDone());
+
+ // Cancel the request
+ cancelled.set(true);
+
+ // Release first permit
+ limiter.release(permit);
+
+ try {
+ future2.get(1, TimeUnit.SECONDS);
+ fail("Expected cancellation exception");
+ } catch (ExecutionException e) {
+ assertTrue(e.getCause() instanceof PermitAcquireCancelledException);
+ }
+ }
+
+ @Test
+ public void testDirectCancellation() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 10, 500, 1000, 10, 500);
+
+ // Acquire all direct memory
+ CompletableFuture<AsyncDualMemoryLimiterPermit> future1 =
+ limiter.acquire(1000, LimitType.DIRECT_MEMORY, () -> false);
+ AsyncDualMemoryLimiterPermit permit = future1.get(1, TimeUnit.SECONDS);
+
+ // Try to acquire more with cancellation
+ AtomicBoolean cancelled = new AtomicBoolean(false);
+ CompletableFuture<AsyncDualMemoryLimiterPermit> future2 =
+ limiter.acquire(100, LimitType.DIRECT_MEMORY, cancelled::get);
+
+ assertFalse(future2.isDone());
+
+ // Cancel the request
+ cancelled.set(true);
+
+ // Release first permit
+ limiter.release(permit);
+
+ try {
+ future2.get(1, TimeUnit.SECONDS);
+ fail("Expected cancellation exception");
+ } catch (ExecutionException e) {
+ assertTrue(e.getCause() instanceof PermitAcquireCancelledException);
+ }
+ }
+
+ @Test
+ public void testUpdateHeapPermitsIncrease() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 10, 5000, 1000, 10, 5000);
+
+ AsyncDualMemoryLimiterPermit permit =
+ limiter.acquire(100, LimitType.HEAP_MEMORY, () -> false)
+ .get(1, TimeUnit.SECONDS);
+
+ assertEquals(permit.getPermits(), 100);
+
+ AsyncDualMemoryLimiterPermit updatedPermit =
+ limiter.update(permit, 200, () -> false).get(1, TimeUnit.SECONDS);
+
+ assertEquals(updatedPermit.getPermits(), 200);
+ assertEquals(updatedPermit.getLimitType(), LimitType.HEAP_MEMORY);
+
+ limiter.release(updatedPermit);
+ }
+
+ @Test
+ public void testUpdateHeapPermitsDecrease() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 10, 5000, 1000, 10, 5000);
+
+ AsyncDualMemoryLimiterPermit permit =
+ limiter.acquire(200, LimitType.HEAP_MEMORY, () -> false)
+ .get(1, TimeUnit.SECONDS);
+
+ assertEquals(permit.getPermits(), 200);
+
+ AsyncDualMemoryLimiterPermit updatedPermit =
+ limiter.update(permit, 100, () -> false).get(1, TimeUnit.SECONDS);
+
+ assertEquals(updatedPermit.getPermits(), 100);
+ assertEquals(updatedPermit.getLimitType(), LimitType.HEAP_MEMORY);
+
+ limiter.release(updatedPermit);
+ }
+
+ @Test
+ public void testUpdateDirectPermitsIncrease() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 10, 5000, 1000, 10, 5000);
+
+ AsyncDualMemoryLimiterPermit permit =
+ limiter.acquire(100, LimitType.DIRECT_MEMORY, () -> false)
+ .get(1, TimeUnit.SECONDS);
+
+ assertEquals(permit.getPermits(), 100);
+
+ AsyncDualMemoryLimiterPermit updatedPermit =
+ limiter.update(permit, 200, () -> false).get(1, TimeUnit.SECONDS);
+
+ assertEquals(updatedPermit.getPermits(), 200);
+ assertEquals(updatedPermit.getLimitType(), LimitType.DIRECT_MEMORY);
+
+ limiter.release(updatedPermit);
+ }
+
+ @Test
+ public void testUpdateDirectPermitsDecrease() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 10, 5000, 1000, 10, 5000);
+
+ AsyncDualMemoryLimiterPermit permit =
+ limiter.acquire(200, LimitType.DIRECT_MEMORY, () -> false)
+ .get(1, TimeUnit.SECONDS);
+
+ assertEquals(permit.getPermits(), 200);
+
+ AsyncDualMemoryLimiterPermit updatedPermit =
+ limiter.update(permit, 100, () -> false).get(1, TimeUnit.SECONDS);
+
+ assertEquals(updatedPermit.getPermits(), 100);
+ assertEquals(updatedPermit.getLimitType(), LimitType.DIRECT_MEMORY);
+
+ limiter.release(updatedPermit);
+ }
+
+ @Test
+ public void testUpdateWithInvalidPermitType() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 10, 5000, 1000, 10, 5000);
+
+ // Create a mock permit with invalid type
+ AsyncDualMemoryLimiterPermit invalidPermit = new AsyncDualMemoryLimiterPermit() {
+ @Override
+ public long getPermits() {
+ return 100;
+ }
+
+ @Override
+ public LimitType getLimitType() {
+ return LimitType.HEAP_MEMORY;
+ }
+ };
+
+ try {
+ limiter.update(invalidPermit, 200, () -> false);
+ fail("Expected IllegalArgumentException");
+ } catch (IllegalArgumentException e) {
+ assertTrue(e.getMessage().contains("Invalid permit type"));
+ }
+ }
+
+ @Test
+ public void testReleaseWithInvalidPermitType() {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 10, 5000, 1000, 10, 5000);
+
+ // Create a mock permit with invalid type
+ AsyncDualMemoryLimiterPermit invalidPermit = new AsyncDualMemoryLimiterPermit() {
+ @Override
+ public long getPermits() {
+ return 100;
+ }
+
+ @Override
+ public LimitType getLimitType() {
+ return LimitType.HEAP_MEMORY;
+ }
+ };
+
+ try {
+ limiter.release(invalidPermit);
+ fail("Expected IllegalArgumentException");
+ } catch (IllegalArgumentException e) {
+ assertTrue(e.getMessage().contains("Invalid permit type"));
+ }
+ }
+
+ @Test
+ public void testConcurrentAcquireAndReleaseHeap() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 100, 5000, 1000, 100, 5000);
+
+ int numThreads = 10;
+ int numOperations = 100;
+ CountDownLatch startLatch = new CountDownLatch(1);
+ CountDownLatch doneLatch = new CountDownLatch(numThreads);
+
+ for (int i = 0; i < numThreads; i++) {
+ new Thread(() -> {
+ try {
+ startLatch.await();
+ for (int j = 0; j < numOperations; j++) {
+ AsyncDualMemoryLimiterPermit permit =
+ limiter.acquire(10, LimitType.HEAP_MEMORY, () -> false)
+ .get(5, TimeUnit.SECONDS);
+ Thread.sleep(1);
+ limiter.release(permit);
+ }
+ } catch (Exception e) {
+ fail("Unexpected exception: " + e.getMessage());
+ } finally {
+ doneLatch.countDown();
+ }
+ }).start();
+ }
+
+ startLatch.countDown();
+ assertTrue(doneLatch.await(30, TimeUnit.SECONDS));
+ }
+
+ @Test
+ public void testConcurrentAcquireAndReleaseDirect() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 100, 5000, 1000, 100, 5000);
+
+ int numThreads = 10;
+ int numOperations = 100;
+ CountDownLatch startLatch = new CountDownLatch(1);
+ CountDownLatch doneLatch = new CountDownLatch(numThreads);
+
+ for (int i = 0; i < numThreads; i++) {
+ new Thread(() -> {
+ try {
+ startLatch.await();
+ for (int j = 0; j < numOperations; j++) {
+ AsyncDualMemoryLimiterPermit permit =
+ limiter.acquire(10, LimitType.DIRECT_MEMORY, () -> false)
+ .get(5, TimeUnit.SECONDS);
+ Thread.sleep(1);
+ limiter.release(permit);
+ }
+ } catch (Exception e) {
+ fail("Unexpected exception: " + e.getMessage());
+ } finally {
+ doneLatch.countDown();
+ }
+ }).start();
+ }
+
+ startLatch.countDown();
+ assertTrue(doneLatch.await(30, TimeUnit.SECONDS));
+ }
+
+ @Test
+ public void testConcurrentMixedMemoryTypes() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 100, 5000, 1000, 100, 5000);
+
+ int numThreads = 20; // 10 for heap, 10 for direct
+ int numOperations = 50;
+ CountDownLatch startLatch = new CountDownLatch(1);
+ CountDownLatch doneLatch = new CountDownLatch(numThreads);
+
+ for (int i = 0; i < numThreads; i++) {
+ final LimitType limitType = i < 10
+ ? LimitType.HEAP_MEMORY
+ : LimitType.DIRECT_MEMORY;
+
+ new Thread(() -> {
+ try {
+ startLatch.await();
+ for (int j = 0; j < numOperations; j++) {
+ AsyncDualMemoryLimiterPermit permit =
+ limiter.acquire(10, limitType, () -> false).get(5, TimeUnit.SECONDS);
+ Thread.sleep(1);
+ limiter.release(permit);
+ }
+ } catch (Exception e) {
+ fail("Unexpected exception: " + e.getMessage());
+ } finally {
+ doneLatch.countDown();
+ }
+ }).start();
+ }
+
+ startLatch.countDown();
+ assertTrue(doneLatch.await(30, TimeUnit.SECONDS));
+ }
+
+ @Test
+ public void testCloseWithOwnExecutor() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 10, 5000, 1000, 10, 5000);
+
+ AsyncDualMemoryLimiterPermit permit =
+ limiter.acquire(100, LimitType.HEAP_MEMORY, () -> false)
+ .get(1, TimeUnit.SECONDS);
+
+ limiter.release(permit);
+ limiter.close();
+
+ // After close, acquire should fail
+ try {
+ limiter.acquire(100, LimitType.HEAP_MEMORY, () -> false)
+ .get(1, TimeUnit.SECONDS);
+ fail("Expected exception after close");
+ } catch (ExecutionException e) {
+ assertTrue(e.getCause() instanceof PermitAcquireAlreadyClosedException);
+ }
+ }
+
+ @Test
+ public void testCloseWithProvidedExecutor() throws Exception {
+ executor = Executors.newSingleThreadScheduledExecutor();
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 10, 5000, 1000, 10, 5000, executor);
+
+ AsyncDualMemoryLimiterPermit permit =
+ limiter.acquire(100, LimitType.HEAP_MEMORY, () -> false)
+ .get(1, TimeUnit.SECONDS);
+
+ limiter.release(permit);
+
+ assertFalse(executor.isShutdown());
+
+ limiter.close();
+
+ // Executor should not be shut down when provided externally
+ assertFalse(executor.isShutdown());
+
+ // After close, acquire should fail
+ try {
+ limiter.acquire(100, LimitType.HEAP_MEMORY, () -> false)
+ .get(1, TimeUnit.SECONDS);
+ fail("Expected exception after close");
+ } catch (ExecutionException e) {
+ assertTrue(e.getCause() instanceof PermitAcquireAlreadyClosedException);
+ }
+ }
+
+ @Test
+ public void testPermitGettersHeap() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 10, 5000, 1000, 10, 5000);
+
+ AsyncDualMemoryLimiterPermit permit =
+ limiter.acquire(250, LimitType.HEAP_MEMORY, () -> false)
+ .get(1, TimeUnit.SECONDS);
+
+ assertEquals(permit.getPermits(), 250);
+ assertEquals(permit.getLimitType(), LimitType.HEAP_MEMORY);
+
+ limiter.release(permit);
+ }
+
+ @Test
+ public void testPermitGettersDirect() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 10, 5000, 1000, 10, 5000);
+
+ AsyncDualMemoryLimiterPermit permit =
+ limiter.acquire(350, LimitType.DIRECT_MEMORY, () -> false)
+ .get(1, TimeUnit.SECONDS);
+
+ assertEquals(permit.getPermits(), 350);
+ assertEquals(permit.getLimitType(), LimitType.DIRECT_MEMORY);
+
+ limiter.release(permit);
+ }
+
+ @Test
+ public void testMultipleReleasesProcessQueueHeap() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 10, 5000, 1000, 10, 5000);
+
+ // Acquire permits that fill the available memory
+ List<AsyncDualMemoryLimiterPermit> permits = new ArrayList<>();
+ for (int i = 0; i < 5; i++) {
+ permits.add(limiter.acquire(200, LimitType.HEAP_MEMORY, () -> false)
+ .get(1, TimeUnit.SECONDS));
+ }
+
+ // Queue up some requests
+ List<CompletableFuture<AsyncDualMemoryLimiterPermit>> futures = new ArrayList<>();
+ for (int i = 0; i < 3; i++) {
+ futures.add(limiter.acquire(150, LimitType.HEAP_MEMORY, () -> false));
+ }
+
+ // All should be pending
+ for (CompletableFuture<AsyncDualMemoryLimiterPermit> future : futures) {
+ assertFalse(future.isDone());
+ }
+
+ // Release permits one by one
+ for (AsyncDualMemoryLimiterPermit permit : permits) {
+ limiter.release(permit);
+ }
+
+ // All queued requests should complete
+ for (CompletableFuture<AsyncDualMemoryLimiterPermit> future : futures) {
+ AsyncDualMemoryLimiterPermit permit = future.get(2, TimeUnit.SECONDS);
+ assertNotNull(permit);
+ limiter.release(permit);
+ }
+ }
+
+ @Test
+ public void testMultipleReleasesProcessQueueDirect() throws Exception {
+ limiter = new AsyncDualMemoryLimiterImpl(1000, 10, 5000, 1000, 10, 5000);
+
+ // Acquire permits that fill the available memory
+ List<AsyncDualMemoryLimiterPermit> permits = new ArrayList<>();
+ for (int i = 0; i < 5; i++) {
+ permits.add(limiter.acquire(200, LimitType.DIRECT_MEMORY, () -> false)
+ .get(1, TimeUnit.SECONDS));
+ }
+
+ // Queue up some requests
+ List<CompletableFuture<AsyncDualMemoryLimiterPermit>> futures = new ArrayList<>();
+ for (int i = 0; i < 3; i++) {
+ futures.add(limiter.acquire(150, LimitType.DIRECT_MEMORY, () -> false));
+ }
+
+ // All should be pending
+ for (CompletableFuture<AsyncDualMemoryLimiterPermit> future : futures) {
+ assertFalse(future.isDone());
+ }
+
+ // Release permits one by one
+ for (AsyncDualMemoryLimiterPermit permit : permits) {
+ limiter.release(permit);
+ }
+
+ // All queued requests should complete
+ for (CompletableFuture<AsyncDualMemoryLimiterPermit> future : futures) {
+ AsyncDualMemoryLimiterPermit permit = future.get(2, TimeUnit.SECONDS);
+ assertNotNull(permit);
+ limiter.release(permit);
+ }
+ }
+}
\ No newline at end of file
diff --git a/pulsar-common/src/test/java/org/apache/pulsar/common/semaphore/AsyncDualMemoryLimiterUtilTest.java b/pulsar-common/src/test/java/org/apache/pulsar/common/semaphore/AsyncDualMemoryLimiterUtilTest.java
new file mode 100644
index 00000000000..6c6d919bf45
--- /dev/null
+++ b/pulsar-common/src/test/java/org/apache/pulsar/common/semaphore/AsyncDualMemoryLimiterUtilTest.java
@@ -0,0 +1,581 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.pulsar.common.semaphore;
+
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.fail;
+import io.netty.buffer.ByteBuf;
+import io.netty.channel.ChannelFuture;
+import io.netty.channel.ChannelHandlerContext;
+import java.time.Duration;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import org.apache.pulsar.common.api.proto.BaseCommand;
+import org.apache.pulsar.common.semaphore.AsyncDualMemoryLimiter.AsyncDualMemoryLimiterPermit;
+import org.apache.pulsar.common.semaphore.AsyncDualMemoryLimiter.LimitType;
+import org.apache.pulsar.common.semaphore.AsyncSemaphore.PermitAcquireCancelledException;
+import org.awaitility.Awaitility;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+public class AsyncDualMemoryLimiterUtilTest {
+
+ private AsyncDualMemoryLimiterImpl limiter;
+
+ @BeforeMethod
+ public void setup() {
+ limiter = new AsyncDualMemoryLimiterImpl(10000, 100, 5000, 10000, 100, 5000);
+ }
+
+ @AfterMethod(alwaysRun = true)
+ public void cleanup() {
+ if (limiter != null) {
+ limiter.close();
+ limiter = null;
+ }
+ }
+
+ @Test
+ public void testWithPermitsFutureSuccessfulExecution() throws Exception {
+ CompletableFuture<AsyncDualMemoryLimiterPermit> permitFuture =
+ limiter.acquire(100, LimitType.HEAP_MEMORY, () -> false);
+
+ AtomicBoolean functionExecuted = new AtomicBoolean(false);
+ AtomicBoolean permitReleased = new AtomicBoolean(false);
+
+ CompletableFuture<String> result = AsyncDualMemoryLimiterUtil.withPermitsFuture(
+ permitFuture,
+ permit -> {
+ functionExecuted.set(true);
+ return CompletableFuture.completedFuture("success");
+ },
+ throwable -> CompletableFuture.completedFuture("error"),
+ permit -> permitReleased.set(true)
+ );
+
+ String value = result.get(1, TimeUnit.SECONDS);
+ assertEquals(value, "success");
+ assertTrue(functionExecuted.get());
+ assertTrue(permitReleased.get());
+ }
+
+ @Test
+ public void testWithPermitsFuturePermitAcquireError() throws Exception {
+ CompletableFuture<AsyncDualMemoryLimiterPermit> permitFuture =
+ CompletableFuture.failedFuture(new RuntimeException("Permit acquire failed"));
+
+ AtomicBoolean functionExecuted = new AtomicBoolean(false);
+ AtomicBoolean errorHandlerExecuted = new AtomicBoolean(false);
+ AtomicBoolean permitReleased = new AtomicBoolean(false);
+ AtomicReference<Throwable> capturedError = new AtomicReference<>();
+
+ CompletableFuture<String> result = AsyncDualMemoryLimiterUtil.withPermitsFuture(
+ permitFuture,
+ permit -> {
+ functionExecuted.set(true);
+ return CompletableFuture.completedFuture("success");
+ },
+ throwable -> {
+ errorHandlerExecuted.set(true);
+ capturedError.set(throwable);
+ return CompletableFuture.completedFuture("handled-error");
+ },
+ permit -> permitReleased.set(true)
+ );
+
+ String value = result.get(1, TimeUnit.SECONDS);
+ assertEquals(value, "handled-error");
+ assertTrue(errorHandlerExecuted.get());
+ assertTrue(capturedError.get() instanceof RuntimeException);
+ assertEquals(capturedError.get().getMessage(), "Permit acquire failed");
+ assertTrue(!functionExecuted.get()); // Function should not be executed
+ assertTrue(!permitReleased.get()); // Permits should not be released
+ }
+
+ @Test
+ public void testWithPermitsFutureFunctionThrowsException() throws Exception {
+ CompletableFuture<AsyncDualMemoryLimiterPermit> permitFuture =
+ limiter.acquire(100, LimitType.HEAP_MEMORY, () -> false);
+
+ AtomicBoolean permitReleased = new AtomicBoolean(false);
+
+ CompletableFuture<String> result = AsyncDualMemoryLimiterUtil.withPermitsFuture(
+ permitFuture,
+ permit -> {
+ throw new RuntimeException("Function failed");
+ },
+ throwable -> CompletableFuture.completedFuture("error"),
+ permit -> permitReleased.set(true)
+ );
+
+ try {
+ result.get(1, TimeUnit.SECONDS);
+ fail("Expected exception");
+ } catch (ExecutionException e) {
+ assertTrue(e.getCause() instanceof RuntimeException);
+ assertEquals(e.getCause().getMessage(), "Function failed");
+ }
+
+ // Permits should be released even when function throws
+ assertTrue(permitReleased.get());
+ }
+
+ @Test
+ public void testWithPermitsFutureFunctionReturnsFailedFuture() throws Exception {
+ CompletableFuture<AsyncDualMemoryLimiterPermit> permitFuture =
+ limiter.acquire(100, LimitType.HEAP_MEMORY, () -> false);
+
+ AtomicBoolean permitReleased = new AtomicBoolean(false);
+
+ CompletableFuture<String> result = AsyncDualMemoryLimiterUtil.withPermitsFuture(
+ permitFuture,
+ permit -> CompletableFuture.failedFuture(new RuntimeException("Async operation failed")),
+ throwable -> CompletableFuture.completedFuture("error"),
+ permit -> permitReleased.set(true)
+ );
+
+ try {
+ result.get(1, TimeUnit.SECONDS);
+ fail("Expected exception");
+ } catch (ExecutionException e) {
+ assertTrue(e.getCause() instanceof RuntimeException);
+ assertEquals(e.getCause().getMessage(), "Async operation failed");
+ }
+
+ // Permits should be released after the returned future completes
+ assertTrue(permitReleased.get());
+ }
+
+ @Test
+ public void testWithPermitsFutureReleasesPermitsAfterAsyncCompletion() throws Exception {
+ CompletableFuture<AsyncDualMemoryLimiterPermit> permitFuture =
+ limiter.acquire(100, LimitType.HEAP_MEMORY, () -> false);
+
+ CompletableFuture<String> asyncOperation = new CompletableFuture<>();
+ AtomicBoolean permitReleased = new AtomicBoolean(false);
+
+ CompletableFuture<String> result = AsyncDualMemoryLimiterUtil.withPermitsFuture(
+ permitFuture,
+ permit -> asyncOperation,
+ throwable -> CompletableFuture.completedFuture("error"),
+ permit -> permitReleased.set(true)
+ );
+
+ // Permits should not be released yet
+ Thread.sleep(100);
+ assertTrue(!permitReleased.get());
+ assertTrue(!result.isDone());
+
+ // Complete the async operation
+ asyncOperation.complete("async-result");
+
+ String value = result.get(1, TimeUnit.SECONDS);
+ assertEquals(value, "async-result");
+ assertTrue(permitReleased.get());
+ }
+
+ @Test
+ public void testWithPermitsFutureWithNullPermit() throws Exception {
+ // Simulate a null permit scenario (edge case)
+ CompletableFuture<AsyncDualMemoryLimiterPermit> permitFuture =
+ CompletableFuture.completedFuture(null);
+
+ AtomicBoolean functionExecuted = new AtomicBoolean(false);
+ AtomicBoolean permitReleased = new AtomicBoolean(false);
+
+ CompletableFuture<String> result = AsyncDualMemoryLimiterUtil.withPermitsFuture(
+ permitFuture,
+ permit -> {
+ functionExecuted.set(true);
+ return CompletableFuture.completedFuture("success");
+ },
+ throwable -> CompletableFuture.completedFuture("error"),
+ permit -> permitReleased.set(true)
+ );
+
+ String value = result.get(1, TimeUnit.SECONDS);
+ assertEquals(value, "success");
+ assertTrue(functionExecuted.get());
+ assertTrue(permitReleased.get());
+ }
+
+ @Test
+ public void testAcquireDirectMemoryPermitsAndWriteAndFlushSuccess() throws Exception {
+ ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
+ ChannelFuture channelFuture = mock(ChannelFuture.class);
+ BaseCommand command = createTestCommand();
+
+ AtomicReference<ByteBuf> capturedBuf = new AtomicReference<>();
+ AtomicBoolean listenerCalled = new AtomicBoolean(false);
+
+ when(ctx.writeAndFlush(any(ByteBuf.class))).thenAnswer(invocation -> {
+ capturedBuf.set(invocation.getArgument(0));
+ return channelFuture;
+ });
+
+ when(channelFuture.addListener(any())).thenAnswer(invocation -> {
+ io.netty.util.concurrent.GenericFutureListener listener = invocation.getArgument(0);
+ listenerCalled.set(true);
+ // Simulate successful write
+ listener.operationComplete(channelFuture);
+ return channelFuture;
+ });
+
+ when(channelFuture.isSuccess()).thenReturn(true);
+
+ AtomicBoolean errorHandlerCalled = new AtomicBoolean(false);
+
+ CompletableFuture<Void> result = AsyncDualMemoryLimiterUtil.acquireDirectMemoryPermitsAndWriteAndFlush(
+ ctx,
+ limiter,
+ () -> false,
+ command,
+ throwable -> errorHandlerCalled.set(true)
+ );
+
+ result.get(1, TimeUnit.SECONDS);
+
+ verify(ctx, times(1)).writeAndFlush(any(ByteBuf.class));
+ assertTrue(listenerCalled.get());
+ assertTrue(!errorHandlerCalled.get());
+ assertNotNull(capturedBuf.get());
+ }
+
+ @Test
+ public void testAcquireDirectMemoryPermitsAndWriteAndFlushPermitAcquireError() throws Exception {
+ // Fill up the limiter to cause permit acquisition to fail
+ limiter.close();
+ limiter = new AsyncDualMemoryLimiterImpl(100, 1, 100, 100, 1, 100);
+
+ // Acquire all memory
+ limiter.acquire(100, LimitType.DIRECT_MEMORY, () -> false).get(1, TimeUnit.SECONDS);
+
+ // Try to acquire more - will be queued and timeout
+ ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
+ BaseCommand command = createTestCommand();
+
+ AtomicBoolean errorHandlerCalled = new AtomicBoolean(false);
+ AtomicReference<Throwable> capturedError = new AtomicReference<>();
+
+ CompletableFuture<Void> result = AsyncDualMemoryLimiterUtil.acquireDirectMemoryPermitsAndWriteAndFlush(
+ ctx,
+ limiter,
+ () -> false,
+ command,
+ throwable -> {
+ errorHandlerCalled.set(true);
+ capturedError.set(throwable);
+ }
+ );
+
+ // Wait for timeout
+ Thread.sleep(200);
+
+ assertTrue(errorHandlerCalled.get());
+ assertNotNull(capturedError.get());
+ verify(ctx, never()).writeAndFlush(any(ByteBuf.class));
+ }
+
+ @Test
+ public void testAcquireDirectMemoryPermitsAndWriteAndFlushSerializationError() throws Exception {
+ ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
+ BaseCommand command = mock(BaseCommand.class);
+
+ // Mock the command to throw exception during serialization
+ when(command.getSerializedSize()).thenReturn(100);
+
+ AtomicBoolean errorHandlerCalled = new AtomicBoolean(false);
+
+ CompletableFuture<Void> result = AsyncDualMemoryLimiterUtil.acquireDirectMemoryPermitsAndWriteAndFlush(
+ ctx,
+ limiter,
+ () -> false,
+ command,
+ throwable -> errorHandlerCalled.set(true)
+ );
+
+ try {
+ result.get(1, TimeUnit.SECONDS);
+ fail("Expected exception");
+ } catch (ExecutionException e) {
+ // Expected exception from serialization
+ assertTrue(e.getCause() != null);
+ }
+ }
+
+ @Test
+ public void testAcquireDirectMemoryPermitsAndWriteAndFlushWriteException() throws Exception {
+ ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
+ BaseCommand command = createTestCommand();
+
+ // Mock writeAndFlush to throw exception
+ when(ctx.writeAndFlush(any(ByteBuf.class)))
+ .thenThrow(new RuntimeException("Write failed"));
+
+ AtomicBoolean errorHandlerCalled = new AtomicBoolean(false);
+
+ CompletableFuture<Void> result = AsyncDualMemoryLimiterUtil.acquireDirectMemoryPermitsAndWriteAndFlush(
+ ctx,
+ limiter,
+ () -> false,
+ command,
+ throwable -> errorHandlerCalled.set(true)
+ );
+
+ try {
+ result.get(1, TimeUnit.SECONDS);
+ fail("Expected exception");
+ } catch (ExecutionException e) {
+ assertTrue(e.getCause() instanceof RuntimeException);
+ assertEquals(e.getCause().getMessage(), "Write failed");
+ }
+ }
+
+ @Test
+ public void testAcquireDirectMemoryPermitsAndWriteAndFlushCancelled() throws Exception {
+ // Fill up the limiter so permit acquisition will be queued
+ limiter.close();
+ limiter = new AsyncDualMemoryLimiterImpl(100, 10, 5000, 100, 10, 5000);
+ AsyncDualMemoryLimiterPermit permits =
+ limiter.acquire(100, LimitType.DIRECT_MEMORY, () -> false).get(1, TimeUnit.SECONDS);
+
+ ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
+ BaseCommand command = createTestCommand();
+
+ AtomicBoolean cancelled = new AtomicBoolean(false);
+ AtomicBoolean errorHandlerCalled = new AtomicBoolean(false);
+
+ CompletableFuture<Void> result = AsyncDualMemoryLimiterUtil.acquireDirectMemoryPermitsAndWriteAndFlush(
+ ctx,
+ limiter,
+ cancelled::get,
+ command,
+ throwable -> errorHandlerCalled.set(true)
+ );
+
+ // Cancel the request
+ cancelled.set(true);
+
+ limiter.release(permits);
+
+ Awaitility.await().untilAsserted(() -> assertTrue(errorHandlerCalled.get()));
+ verify(ctx, never()).writeAndFlush(any(ByteBuf.class));
+ assertTrue(result.isCompletedExceptionally());
+ assertThatThrownBy(() -> result.get()).hasCauseInstanceOf(PermitAcquireCancelledException.class);
+ }
+
+ @Test
+ public void testAcquireDirectMemoryPermitsAndWriteAndFlushMultipleConcurrent() throws Exception {
+ ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
+ ChannelFuture channelFuture = mock(ChannelFuture.class);
+
+ when(ctx.writeAndFlush(any(ByteBuf.class))).thenReturn(channelFuture);
+ when(channelFuture.addListener(any())).thenAnswer(invocation -> {
+ io.netty.util.concurrent.GenericFutureListener listener = invocation.getArgument(0);
+ listener.operationComplete(channelFuture);
+ return channelFuture;
+ });
+ when(channelFuture.isSuccess()).thenReturn(true);
+
+ int numRequests = 10;
+ CompletableFuture<Void>[] futures = new CompletableFuture[numRequests];
+
+ for (int i = 0; i < numRequests; i++) {
+ BaseCommand command = createTestCommand();
+ futures[i] = AsyncDualMemoryLimiterUtil.acquireDirectMemoryPermitsAndWriteAndFlush(
+ ctx,
+ limiter,
+ () -> false,
+ command,
+ throwable -> {
+ }
+ );
+ }
+
+ // Wait for all to complete
+ CompletableFuture.allOf(futures).get(5, TimeUnit.SECONDS);
+
+ verify(ctx, times(numRequests)).writeAndFlush(any(ByteBuf.class));
+ }
+
+ @Test
+ public void testAcquireDirectMemoryPermitsAndWriteAndFlushReleasesOnWriteFailure() throws Exception {
+ ChannelHandlerContext ctx = mock(ChannelHandlerContext.class);
+ ChannelFuture channelFuture = mock(ChannelFuture.class);
+ BaseCommand command = createTestCommand();
+
+ when(ctx.writeAndFlush(any(ByteBuf.class))).thenReturn(channelFuture);
+ when(channelFuture.addListener(any())).thenAnswer(invocation -> {
+ io.netty.util.concurrent.GenericFutureListener listener = invocation.getArgument(0);
+ // Simulate write failure
+ when(channelFuture.isSuccess()).thenReturn(false);
+ when(channelFuture.cause()).thenReturn(new RuntimeException("Write to socket failed"));
+ listener.operationComplete(channelFuture);
+ return channelFuture;
+ });
+
+ AtomicBoolean errorHandlerCalled = new AtomicBoolean(false);
+
+ CompletableFuture<Void> result = AsyncDualMemoryLimiterUtil.acquireDirectMemoryPermitsAndWriteAndFlush(
+ ctx,
+ limiter,
+ () -> false,
+ command,
+ throwable -> errorHandlerCalled.set(true)
+ );
+
+ result.get(1, TimeUnit.SECONDS);
+
+ verify(ctx, times(1)).writeAndFlush(any(ByteBuf.class));
+ // Error handler should not be called for write failures (only for permit acquire failures)
+ assertTrue(!errorHandlerCalled.get());
+ }
+
+ @Test
+ public void testWithPermitsFutureMultipleConcurrent() throws Exception {
+ int numOperations = 20;
+ CompletableFuture<String>[] futures = new CompletableFuture[numOperations];
+ AtomicInteger releaseCount = new AtomicInteger(0);
+
+ for (int i = 0; i < numOperations; i++) {
+ final int index = i;
+ CompletableFuture<AsyncDualMemoryLimiterPermit> permitFuture =
+ limiter.acquire(50, LimitType.HEAP_MEMORY, () -> false);
+
+ futures[i] = AsyncDualMemoryLimiterUtil.withPermitsFuture(
+ permitFuture,
+ permit -> CompletableFuture.completedFuture("result-" + index),
+ throwable -> CompletableFuture.completedFuture("error-" + index),
+ permit -> releaseCount.incrementAndGet()
+ );
+ }
+
+ CompletableFuture.allOf(futures).get(5, TimeUnit.SECONDS);
+
+ assertEquals(releaseCount.get(), numOperations);
+
+ for (int i = 0; i < numOperations; i++) {
+ assertEquals(futures[i].get(), "result-" + i);
+ }
+ }
+
+ @Test
+ public void testWithAcquiredAndUpdatedPermitsDoesntLeakPermitsWhenUpdatedSizeIsLarger() throws Exception {
+ for (int i = 0; i < 10_000; i++) {
+ doTestWithAcquiredPermitsAndUpdatedPermits(50, 100);
+ }
+ assertThat(limiter.getLimiter(LimitType.HEAP_MEMORY).getAcquiredPermits()).isEqualTo(0);
+ assertThat(limiter.getLimiter(LimitType.HEAP_MEMORY).getAvailablePermits()).isEqualTo(10000);
+ }
+
+ @Test
+ public void testWithAcquiredAndUpdatedPermitsDoesntLeakPermitsWhenUpdatedSizeIsSmaller() throws Exception {
+ for (int i = 0; i < 10_000; i++) {
+ doTestWithAcquiredPermitsAndUpdatedPermits(100, 50);
+ }
+ assertThat(limiter.getLimiter(LimitType.HEAP_MEMORY).getAcquiredPermits()).isEqualTo(0);
+ assertThat(limiter.getLimiter(LimitType.HEAP_MEMORY).getAvailablePermits()).isEqualTo(10000);
+ }
+
+ private void doTestWithAcquiredPermitsAndUpdatedPermits(int initialMemorySize, int newMemorySize) {
+ CompletableFuture<Void> result = limiter.withAcquiredPermits(
+ initialMemorySize, // estimated permits
+ LimitType.HEAP_MEMORY,
+ () -> false,
+ firstPermit -> {
+ // Simulate getting actual size and updating permits
+ return limiter.withUpdatedPermits(
+ firstPermit,
+ newMemorySize,
+ () -> false,
+ secondPermit -> {
+ return CompletableFuture.completedFuture(null);
+ },
+ throwable -> CompletableFuture.failedFuture(throwable)
+ );
+ },
+ throwable -> CompletableFuture.failedFuture(throwable)
+ );
+ assertThat(result).succeedsWithin(Duration.ofSeconds(1));
+ }
+
+ @Test
+ public void testWithAcquiredPermitsDoesntLeakPermitsWhenExceptionIsThrown() throws Exception {
+ assertThat(limiter.getLimiter(LimitType.HEAP_MEMORY).getAvailablePermits()).isEqualTo(10000);
+ CompletableFuture<Void> result = limiter.withAcquiredPermits(
+ 90, // estimated permits
+ LimitType.HEAP_MEMORY,
+ () -> false,
+ firstPermit -> {
+ throw new RuntimeException("Exception in withAcquiredPermits");
+ },
+ throwable -> CompletableFuture.failedFuture(throwable)
+ );
+ assertThat(result).failsWithin(Duration.ofSeconds(1));
+ assertThat(limiter.getLimiter(LimitType.HEAP_MEMORY).getAcquiredPermits()).isEqualTo(0);
+ assertThat(limiter.getLimiter(LimitType.HEAP_MEMORY).getAvailablePermits()).isEqualTo(10000);
+ }
+
+ @Test
+ public void testWithAcquiredAndUpdatedPermitsDoesntLeakPermitsWhenExceptionIsThrown() throws Exception {
+ assertThat(limiter.getLimiter(LimitType.HEAP_MEMORY).getAvailablePermits()).isEqualTo(10000);
+ CompletableFuture<Void> result = limiter.withAcquiredPermits(
+ 90, // estimated permits
+ LimitType.HEAP_MEMORY,
+ () -> false,
+ firstPermit -> {
+ // Simulate getting actual size and updating permits
+ return limiter.withUpdatedPermits(
+ firstPermit,
+ 100,
+ () -> false,
+ secondPermit -> {
+ throw new RuntimeException("Exception in withUpdatedPermits");
+ },
+ throwable -> CompletableFuture.failedFuture(throwable)
+ );
+ },
+ throwable -> CompletableFuture.failedFuture(throwable)
+ );
+ assertThat(result).failsWithin(Duration.ofSeconds(1));
+ assertThat(limiter.getLimiter(LimitType.HEAP_MEMORY).getAcquiredPermits()).isEqualTo(0);
+ assertThat(limiter.getLimiter(LimitType.HEAP_MEMORY).getAvailablePermits()).isEqualTo(10000);
+ }
+
+ private BaseCommand createTestCommand() {
+ BaseCommand command = new BaseCommand().setType(BaseCommand.Type.PING);
+ command.setPing();
+ return command;
+ }
+}
\ No newline at end of file
diff --git a/pulsar-common/src/test/java/org/apache/pulsar/common/semaphore/AsyncSemaphoreImplTest.java b/pulsar-common/src/test/java/org/apache/pulsar/common/semaphore/AsyncSemaphoreImplTest.java
new file mode 100644
index 00000000000..ec2e01e1a71
--- /dev/null
+++ b/pulsar-common/src/test/java/org/apache/pulsar/common/semaphore/AsyncSemaphoreImplTest.java
@@ -0,0 +1,669 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.pulsar.common.semaphore;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.fail;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import org.apache.pulsar.common.semaphore.AsyncSemaphore.AsyncSemaphorePermit;
+import org.apache.pulsar.common.semaphore.AsyncSemaphore.PermitAcquireAlreadyClosedException;
+import org.apache.pulsar.common.semaphore.AsyncSemaphore.PermitAcquireCancelledException;
+import org.apache.pulsar.common.semaphore.AsyncSemaphore.PermitAcquireQueueFullException;
+import org.apache.pulsar.common.semaphore.AsyncSemaphore.PermitAcquireTimeoutException;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.Test;
+
+public class AsyncSemaphoreImplTest {
+
+ private AsyncSemaphoreImpl semaphore;
+
+ @AfterMethod(alwaysRun = true)
+ public void cleanup() {
+ if (semaphore != null) {
+ semaphore.close();
+ semaphore = null;
+ }
+ }
+
+ @Test
+ public void testAcquireAndReleaseSinglePermit() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(1, 10, 5000);
+
+ CompletableFuture<AsyncSemaphorePermit> future = semaphore.acquire(1, () -> false);
+ AsyncSemaphorePermit permit = future.get(1, TimeUnit.SECONDS);
+
+ assertNotNull(permit);
+ assertEquals(permit.getPermits(), 1);
+
+ semaphore.release(permit);
+ }
+
+ @Test
+ public void testAcquireMultiplePermits() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(10, 10, 5000);
+
+ CompletableFuture<AsyncSemaphorePermit> future = semaphore.acquire(5, () -> false);
+ AsyncSemaphorePermit permit = future.get(1, TimeUnit.SECONDS);
+
+ assertNotNull(permit);
+ assertEquals(permit.getPermits(), 5);
+
+ semaphore.release(permit);
+ }
+
+ @Test
+ public void testQueueingWhenNoPermitsAvailable() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(5, 10, 5000);
+
+ // Acquire all permits
+ CompletableFuture<AsyncSemaphorePermit> future1 = semaphore.acquire(5, () -> false);
+ AsyncSemaphorePermit permit1 = future1.get(1, TimeUnit.SECONDS);
+
+ // Try to acquire more - should be queued
+ CompletableFuture<AsyncSemaphorePermit> future2 = semaphore.acquire(3, () -> false);
+ assertFalse(future2.isDone());
+
+ // Release permits
+ semaphore.release(permit1);
+
+ // Now the queued request should complete
+ AsyncSemaphorePermit permit2 = future2.get(1, TimeUnit.SECONDS);
+ assertNotNull(permit2);
+ assertEquals(permit2.getPermits(), 3);
+
+ semaphore.release(permit2);
+ }
+
+ @Test
+ public void testQueueFullException() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(1, 2, 5000);
+
+ // Acquire the only permit
+ CompletableFuture<AsyncSemaphorePermit> future1 = semaphore.acquire(1, () -> false);
+ future1.get(1, TimeUnit.SECONDS);
+
+ // Fill the queue
+ CompletableFuture<AsyncSemaphorePermit> future2 = semaphore.acquire(1, () -> false);
+ CompletableFuture<AsyncSemaphorePermit> future3 = semaphore.acquire(1, () -> false);
+
+ // This should fail with queue full
+ CompletableFuture<AsyncSemaphorePermit> future4 = semaphore.acquire(1, () -> false);
+
+ assertTrue(future4.isCompletedExceptionally());
+ try {
+ future4.get();
+ fail("Should have thrown exception");
+ } catch (ExecutionException e) {
+ assertTrue(e.getCause() instanceof PermitAcquireQueueFullException);
+ }
+ }
+
+ @Test
+ public void testTimeoutException() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(1, 10, 100); // 100ms timeout
+
+ // Acquire the only permit
+ CompletableFuture<AsyncSemaphorePermit> future1 = semaphore.acquire(1, () -> false);
+ AsyncSemaphorePermit permit1 = future1.get(1, TimeUnit.SECONDS);
+
+ // Try to acquire another permit - should timeout
+ CompletableFuture<AsyncSemaphorePermit> future2 = semaphore.acquire(1, () -> false);
+
+ Thread.sleep(200); // Wait for timeout
+
+ assertTrue(future2.isCompletedExceptionally());
+ try {
+ future2.get();
+ fail("Should have thrown exception");
+ } catch (ExecutionException e) {
+ assertTrue(e.getCause() instanceof PermitAcquireTimeoutException);
+ }
+
+ semaphore.release(permit1);
+ }
+
+ @Test
+ public void testCancellation() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(1, 10, 5000);
+
+ // Acquire the only permit
+ CompletableFuture<AsyncSemaphorePermit> future1 = semaphore.acquire(1, () -> false);
+ AsyncSemaphorePermit permit1 = future1.get(1, TimeUnit.SECONDS);
+
+ // Try to acquire with cancellation flag
+ AtomicBoolean cancelled = new AtomicBoolean(false);
+ CompletableFuture<AsyncSemaphorePermit> future2 = semaphore.acquire(1, cancelled::get);
+
+ assertFalse(future2.isDone());
+
+ // Cancel the request
+ cancelled.set(true);
+
+ // Release the first permit
+ semaphore.release(permit1);
+
+ // Give time for the queue to process
+ Thread.sleep(100);
+
+ assertTrue(future2.isCompletedExceptionally());
+ try {
+ future2.get();
+ fail("Should have thrown exception");
+ } catch (ExecutionException e) {
+ assertTrue(e.getCause() instanceof PermitAcquireCancelledException);
+ }
+ }
+
+ @Test
+ public void testInvalidPermits() {
+ semaphore = new AsyncSemaphoreImpl(10, 10, 5000);
+
+ assertThatThrownBy(() ->
+ semaphore.acquire(-1, () -> false)
+ ).isInstanceOf(IllegalArgumentException.class)
+ .hasMessage("Invalid negative permits value: -1");
+ }
+
+ @Test
+ public void testInvalidPermitsExceedingMaxPermits() {
+ semaphore = new AsyncSemaphoreImpl(10, 10, 5000);
+
+ assertThatThrownBy(() ->
+ semaphore.acquire(11, () -> false)
+ ).isInstanceOf(IllegalArgumentException.class)
+ .hasMessage("Requested permits=11 is larger than maxPermits=10");
+ }
+
+ @Test
+ public void testClose() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(5, 10, 5000);
+
+ // Acquire some permits
+ CompletableFuture<AsyncSemaphorePermit> future1 = semaphore.acquire(3, () -> false);
+ future1.get(1, TimeUnit.SECONDS);
+
+ // Queue another request
+ CompletableFuture<AsyncSemaphorePermit> future2 = semaphore.acquire(5, () -> false);
+ assertFalse(future2.isDone());
+
+ // Close the semaphore
+ semaphore.close();
+
+ // Queued request should fail
+ assertTrue(future2.isCompletedExceptionally());
+ try {
+ future2.get();
+ fail("Should have thrown exception");
+ } catch (ExecutionException e) {
+ assertTrue(e.getCause() instanceof PermitAcquireAlreadyClosedException);
+ }
+
+ // New acquisitions should fail
+ CompletableFuture<AsyncSemaphorePermit> future3 = semaphore.acquire(1, () -> false);
+ assertTrue(future3.isCompletedExceptionally());
+ try {
+ future3.get();
+ fail("Should have thrown exception");
+ } catch (ExecutionException e) {
+ assertTrue(e.getCause() instanceof PermitAcquireAlreadyClosedException);
+ }
+ }
+
+ @Test
+ public void testUpdatePermitsIncrease() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(10, 10, 5000);
+
+ // Acquire initial permits
+ CompletableFuture<AsyncSemaphorePermit> future1 = semaphore.acquire(5, () -> false);
+ AsyncSemaphorePermit permit1 = future1.get(1, TimeUnit.SECONDS);
+ assertEquals(permit1.getPermits(), 5);
+
+ // Update to more permits
+ CompletableFuture<AsyncSemaphorePermit> future2 = semaphore.update(permit1, 8, () -> false);
+ AsyncSemaphorePermit permit2 = future2.get(1, TimeUnit.SECONDS);
+
+ assertNotNull(permit2);
+ assertEquals(permit2.getPermits(), 8);
+
+ CompletableFuture<AsyncSemaphorePermit> future3 = semaphore.acquire(2, () -> false);
+ AsyncSemaphorePermit permit3 = future3.get(1, TimeUnit.SECONDS);
+ assertNotNull(permit3);
+ assertEquals(permit3.getPermits(), 2);
+
+ CompletableFuture<AsyncSemaphorePermit> future4 = semaphore.acquire(1, () -> false);
+ Thread.sleep(1000);
+ // no more permits available, this won't complete
+ assertThat(future4).isNotDone();
+
+ // release permits
+ semaphore.release(permit2);
+
+ // now future4 should complete
+ assertThat(future4).succeedsWithin(1, TimeUnit.SECONDS)
+ .satisfies(p -> assertEquals(p.getPermits(), 1));
+ }
+
+ @Test
+ public void testUpdatePermitsDecrease() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(10, 10, 5000);
+
+ // Acquire initial permits
+ CompletableFuture<AsyncSemaphorePermit> future1 = semaphore.acquire(8, () -> false);
+ AsyncSemaphorePermit permit1 = future1.get(1, TimeUnit.SECONDS);
+ assertEquals(permit1.getPermits(), 8);
+
+ // Update to fewer permits (should be immediate)
+ CompletableFuture<AsyncSemaphorePermit> future2 = semaphore.update(permit1, 5, () -> false);
+ AsyncSemaphorePermit permit2 = future2.get(100, TimeUnit.MILLISECONDS);
+
+ assertNotNull(permit2);
+ assertEquals(permit2.getPermits(), 5);
+ assertTrue(future2.isDone());
+
+ semaphore.release(permit2);
+ }
+
+ @Test
+ public void testUpdateWithInvalidPermits() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(10, 10, 5000);
+
+ CompletableFuture<AsyncSemaphorePermit> future = semaphore.acquire(5, () -> false);
+ AsyncSemaphorePermit permit = future.get(1, TimeUnit.SECONDS);
+
+ assertThatThrownBy(() ->
+ semaphore.update(permit, -1, () -> false)
+ ).isInstanceOf(IllegalArgumentException.class)
+ .hasMessage("Invalid negative permits value: -1");
+
+ semaphore.release(permit);
+ }
+
+ @Test
+ public void testUpdateWithInvalidPermitsExceedingMaxPermits() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(10, 10, 5000);
+
+ CompletableFuture<AsyncSemaphorePermit> future = semaphore.acquire(5, () -> false);
+ AsyncSemaphorePermit permit = future.get(1, TimeUnit.SECONDS);
+
+ assertThatThrownBy(() ->
+ semaphore.update(permit, 11, () -> false)
+ ).isInstanceOf(IllegalArgumentException.class)
+ .hasMessage("Requested permits=11 is larger than maxPermits=10");
+
+ semaphore.release(permit);
+ }
+
+ @Test
+ public void testUpdateWithZeroPermitsShouldReleasePermits() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(10, 10, 5000);
+
+ CompletableFuture<AsyncSemaphorePermit> future = semaphore.acquire(5, () -> false);
+ AsyncSemaphorePermit permit = future.get(1, TimeUnit.SECONDS);
+
+ AsyncSemaphorePermit updatedPermit = semaphore.update(permit, 0, () -> false).get(1, TimeUnit.SECONDS);
+
+ AsyncSemaphorePermit permit2 = semaphore.acquire(10, () -> false).get(1, TimeUnit.SECONDS);
+
+ semaphore.release(updatedPermit);
+ semaphore.release(permit2);
+ }
+
+ @Test
+ public void testConcurrentAcquireAndRelease() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(10, 100, 5000);
+
+ List<CompletableFuture<Void>> futures = new ArrayList<>();
+
+ // Start multiple threads acquiring and releasing permits
+ for (int i = 0; i < 20; i++) {
+ CompletableFuture<Void> future = CompletableFuture.runAsync(() -> {
+ try {
+ for (int j = 0; j < 10; j++) {
+ CompletableFuture<AsyncSemaphorePermit> permitFuture =
+ semaphore.acquire(1, () -> false);
+ AsyncSemaphorePermit permit = permitFuture.get(5, TimeUnit.SECONDS);
+ Thread.sleep(10); // Hold permit briefly
+ semaphore.release(permit);
+ }
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ });
+ futures.add(future);
+ }
+
+ // Wait for all threads to complete
+ CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]))
+ .get(30, TimeUnit.SECONDS);
+ }
+
+ @Test
+ public void testTimeoutProcessesNextRequest() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(5, 10, 250);
+
+ // Acquire all permits
+ CompletableFuture<AsyncSemaphorePermit> future1 = semaphore.acquire(5, () -> false);
+ AsyncSemaphorePermit permit1 = future1.get(1, TimeUnit.SECONDS);
+
+ // Request that will timeout (needs more permits than available)
+ CompletableFuture<AsyncSemaphorePermit> future2 = semaphore.acquire(5, () -> false);
+
+ // make requests 250ms apart
+ Thread.sleep(250);
+
+ // Request that can be satisfied
+ CompletableFuture<AsyncSemaphorePermit> future3 = semaphore.acquire(3, () -> false);
+
+ // Release permits
+ semaphore.release(permit1);
+
+ // Expect the second request to timeout
+ Throwable throwable = future2.handle((permit, t) -> t).join();
+ assertThat(throwable).isInstanceOf(PermitAcquireTimeoutException.class);
+
+ // The third request should have succeeded now
+ AsyncSemaphorePermit permit3 = future3.join();
+ assertNotNull(permit3);
+ assertEquals(permit3.getPermits(), 3);
+
+ semaphore.release(permit3);
+
+ }
+
+ @Test
+ public void testMultipleReleasesProcessQueue() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(10, 10, 5000);
+
+ // Acquire all permits in small chunks
+ List<AsyncSemaphorePermit> permits = new ArrayList<>();
+ for (int i = 0; i < 5; i++) {
+ CompletableFuture<AsyncSemaphorePermit> future = semaphore.acquire(2, () -> false);
+ permits.add(future.get(1, TimeUnit.SECONDS));
+ }
+
+ // Queue multiple requests
+ CompletableFuture<AsyncSemaphorePermit> queued1 = semaphore.acquire(3, () -> false);
+ CompletableFuture<AsyncSemaphorePermit> queued2 = semaphore.acquire(4, () -> false);
+ CompletableFuture<AsyncSemaphorePermit> queued3 = semaphore.acquire(2, () -> false);
+
+ assertFalse(queued1.isDone());
+ assertFalse(queued2.isDone());
+ assertFalse(queued3.isDone());
+
+ // Release permits one by one
+ semaphore.release(permits.get(0)); // 2 available
+ semaphore.release(permits.get(1)); // 4 available
+
+ // First queued request should complete
+ Thread.sleep(100);
+ assertTrue(queued1.isDone());
+ assertFalse(queued2.isDone());
+
+ semaphore.release(permits.get(2)); // 3 available (4 - 3 from queued1 + 2)
+
+ Thread.sleep(100);
+ assertFalse(queued2.isDone()); // Still needs 4
+
+ semaphore.release(permits.get(3)); // 5 available
+
+ // Second queued request should complete
+ Thread.sleep(100);
+ assertTrue(queued2.isDone());
+
+ semaphore.release(permits.get(4)); // 3 available (5 - 4 from queued2 + 2)
+
+ // Third queued request should complete
+ Thread.sleep(100);
+ assertTrue(queued3.isDone());
+ }
+
+ @Test
+ public void testUnboundedSemaphoreAcquireSinglePermit() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(0, 10, 5000);
+
+ CompletableFuture<AsyncSemaphorePermit> future = semaphore.acquire(1, () -> false);
+ AsyncSemaphorePermit permit = future.get(100, TimeUnit.MILLISECONDS);
+
+ assertNotNull(permit);
+ assertEquals(permit.getPermits(), 1);
+ assertTrue(future.isDone());
+ }
+
+ @Test
+ public void testUnboundedSemaphoreAcquireMultiplePermits() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(0, 10, 5000);
+
+ CompletableFuture<AsyncSemaphorePermit> future = semaphore.acquire(100, () -> false);
+ AsyncSemaphorePermit permit = future.get(100, TimeUnit.MILLISECONDS);
+
+ assertNotNull(permit);
+ assertEquals(permit.getPermits(), 100);
+ assertTrue(future.isDone());
+ }
+
+ @Test
+ public void testUnboundedSemaphoreNoQueueing() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(0, 10, 5000);
+
+ List<CompletableFuture<AsyncSemaphorePermit>> futures = new ArrayList<>();
+
+ // Acquire many permits concurrently - all should complete immediately
+ for (int i = 0; i < 100; i++) {
+ futures.add(semaphore.acquire(10, () -> false));
+ }
+
+ // All futures should complete immediately without queueing
+ for (CompletableFuture<AsyncSemaphorePermit> future : futures) {
+ AsyncSemaphorePermit permit = future.get(100, TimeUnit.MILLISECONDS);
+ assertNotNull(permit);
+ assertEquals(permit.getPermits(), 10);
+ assertTrue(future.isDone());
+ }
+ }
+
+ @Test
+ public void testUnboundedSemaphoreReleaseIsNoop() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(0, 10, 5000);
+
+ CompletableFuture<AsyncSemaphorePermit> future = semaphore.acquire(50, () -> false);
+ AsyncSemaphorePermit permit = future.get(100, TimeUnit.MILLISECONDS);
+
+ // Release should not throw exception but is a no-op
+ semaphore.release(permit);
+
+ // Should still be able to acquire more permits
+ CompletableFuture<AsyncSemaphorePermit> future2 = semaphore.acquire(100, () -> false);
+ AsyncSemaphorePermit permit2 = future2.get(100, TimeUnit.MILLISECONDS);
+ assertNotNull(permit2);
+ }
+
+ @Test
+ public void testUnboundedSemaphoreAvailablePermits() {
+ semaphore = new AsyncSemaphoreImpl(0, 10, 5000);
+
+ assertEquals(semaphore.getAvailablePermits(), Long.MAX_VALUE);
+ }
+
+ @Test
+ public void testUnboundedSemaphoreAcquiredPermits() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(0, 10, 5000);
+
+ // Acquire some permits
+ CompletableFuture<AsyncSemaphorePermit> future = semaphore.acquire(100, () -> false);
+ future.get(100, TimeUnit.MILLISECONDS);
+
+ // Acquired permits should always be 0 for unbounded semaphore
+ assertEquals(semaphore.getAcquiredPermits(), 0);
+ }
+
+ @Test
+ public void testUnboundedSemaphoreQueueSize() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(0, 10, 5000);
+
+ // Acquire multiple permits
+ for (int i = 0; i < 10; i++) {
+ semaphore.acquire(1, () -> false).get(100, TimeUnit.MILLISECONDS);
+ }
+
+ // Queue size should always be 0 since requests complete immediately
+ assertEquals(semaphore.getQueueSize(), 0);
+ }
+
+ @Test
+ public void testUnboundedSemaphoreUpdate() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(0, 10, 5000);
+
+ CompletableFuture<AsyncSemaphorePermit> future = semaphore.acquire(10, () -> false);
+ AsyncSemaphorePermit permit = future.get(100, TimeUnit.MILLISECONDS);
+
+ // Update should complete immediately
+ CompletableFuture<AsyncSemaphorePermit> future2 = semaphore.update(permit, 50, () -> false);
+ AsyncSemaphorePermit permit2 = future2.get(100, TimeUnit.MILLISECONDS);
+
+ assertNotNull(permit2);
+ assertEquals(permit2.getPermits(), 50);
+ assertTrue(future2.isDone());
+ }
+
+ @Test
+ public void testUnboundedSemaphoreUpdateToZero() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(0, 10, 5000);
+
+ CompletableFuture<AsyncSemaphorePermit> future = semaphore.acquire(100, () -> false);
+ AsyncSemaphorePermit permit = future.get(100, TimeUnit.MILLISECONDS);
+
+ // Update to zero should complete immediately
+ CompletableFuture<AsyncSemaphorePermit> future2 = semaphore.update(permit, 0, () -> false);
+ AsyncSemaphorePermit permit2 = future2.get(100, TimeUnit.MILLISECONDS);
+
+ assertNotNull(permit2);
+ assertEquals(permit2.getPermits(), 0);
+ assertTrue(future2.isDone());
+ }
+
+ @Test
+ public void testUnboundedSemaphoreNoTimeout() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(0, 10, 100); // Short timeout
+
+ // Even with short timeout, requests should complete immediately
+ CompletableFuture<AsyncSemaphorePermit> future = semaphore.acquire(1000, () -> false);
+ AsyncSemaphorePermit permit = future.get(50, TimeUnit.MILLISECONDS);
+
+ assertNotNull(permit);
+ assertEquals(permit.getPermits(), 1000);
+
+ Thread.sleep(150); // Wait longer than timeout
+
+ // Permit should still be valid
+ assertEquals(permit.getPermits(), 1000);
+ }
+
+ @Test
+ public void testUnboundedSemaphoreCancellationIgnored() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(0, 10, 5000);
+
+ AtomicBoolean cancelled = new AtomicBoolean(true);
+
+ // Even though cancelled is true, unbounded semaphore completes immediately
+ CompletableFuture<AsyncSemaphorePermit> future = semaphore.acquire(10, cancelled::get);
+ AsyncSemaphorePermit permit = future.get(100, TimeUnit.MILLISECONDS);
+
+ assertNotNull(permit);
+ assertEquals(permit.getPermits(), 10);
+ assertTrue(future.isDone());
+ assertFalse(future.isCompletedExceptionally());
+ }
+
+ @Test
+ public void testPermitsArentAcquiredInUpdateWhenCancelled() throws Exception {
+ semaphore = new AsyncSemaphoreImpl(10, 10, 5000);
+
+ // setup
+ AtomicBoolean cancelled = new AtomicBoolean(false);
+ // acquire 5 permits
+ CompletableFuture<AsyncSemaphorePermit> future = semaphore.acquire(5, cancelled::get);
+ assertThat(future).succeedsWithin(Duration.ofSeconds(1))
+ .satisfies(p -> assertThat(p.getPermits()).isEqualTo(5));
+ AsyncSemaphorePermit permit = future.join();
+ assertThat(semaphore.getAcquiredPermits()).isEqualTo(5);
+
+ // when permits are update when the request is already cancelled
+ cancelled.set(true);
+ CompletableFuture<AsyncSemaphorePermit> future2 = semaphore.update(permit, 10, cancelled::get);
+ assertThat(future2).failsWithin(Duration.ofSeconds(1))
+ .withThrowableThat().havingRootCause()
+ .isInstanceOf(AsyncSemaphore.PermitAcquireCancelledException.class);
+ assertThat(semaphore.getAcquiredPermits()).isEqualTo(5);
+
+ // then no permits should be acquired so that
+ // when a new acquisition with update is made, it should succeed
+ cancelled.set(false);
+ CompletableFuture<AsyncSemaphorePermit> future3 = semaphore.update(permit, 10, cancelled::get);
+ assertThat(future3).succeedsWithin(Duration.ofSeconds(1))
+ .satisfies(p -> assertThat(p.getPermits()).isEqualTo(10));
+ assertThat(semaphore.getAcquiredPermits()).isEqualTo(10);
+
+ // when original permit is released, it shouldn't reduce acquired permits
+ semaphore.release(permit);
+ assertThat(semaphore.getAcquiredPermits()).isEqualTo(10);
+
+ // when updated permit is released, it should reduce acquired permits
+ semaphore.release(future3.join());
+ assertThat(semaphore.getAcquiredPermits()).isEqualTo(0);
+ }
+
+ @Test
+ public void testSupportCompletableFutureCancel() throws Exception {
+ int timeoutMillis = 200;
+ semaphore = new AsyncSemaphoreImpl(10, 10, timeoutMillis);
+
+ // setup
+ AtomicBoolean cancelled = new AtomicBoolean(false);
+ // acquire 5 permits
+ CompletableFuture<AsyncSemaphorePermit> future = semaphore.acquire(5, cancelled::get);
+ assertThat(future).succeedsWithin(Duration.ofSeconds(1))
+ .satisfies(p -> assertThat(p.getPermits()).isEqualTo(5));
+ assertThat(semaphore.getAcquiredPermits()).isEqualTo(5);
+
+ // attempt to acquire 5 permits
+ CompletableFuture<AsyncSemaphorePermit> future2 = semaphore.acquire(10, cancelled::get);
+
+ // attempt to acquire 5 permits
+ CompletableFuture<AsyncSemaphorePermit> future3 = semaphore.acquire(5, cancelled::get);
+ // cancel future2
+ future2.cancel(true);
+ assertThat(future2).isCancelled();
+ // now future3 should succeed
+ assertThat(future3).succeedsWithin(Duration.ofSeconds(1))
+ .satisfies(p -> assertThat(p.getPermits()).isEqualTo(5));
+
+ assertThat(semaphore.getAcquiredPermits()).isEqualTo(10);
+ }
+}
\ No newline at end of file
diff --git a/pulsar-proxy/pom.xml b/pulsar-proxy/pom.xml
index 9531879ea8e..42d16e22d48 100644
--- a/pulsar-proxy/pom.xml
+++ b/pulsar-proxy/pom.xml
@@ -215,6 +215,18 @@
<version>${consolecaptor.version}</version>
<scope>test</scope>
</dependency>
+ <dependency>
+ <groupId>io.rest-assured</groupId>
+ <artifactId>rest-assured</artifactId>
+ <scope>test</scope>
+ </dependency>
+ <dependency>
+ <groupId>${project.groupId}</groupId>
+ <artifactId>pulsar-broker-common</artifactId>
+ <version>${project.version}</version>
+ <type>test-jar</type>
+ <scope>test</scope>
+ </dependency>
</dependencies>
<build>
<plugins>
diff --git a/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/LookupProxyHandler.java b/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/LookupProxyHandler.java
index c38e2ba08c1..49a59444b62 100644
--- a/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/LookupProxyHandler.java
+++ b/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/LookupProxyHandler.java
@@ -18,6 +18,7 @@
*/
package org.apache.pulsar.proxy.server;
+import static org.apache.pulsar.common.semaphore.AsyncDualMemoryLimiterUtil.acquireDirectMemoryPermitsAndWriteAndFlush;
import io.netty.buffer.ByteBuf;
import io.netty.channel.ChannelHandlerContext;
import io.prometheus.client.Counter;
@@ -26,20 +27,29 @@ import java.net.SocketAddress;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Semaphore;
+import java.util.function.BooleanSupplier;
import org.apache.commons.lang3.StringUtils;
+import org.apache.pulsar.broker.topiclistlimit.TopicListMemoryLimiter;
+import org.apache.pulsar.broker.topiclistlimit.TopicListSizeResultCache;
import org.apache.pulsar.client.api.PulsarClientException;
import org.apache.pulsar.client.impl.BinaryProtoLookupService;
+import org.apache.pulsar.client.impl.ClientCnx;
+import org.apache.pulsar.common.api.proto.BaseCommand;
import org.apache.pulsar.common.api.proto.CommandGetSchema;
import org.apache.pulsar.common.api.proto.CommandGetTopicsOfNamespace;
import org.apache.pulsar.common.api.proto.CommandLookupTopic;
import org.apache.pulsar.common.api.proto.CommandLookupTopicResponse.LookupType;
import org.apache.pulsar.common.api.proto.CommandPartitionedTopicMetadata;
import org.apache.pulsar.common.api.proto.ServerError;
+import org.apache.pulsar.common.lookup.GetTopicsResult;
import org.apache.pulsar.common.naming.TopicName;
import org.apache.pulsar.common.protocol.Commands;
import org.apache.pulsar.common.protocol.schema.BytesSchemaVersion;
import org.apache.pulsar.common.protocol.schema.SchemaVersion;
+import org.apache.pulsar.common.semaphore.AsyncDualMemoryLimiter;
+import org.apache.pulsar.common.semaphore.AsyncDualMemoryLimiterImpl;
import org.apache.pulsar.common.util.netty.NettyChannelUtil;
import org.apache.pulsar.policies.data.loadbalancer.ServiceLookupData;
import org.slf4j.Logger;
@@ -84,6 +94,8 @@ public class LookupProxyHandler {
"Counter of getTopicsOfNamespace requests rejected due to throttling")
.create().register();
private final Semaphore lookupRequestSemaphore;
+ private final AsyncDualMemoryLimiterImpl maxTopicListInFlightLimiter;
+ private final TopicListSizeResultCache topicListSizeResultCache;
public LookupProxyHandler(ProxyService proxy, ProxyConnection proxyConnection) {
this.discoveryProvider = proxy.getDiscoveryProvider();
@@ -93,6 +105,8 @@ public class LookupProxyHandler {
this.connectWithTLS = proxy.getConfiguration().isTlsEnabledWithBroker();
this.brokerServiceURL = this.connectWithTLS ? proxy.getConfiguration().getBrokerServiceURLTLS()
: proxy.getConfiguration().getBrokerServiceURL();
+ this.maxTopicListInFlightLimiter = proxy.getMaxTopicListInFlightLimiter();
+ this.topicListSizeResultCache = proxy.getTopicListSizeResultCache();
}
public void handleLookup(CommandLookupTopic lookup) {
@@ -326,7 +340,7 @@ public class LookupProxyHandler {
if (log.isDebugEnabled()) {
log.debug("Getting connections to '{}' for getting TopicsOfNamespace '{}' with clientReq Id '{}'",
- addr, namespaceName, clientRequestId);
+ addr, namespaceName, clientRequestId);
}
proxyConnection.getConnectionPool().getConnection(addr).thenAccept(clientCnx -> {
// Connected to backend broker
@@ -334,20 +348,8 @@ public class LookupProxyHandler {
ByteBuf command;
command = Commands.newGetTopicsOfNamespaceRequest(namespaceName, requestId, mode,
topicsPattern, topicsHash);
- clientCnx.newGetTopicsOfNamespace(command, requestId).whenComplete((r, t) -> {
- if (t != null) {
- log.warn("[{}] Failed to get TopicsOfNamespace {}: {}",
- clientAddress, namespaceName, t.getMessage());
- writeAndFlush(
- Commands.newError(clientRequestId, getServerError(t), t.getMessage()));
- } else {
- writeAndFlush(
- Commands.newGetTopicsOfNamespaceResponse(r.getNonPartitionedOrPartitionTopics(),
- r.getTopicsHash(), r.isFiltered(),
- r.isChanged(), clientRequestId));
- }
- });
+ internalPerformGetTopicsOfNamespace(clientRequestId, namespaceName, mode, clientCnx, command, requestId);
proxyConnection.getConnectionPool().releaseConnection(clientCnx);
}).exceptionally(ex -> {
// Failed to connect to backend broker
@@ -357,6 +359,70 @@ public class LookupProxyHandler {
});
}
+ private void internalPerformGetTopicsOfNamespace(long clientRequestId, String namespaceName,
+ CommandGetTopicsOfNamespace.Mode mode, ClientCnx clientCnx,
+ ByteBuf command, long requestId) {
+ BooleanSupplier isPermitRequestCancelled = () -> !proxyConnection.ctx().channel().isActive();
+ TopicListSizeResultCache.ResultHolder
+ listSizeHolder = topicListSizeResultCache.getTopicListSize(namespaceName, mode);
+ listSizeHolder.getSizeAsync().thenAccept(initialSize -> {
+ maxTopicListInFlightLimiter.withAcquiredPermits(initialSize,
+ AsyncDualMemoryLimiter.LimitType.HEAP_MEMORY, isPermitRequestCancelled, initialPermits -> {
+ return clientCnx.newGetTopicsOfNamespace(command, requestId).whenComplete((r, t) -> {
+ if (t != null) {
+ log.warn("[{}] Failed to get TopicsOfNamespace {}: {}", clientAddress, namespaceName,
+ t.getMessage());
+ listSizeHolder.resetIfInitializing();
+ writeAndFlush(Commands.newError(clientRequestId, getServerError(t), t.getMessage()));
+ } else {
+ long actualSize = TopicListMemoryLimiter.estimateTopicListSize(
+ r.getNonPartitionedOrPartitionTopics());
+ listSizeHolder.updateSize(actualSize);
+ maxTopicListInFlightLimiter.withUpdatedPermits(initialPermits, actualSize,
+ isPermitRequestCancelled, permits -> {
+ return handleWritingGetTopicsResponse(clientRequestId, r,
+ isPermitRequestCancelled);
+ }, t2 -> {
+ log.warn("[{}] Failed to acquire actual heap memory permits for "
+ + "GetTopicsOfNamespace: {}", clientAddress, t2.getMessage());
+ writeAndFlush(
+ Commands.newError(clientRequestId, ServerError.TooManyRequests,
+ "Failed due to heap memory limit exceeded"));
+
+ return CompletableFuture.completedFuture(null);
+ });
+ }
+ }).thenApply(__ -> null);
+ }, t -> {
+ log.warn("[{}] Failed to acquire initial heap memory permits for GetTopicsOfNamespace: {}",
+ clientAddress, t.getMessage());
+ listSizeHolder.resetIfInitializing();
+ writeAndFlush(Commands.newError(clientRequestId, ServerError.TooManyRequests,
+ "Failed due to heap memory limit exceeded"));
+
+ return CompletableFuture.completedFuture(null);
+ }).exceptionally(ex -> {
+ listSizeHolder.resetIfInitializing();
+ writeAndFlush(Commands.newError(clientRequestId, getServerError(ex), ex.getMessage()));
+ return null;
+ });
+ });
+ }
+
+ private CompletableFuture<Void> handleWritingGetTopicsResponse(long clientRequestId, GetTopicsResult r,
+ BooleanSupplier isCancelled) {
+ BaseCommand responseCommand = Commands.newGetTopicsOfNamespaceResponseCommand(
+ r.getNonPartitionedOrPartitionTopics(), r.getTopicsHash(), r.isFiltered(),
+ r.isChanged(), clientRequestId);
+ return acquireDirectMemoryPermitsAndWriteAndFlush(proxyConnection.ctx(), maxTopicListInFlightLimiter,
+ isCancelled, responseCommand, t -> {
+ log.warn("[{}] Failed to acquire actual direct memory permits for GetTopicsOfNamespace: {}",
+ clientAddress, t.getMessage());
+ writeAndFlush(Commands.newError(clientRequestId, ServerError.TooManyRequests,
+ "Failed due to direct memory limit exceeded"));
+ });
+ }
+
public void handleGetSchema(CommandGetSchema commandGetSchema) {
GET_SCHEMA_REQUESTS.inc();
if (log.isDebugEnabled()) {
diff --git a/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyConfiguration.java b/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyConfiguration.java
index 6db1b302c66..9fd91a591cb 100644
--- a/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyConfiguration.java
+++ b/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyConfiguration.java
@@ -480,6 +480,42 @@ public class ProxyConfiguration implements PulsarConfiguration {
)
private int maxConcurrentLookupRequests = 50000;
+ @FieldContext(
+ category = CATEGORY_RATE_LIMITING,
+ doc = "Maximum heap memory for inflight topic list operations (MB).\n"
+ + "Default: 100 MB (supports ~1M topic names assuming 100 bytes each)")
+ private int maxTopicListInFlightHeapMemSizeMB = 100;
+
+ @FieldContext(
+ category = CATEGORY_RATE_LIMITING,
+ doc = "Maximum direct memory for inflight topic list responses (MB).\n"
+ + "Default: 100 MB (network buffers for serialized responses)")
+ private int maxTopicListInFlightDirectMemSizeMB = 100;
+
+ @FieldContext(
+ category = CATEGORY_RATE_LIMITING,
+ doc = "Timeout for acquiring heap memory permits (milliseconds).\n"
+ + "Default: 25000 (25 seconds)")
+ private int maxTopicListInFlightHeapMemSizePermitsAcquireTimeoutMillis = 25000;
+
+ @FieldContext(
+ category = CATEGORY_RATE_LIMITING,
+ doc = "Maximum queue size for heap memory permit requests.\n"
+ + "Default: 10000 (prevent unbounded queueing)")
+ private int maxTopicListInFlightHeapMemSizePermitsAcquireQueueSize = 10000;
+
+ @FieldContext(
+ category = CATEGORY_RATE_LIMITING,
+ doc = "Timeout for acquiring direct memory permits (milliseconds).\n"
+ + "Default: 25000 (25 seconds)")
+ private int maxTopicListInFlightDirectMemSizePermitsAcquireTimeoutMillis = 25000;
+
+ @FieldContext(
+ category = CATEGORY_RATE_LIMITING,
+ doc = "Maximum queue size for direct memory permit requests.\n"
+ + "Default: 10000 (prevent unbounded queueing)")
+ private int maxTopicListInFlightDirectMemSizePermitsAcquireQueueSize = 10000;
+
@FieldContext(
category = CATEGORY_CLIENT_AUTHENTICATION,
doc = "The authentication plugin used by the Pulsar proxy to authenticate with Pulsar brokers"
diff --git a/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyService.java b/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyService.java
index 11afb68398e..251e25a0e6b 100644
--- a/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyService.java
+++ b/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyService.java
@@ -35,6 +35,7 @@ import io.netty.resolver.dns.DnsAddressResolverGroup;
import io.netty.resolver.dns.DnsNameResolverBuilder;
import io.netty.util.concurrent.DefaultThreadFactory;
import io.netty.util.concurrent.Future;
+import io.prometheus.client.CollectorRegistry;
import io.prometheus.client.Counter;
import io.prometheus.client.Gauge;
import java.io.Closeable;
@@ -62,10 +63,13 @@ import org.apache.pulsar.broker.limiter.ConnectionController;
import org.apache.pulsar.broker.resources.PulsarResources;
import org.apache.pulsar.broker.stats.prometheus.PrometheusMetricsServlet;
import org.apache.pulsar.broker.stats.prometheus.PrometheusRawMetricsProvider;
+import org.apache.pulsar.broker.topiclistlimit.TopicListMemoryLimiter;
+import org.apache.pulsar.broker.topiclistlimit.TopicListSizeResultCache;
import org.apache.pulsar.broker.web.plugin.servlet.AdditionalServlets;
import org.apache.pulsar.client.api.Authentication;
import org.apache.pulsar.common.allocator.PulsarByteBufAllocator;
import org.apache.pulsar.common.configuration.PulsarConfigurationLoader;
+import org.apache.pulsar.common.semaphore.AsyncDualMemoryLimiterImpl;
import org.apache.pulsar.common.util.netty.DnsResolverUtil;
import org.apache.pulsar.common.util.netty.EventLoopUtil;
import org.apache.pulsar.metadata.api.MetadataStoreException;
@@ -149,11 +153,17 @@ public class ProxyService implements Closeable {
private PrometheusMetricsServlet metricsServlet;
private List<PrometheusRawMetricsProvider> pendingMetricsProviders;
@Getter
- private PulsarProxyOpenTelemetry openTelemetry;
+ private final PulsarProxyOpenTelemetry openTelemetry;
@Getter
private final ConnectionController connectionController;
+ @Getter
+ private final AsyncDualMemoryLimiterImpl maxTopicListInFlightLimiter;
+
+ @Getter
+ private final TopicListSizeResultCache topicListSizeResultCache = new TopicListSizeResultCache();
+
private boolean gracefulShutdown = true;
public ProxyService(ProxyConfiguration proxyConfig,
@@ -212,6 +222,18 @@ public class ProxyService implements Closeable {
this.connectionController = new ConnectionController.DefaultConnectionController(
proxyConfig.getMaxConcurrentInboundConnections(),
proxyConfig.getMaxConcurrentInboundConnectionsPerIp());
+
+ this.openTelemetry = new PulsarProxyOpenTelemetry(proxyConfig);
+
+ // Initialize topic list memory limiter
+ this.maxTopicListInFlightLimiter = new TopicListMemoryLimiter(
+ CollectorRegistry.defaultRegistry, "pulsar_proxy_", openTelemetry.getMeter(),
+ proxyConfig.getMaxTopicListInFlightHeapMemSizeMB() * 1024L * 1024L,
+ proxyConfig.getMaxTopicListInFlightHeapMemSizePermitsAcquireQueueSize(),
+ proxyConfig.getMaxTopicListInFlightHeapMemSizePermitsAcquireTimeoutMillis(),
+ proxyConfig.getMaxTopicListInFlightDirectMemSizeMB() * 1024L * 1024L,
+ proxyConfig.getMaxTopicListInFlightDirectMemSizePermitsAcquireQueueSize(),
+ proxyConfig.getMaxTopicListInFlightDirectMemSizePermitsAcquireTimeoutMillis());
}
public void start() throws Exception {
@@ -287,7 +309,6 @@ public class ProxyService implements Closeable {
}
createMetricsServlet();
- openTelemetry = new PulsarProxyOpenTelemetry(proxyConfig);
// Initialize the message protocol handlers.
// start the protocol handlers only after the broker is ready,
@@ -391,6 +412,8 @@ public class ProxyService implements Closeable {
dnsAddressResolverGroup.close();
+ maxTopicListInFlightLimiter.close();
+
if (discoveryProvider != null) {
discoveryProvider.close();
}
diff --git a/pulsar-proxy/src/test/java/org/apache/pulsar/proxy/server/InvalidProxyConfigForAuthorizationTest.java b/pulsar-proxy/src/test/java/org/apache/pulsar/proxy/server/InvalidProxyConfigForAuthorizationTest.java
index e97c355ebe5..1bd6e712c94 100644
--- a/pulsar-proxy/src/test/java/org/apache/pulsar/proxy/server/InvalidProxyConfigForAuthorizationTest.java
+++ b/pulsar-proxy/src/test/java/org/apache/pulsar/proxy/server/InvalidProxyConfigForAuthorizationTest.java
@@ -32,6 +32,7 @@ public class InvalidProxyConfigForAuthorizationTest {
ProxyConfiguration proxyConfiguration = new ProxyConfiguration();
proxyConfiguration.setAuthorizationEnabled(true);
proxyConfiguration.setAuthenticationEnabled(false);
+ proxyConfiguration.setClusterName("test");
try (ProxyService proxyService = new ProxyService(proxyConfiguration,
Mockito.mock(AuthenticationService.class), Mockito.mock(Authentication.class))) {
proxyService.start();
diff --git a/pulsar-proxy/src/test/java/org/apache/pulsar/proxy/server/ProxyPatternConsumerBackPressureMultipleConsumersTest.java b/pulsar-proxy/src/test/java/org/apache/pulsar/proxy/server/ProxyPatternConsumerBackPressureMultipleConsumersTest.java
new file mode 100644
index 00000000000..3263cc2985c
--- /dev/null
+++ b/pulsar-proxy/src/test/java/org/apache/pulsar/proxy/server/ProxyPatternConsumerBackPressureMultipleConsumersTest.java
@@ -0,0 +1,92 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.pulsar.proxy.server;
+
+import static org.mockito.Mockito.doReturn;
+import java.util.Optional;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.pulsar.broker.authentication.AuthenticationService;
+import org.apache.pulsar.client.api.Authentication;
+import org.apache.pulsar.client.api.AuthenticationFactory;
+import org.apache.pulsar.client.api.PatternConsumerBackPressureMultipleConsumersTest;
+import org.apache.pulsar.common.configuration.PulsarConfigurationLoader;
+import org.apache.pulsar.metadata.impl.ZKMetadataStore;
+import org.mockito.Mockito;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+
+@Slf4j
+public class ProxyPatternConsumerBackPressureMultipleConsumersTest extends
+ PatternConsumerBackPressureMultipleConsumersTest {
+
+ protected ProxyService proxyService;
+ protected ProxyConfiguration proxyConfig = new ProxyConfiguration();
+ protected Authentication proxyClientAuthentication;
+
+ @Override
+ @BeforeMethod
+ protected void setup() throws Exception {
+ super.setup();
+ initializeProxyConfig();
+ proxyService = Mockito.spy(new ProxyService(proxyConfig, new AuthenticationService(
+ PulsarConfigurationLoader.convertFrom(proxyConfig)), proxyClientAuthentication));
+ doReturn(registerCloseable(new ZKMetadataStore(mockZooKeeper))).when(proxyService).createLocalMetadataStore();
+ doReturn(registerCloseable(new ZKMetadataStore(mockZooKeeperGlobal))).when(proxyService)
+ .createConfigurationMetadataStore();
+ proxyService.start();
+ }
+
+ protected void initializeProxyConfig() throws Exception {
+ proxyConfig.setServicePort(Optional.ofNullable(0));
+ proxyConfig.setBrokerProxyAllowedTargetPorts("*");
+ proxyConfig.setMetadataStoreUrl(DUMMY_VALUE);
+ proxyConfig.setConfigurationMetadataStoreUrl(GLOBAL_DUMMY_VALUE);
+ proxyConfig.setClusterName(configClusterName);
+
+ proxyClientAuthentication = AuthenticationFactory.create(
+ proxyConfig.getBrokerClientAuthenticationPlugin(),
+ proxyConfig.getBrokerClientAuthenticationParameters());
+ proxyClientAuthentication.start();
+ }
+
+ @Override
+ @AfterMethod(alwaysRun = true)
+ protected void cleanup() throws Exception {
+ proxyService.close();
+ if (proxyClientAuthentication != null) {
+ proxyClientAuthentication.close();
+ }
+ super.cleanup();
+ }
+
+ @Override
+ protected String getClientServiceUrl() {
+ return proxyService.getServiceUrl();
+ }
+
+ @Override
+ protected int getDirectMemoryRequiredMB() {
+ return 2 * super.getDirectMemoryRequiredMB();
+ }
+
+ @Override
+ protected String getMetricsPrefix() {
+ return "pulsar_proxy_topic_list_";
+ }
+}