Key_Shared subscription on non-persistent topic

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

@@ -20,11 +20,18 @@
 
 import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
 
+import io.netty.buffer.ByteBuf;
+import org.apache.commons.lang3.StringUtils;
+import org.apache.pulsar.broker.service.persistent.PersistentStickyKeyDispatcherMultipleConsumers;
+import org.apache.pulsar.common.api.Commands;
+import org.apache.pulsar.common.api.proto.PulsarApi;
 import org.apache.pulsar.common.api.proto.PulsarApi.CommandSubscribe.SubType;
 import org.apache.pulsar.utils.CopyOnWriteArrayList;
 
 import com.carrotsearch.hppc.ObjectHashSet;
 import com.carrotsearch.hppc.ObjectSet;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  */
@@ -205,4 +212,23 @@ private int getFirstConsumerIndexOfPriority(int targetPriority) {
         return -1;
     }
 
+    public static final String NONE_KEY = "NONE_KEY";
+    protected byte[] peekStickyKey(ByteBuf metadataAndPayload) {
+        metadataAndPayload.markReaderIndex();
+        PulsarApi.MessageMetadata metadata = Commands.parseMessageMetadata(metadataAndPayload);
+        metadataAndPayload.resetReaderIndex();
+        String key = metadata.getPartitionKey();
+        if (log.isDebugEnabled()) {
+            log.debug("Parse message metadata, partition key is {}, ordering key is {}", key, metadata.getOrderingKey());
+        }
+        if (StringUtils.isNotBlank(key) || metadata.hasOrderingKey()) {
+            return metadata.hasOrderingKey() ? metadata.getOrderingKey().toByteArray() : key.getBytes();
+        }
+        metadata.recycle();
+        return NONE_KEY.getBytes();
+    }
+
+    private static final Logger log = LoggerFactory.getLogger(PersistentStickyKeyDispatcherMultipleConsumers.class);
+
+
 }

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/nonpersistent/NonPersistentDispatcherMultipleConsumers.java

