[improve][broker] PIP-442: Add memory limits for topic list watcher (part 2) (#25070)

Author: lhotari

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/PulsarCommandSender.java

@@ -20,10 +20,11 @@
 
 
 import io.netty.util.concurrent.Future;
+import java.util.Collection;
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
-import java.util.function.Consumer;
+import java.util.function.Function;
 import org.apache.bookkeeper.mledger.Entry;
 import org.apache.pulsar.client.api.transaction.TxnID;
 import org.apache.pulsar.common.api.proto.CommandLookupTopicResponse;
@@ -55,7 +56,8 @@ void sendSendReceiptResponse(long producerId, long sequenceId, long highestId, l
 
     CompletableFuture<Void> sendGetTopicsOfNamespaceResponse(List<String> topics, String topicsHash, boolean filtered,
                                                              boolean changed, long requestId,
-                                                             Consumer<Throwable> permitAcquireErrorHandler);
+                                                             Function<Throwable, CompletableFuture<Void>>
+                                                                     permitAcquireErrorHandler);
 
     void sendGetSchemaResponse(long requestId, SchemaInfo schema, SchemaVersion version);
 
@@ -96,8 +98,13 @@ Future<Void> sendMessagesToConsumer(long consumerId, String topicName, Subscript
 
     void sendEndTxnErrorResponse(long requestId, TxnID txnID, ServerError error, String message);
 
-    void sendWatchTopicListSuccess(long requestId, long watcherId, String topicsHash, List<String> topics);
+    CompletableFuture<Void> sendWatchTopicListSuccess(long requestId, long watcherId, String topicsHash,
+                                                      Collection<String> topics,
+                                                      Function<Throwable, CompletableFuture<Void>>
+                                                              permitAcquireErrorHandler);
 
-    void sendWatchTopicListUpdate(long watcherId,
-                                         List<String> newTopics, List<String> deletedTopics, String topicsHash);
+    CompletableFuture<Void> sendWatchTopicListUpdate(long watcherId,
+                                  List<String> newTopics, List<String> deletedTopics, String topicsHash,
+                                  Function<Throwable, CompletableFuture<Void>>
+                                                             permitAcquireErrorHandler);
 }

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/PulsarCommandSenderImpl.java

@@ -24,10 +24,11 @@
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelPromise;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
-import java.util.function.Consumer;
+import java.util.function.Function;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.bookkeeper.mledger.Entry;
 import org.apache.pulsar.broker.intercept.BrokerInterceptor;
@@ -129,7 +130,9 @@ public void sendSendError(long producerId, long sequenceId, ServerError error, S
     @Override
     public CompletableFuture<Void> sendGetTopicsOfNamespaceResponse(List<String> topics, String topicsHash,
                                                                     boolean filtered, boolean changed, long requestId,
-                                                                    Consumer<Throwable> permitAcquireErrorHandler) {
+                                                                    Function<Throwable,
+                                                                            CompletableFuture<Void>>
+                                                                                permitAcquireErrorHandler) {
         BaseCommand command = Commands.newGetTopicsOfNamespaceResponseCommand(topics, topicsHash,
                 filtered, changed, requestId);
         safeIntercept(command, cnx);
@@ -366,27 +369,32 @@ public void sendEndTxnErrorResponse(long requestId, TxnID txnID, ServerError err
 
     /***
      * @param topics topic names which are matching, the topic name contains the partition suffix.
+     * @return a CompletableFuture&lt;Void&gt; that completes when the operation finishes
      */
     @Override
-    public void sendWatchTopicListSuccess(long requestId, long watcherId, String topicsHash, List<String> topics) {
+    public CompletableFuture<Void> sendWatchTopicListSuccess(long requestId, long watcherId, String topicsHash,
+                                                             Collection<String> topics,
+                                                             Function<Throwable, CompletableFuture<Void>>
+                                                                         permitAcquireErrorHandler) {
         BaseCommand command = Commands.newWatchTopicListSuccess(requestId, watcherId, topicsHash, topics);
-        interceptAndWriteCommand(command);
+        safeIntercept(command, cnx);
+        return acquireDirectMemoryPermitsAndWriteAndFlush(cnx.ctx(), maxTopicListInFlightLimiter, () -> !cnx.isActive(),
+                command, permitAcquireErrorHandler);
     }
 
     /***
      * {@inheritDoc}
+     * @return a CompletableFuture that completes when the watch topic list update operation finishes
      */
     @Override
-    public void sendWatchTopicListUpdate(long watcherId,
-                                         List<String> newTopics, List<String> deletedTopics, String topicsHash) {
+    public CompletableFuture<Void> sendWatchTopicListUpdate(long watcherId, List<String> newTopics,
+                                                            List<String> deletedTopics, String topicsHash,
+                                                            Function<Throwable, CompletableFuture<Void>>
+                                                                        permitAcquireErrorHandler) {
         BaseCommand command = Commands.newWatchTopicUpdate(watcherId, newTopics, deletedTopics, topicsHash);
-        interceptAndWriteCommand(command);
-    }
-
-    private void interceptAndWriteCommand(BaseCommand command) {
         safeIntercept(command, cnx);
-        ByteBuf outBuf = Commands.serializeWithSize(command);
-        writeAndFlush(outBuf);
+        return acquireDirectMemoryPermitsAndWriteAndFlush(cnx.ctx(), maxTopicListInFlightLimiter, () -> !cnx.isActive(),
+                command, permitAcquireErrorHandler);
     }
 
     private void writeAndFlush(ByteBuf outBuf) {

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/ServerCnx.java

@@ -2644,6 +2644,7 @@ private void internalHandleGetTopicsOfNamespace(String namespace, NamespaceName
                                                             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 "