@@ -46,11 +46,11 @@ public class NonPersistentDispatcherMultipleConsumers extends AbstractDispatcher
         implements NonPersistentDispatcher {
 
     private final NonPersistentTopic topic;
-    private final Subscription subscription;
+    protected final Subscription subscription;
 
     private CompletableFuture<Void> closeFuture = null;
     private final String name;
-    private final Rate msgDrop;
+    protected final Rate msgDrop;
     protected static final AtomicIntegerFieldUpdater<NonPersistentDispatcherMultipleConsumers> TOTAL_AVAILABLE_PERMITS_UPDATER = AtomicIntegerFieldUpdater
             .newUpdater(NonPersistentDispatcherMultipleConsumers.class, "totalAvailablePermits");
     @SuppressWarnings("unused")

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/nonpersistent/NonPersistentStickyKeyDispatcherMultipleConsumers.java

@@ -0,0 +1,102 @@
+/**
+ * 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.service.nonpersistent;
+
+import org.apache.bookkeeper.mledger.Entry;
+import org.apache.pulsar.broker.service.BrokerServiceException;
+import org.apache.pulsar.broker.service.Consumer;
+import org.apache.pulsar.broker.service.EntryBatchSizes;
+import org.apache.pulsar.broker.service.HashRangeStickyKeyConsumerSelector;
+import org.apache.pulsar.broker.service.SendMessageInfo;
+import org.apache.pulsar.broker.service.StickyKeyConsumerSelector;
+import org.apache.pulsar.broker.service.Subscription;
+import org.apache.pulsar.common.api.Commands;
+import org.apache.pulsar.common.api.proto.PulsarApi.CommandSubscribe.SubType;
+import org.apache.pulsar.common.util.Murmur3_32Hash;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+
+public class NonPersistentStickyKeyDispatcherMultipleConsumers extends NonPersistentDispatcherMultipleConsumers {
+
+    private final StickyKeyConsumerSelector selector;
+
+    public NonPersistentStickyKeyDispatcherMultipleConsumers(NonPersistentTopic topic, Subscription subscription) {
+        super(topic, subscription);
+        //TODO: Consumer selector Pluggable
+        selector = new HashRangeStickyKeyConsumerSelector();
+    }
+
+    @Override
+    public synchronized void addConsumer(Consumer consumer) throws BrokerServiceException {
+        super.addConsumer(consumer);
+        selector.addConsumer(consumer);
+    }
+
+    @Override
+    public synchronized void removeConsumer(Consumer consumer) throws BrokerServiceException {
+        super.removeConsumer(consumer);
+        selector.removeConsumer(consumer);
+    }
+
+    @Override
+    public SubType getType() {
+        return SubType.Key_Shared;
+    }
+
+    @Override
+    public void sendMessages(List<Entry> entries) {
+        if (entries.size() > 0) {
+            final Map<Integer, List<Entry>> groupedEntries = new HashMap<>();
+            for (Entry entry : entries) {
+                int key = Murmur3_32Hash.getInstance().makeHash(peekStickyKey(entry.getDataBuffer()));
+                groupedEntries.putIfAbsent(key, new ArrayList<>());
+                groupedEntries.get(key).add(entry);
+            }
+            final Iterator<Map.Entry<Integer, List<Entry>>> iterator = groupedEntries.entrySet().iterator();
+            while (iterator.hasNext()) {
+                final Map.Entry<Integer, List<Entry>> entriesWithSameKey = iterator.next();
+                //TODO: None key policy
+                Consumer consumer = null;
+                if (selector instanceof HashRangeStickyKeyConsumerSelector) {
+                    consumer = ((HashRangeStickyKeyConsumerSelector)selector).select(entriesWithSameKey.getKey());
+                }
+                if (consumer != null) {
+                    SendMessageInfo sendMessageInfo = SendMessageInfo.getThreadLocal();
+                    EntryBatchSizes batchSizes = EntryBatchSizes.get(entriesWithSameKey.getValue().size());
+                    filterEntriesForConsumer(entriesWithSameKey.getValue(), batchSizes, sendMessageInfo);
+                    consumer.sendMessages(entriesWithSameKey.getValue(), batchSizes, sendMessageInfo.getTotalMessages(),
+                            sendMessageInfo.getTotalBytes(), getRedeliveryTracker());
+                    TOTAL_AVAILABLE_PERMITS_UPDATER.addAndGet(this, -sendMessageInfo.getTotalMessages());
+                } else {
+                    entries.forEach(entry -> {
+                        int totalMsgs = Commands.getNumberOfMessagesInBatch(entry.getDataBuffer(), subscription.toString(), -1);
+                        if (totalMsgs > 0) {
+                            msgDrop.recordEvent(totalMsgs);
+                        }
+                        entry.release();
+                    });
+                }
+            }
+        }
+    }
+}

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/nonpersistent/NonPersistentSubscription.java

@@ -111,6 +111,11 @@ public synchronized void addConsumer(Consumer consumer) throws BrokerServiceExce
                             topic, this);
                 }
                 break;
+            case Key_Shared:
+                if (dispatcher == null || dispatcher.getType() != SubType.Key_Shared) {
+                    dispatcher = new NonPersistentStickyKeyDispatcherMultipleConsumers(topic, this);
+                }
+                break;
             default:
                 throw new ServerMetadataException("Unsupported subscription type");
             }

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumers.java

@@ -18,8 +18,6 @@
  */
 package org.apache.pulsar.broker.service.persistent;
 
-import io.netty.buffer.ByteBuf;
-
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.Iterator;
@@ -31,25 +29,20 @@
 import org.apache.bookkeeper.mledger.Entry;
 import org.apache.bookkeeper.mledger.ManagedCursor;
 import org.apache.bookkeeper.mledger.Position;
-import org.apache.commons.lang3.StringUtils;
 import org.apache.pulsar.broker.service.BrokerServiceException;
 import org.apache.pulsar.broker.service.Consumer;
 import org.apache.pulsar.broker.service.EntryBatchSizes;
 import org.apache.pulsar.broker.service.HashRangeStickyKeyConsumerSelector;
 import org.apache.pulsar.broker.service.SendMessageInfo;
 import org.apache.pulsar.broker.service.StickyKeyConsumerSelector;
 import org.apache.pulsar.broker.service.Subscription;
-import org.apache.pulsar.common.api.Commands;
 import org.apache.pulsar.common.api.proto.PulsarApi.CommandSubscribe.SubType;
-import org.apache.pulsar.common.api.proto.PulsarApi.MessageMetadata;
 import org.apache.pulsar.common.util.Murmur3_32Hash;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 public class PersistentStickyKeyDispatcherMultipleConsumers extends PersistentDispatcherMultipleConsumers {
 
-    public static final String NONE_KEY = "NONE_KEY";
-
     private final StickyKeyConsumerSelector selector;
 
     PersistentStickyKeyDispatcherMultipleConsumers(PersistentTopic topic, ManagedCursor cursor, Subscription subscription) {
@@ -171,21 +164,6 @@ protected Set<? extends Position> asyncReplayEntries(Set<? extends Position> pos
         return cursor.asyncReplayEntries(positions, this, ReadType.Replay, true);
     }
 
-    private byte[] peekStickyKey(ByteBuf metadataAndPayload) {
-        metadataAndPayload.markReaderIndex();
-        MessageMetadata metadata = Commands.parseMessageMetadata(metadataAndPayload);
-        metadataAndPayload.resetReaderIndex();
-        String key = metadata.getPartitionKey();
-        if (log.isDebugEnabled()) {
-            log.debug("Parse message metadata, partition key is {}, ordering key is {}", key, metadata.getOrderingKey());
-        }
-        if (StringUtils.isNotBlank(key) || metadata.hasOrderingKey()) {
-            return metadata.hasOrderingKey() ? metadata.getOrderingKey().toByteArray() : key.getBytes();
-        }
-        metadata.recycle();
-        return NONE_KEY.getBytes();
-    }
-
     private static final Logger log = LoggerFactory.getLogger(PersistentStickyKeyDispatcherMultipleConsumers.class);
 
 }