[PIP 96] Add message payload processor for Pulsar client

pulsar-broker/src/test/java/org/apache/pulsar/client/processor/CustomBatchFormat.java

@@ -0,0 +1,104 @@
+/**
+ * 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.processor;
+
+import io.netty.buffer.ByteBuf;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import lombok.AllArgsConstructor;
+import lombok.Getter;
+import org.apache.pulsar.client.api.Schema;
+import org.apache.pulsar.common.allocator.PulsarByteBufAllocator;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+/**
+ * A batch message whose format is customized.
+ *
+ * 1. First 2 bytes represent the number of messages.
+ * 2. Each message is a string, whose format is
+ *   1. First 2 bytes represent the length `N`.
+ *   2. Followed N bytes are the bytes of the string.
+ */
+public class CustomBatchFormat {
+
+    public static final String KEY = "entry.format";
+    public static final String VALUE = "custom";
+
+    @AllArgsConstructor
+    @Getter
+    public static class Metadata {
+        private final int numMessages;
+    }
+
+    public static ByteBuf serialize(Iterable<String> strings) {
+        final ByteBuf buf = PulsarByteBufAllocator.DEFAULT.buffer(1024);
+        buf.writeShort(0);
+        short numMessages = 0;
+        for (String s : strings) {
+            writeString(buf, s);
+            numMessages++;
+        }
+        buf.setShort(0, numMessages);
+        return buf;
+    }
+
+    private static void writeString(final ByteBuf buf, final String s) {
+        final byte[] bytes = Schema.STRING.encode(s);
+        buf.writeShort(bytes.length);
+        buf.writeBytes(bytes);
+    }
+
+    public static Metadata readMetadata(final ByteBuf buf) {
+        return new Metadata(buf.readShort());
+    }
+
+    public static byte[] readMessage(final ByteBuf buf) {
+        final short length = buf.readShort();
+        final byte[] bytes = new byte[length];
+        buf.readBytes(bytes);
+        return bytes;
+    }
+
+    @Test
+    public void testMultipleStrings() {
+        final List<List<String>> inputs = new ArrayList<>();
+        inputs.add(Collections.emptyList());
+        inputs.add(Collections.singletonList("java"));
+        inputs.add(Arrays.asList("hello", "world", "java"));
+
+        for (List<String> input : inputs) {
+            final ByteBuf buf = serialize(input);
+
+            final Metadata metadata = readMetadata(buf);
+            final List<String> parsedTokens = new ArrayList<>();
+            for (int i = 0; i < metadata.getNumMessages(); i++) {
+                parsedTokens.add(Schema.STRING.decode(readMessage(buf)));
+            }
+
+            Assert.assertEquals(parsedTokens, input);
+            Assert.assertEquals(parsedTokens.size(), input.size());
+
+            Assert.assertEquals(buf.refCnt(), 1);
+            buf.release();
+        }
+    }
+}

pulsar-broker/src/test/java/org/apache/pulsar/client/processor/CustomBatchPayloadProcessor.java

@@ -0,0 +1,61 @@
+/**
+ * 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.processor;
+
+import io.netty.buffer.ByteBuf;
+import java.util.function.Consumer;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.pulsar.client.api.MessagePayloadContext;
+import org.apache.pulsar.client.api.Message;
+import org.apache.pulsar.client.api.MessagePayload;
+import org.apache.pulsar.client.api.MessagePayloadFactory;
+import org.apache.pulsar.client.api.MessagePayloadProcessor;
+import org.apache.pulsar.client.api.Schema;
+import org.apache.pulsar.client.impl.MessagePayloadUtils;
+
+@Slf4j
+public class CustomBatchPayloadProcessor implements MessagePayloadProcessor {
+
+    @Override
+    public <T> void process(MessagePayload payload, MessagePayloadContext context, Schema<T> schema,
+                            Consumer<Message<T>> messageConsumer) throws Exception {
+        final String value = context.getProperty(CustomBatchFormat.KEY);
+        if (value == null || !value.equals(CustomBatchFormat.VALUE)) {
+            DEFAULT.process(payload, context, schema, messageConsumer);
+            return;
+        }
+
+        final ByteBuf buf = MessagePayloadUtils.convertToByteBuf(payload);
+        try {
+            final int numMessages = CustomBatchFormat.readMetadata(buf).getNumMessages();
+            for (int i = 0; i < numMessages; i++) {
+                final MessagePayload singlePayload =
+                        MessagePayloadFactory.DEFAULT.wrap(CustomBatchFormat.readMessage(buf));
+                try {
+                    messageConsumer.accept(
+                            context.getMessageAt(i, numMessages, singlePayload, false, schema));
+                } finally {
+                    singlePayload.release();
+                }
+            }
+        } finally {
+            buf.release();
+        }
+    }
+}

pulsar-broker/src/test/java/org/apache/pulsar/client/processor/CustomBatchProducer.java

@@ -0,0 +1,70 @@
+/**
+ * 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.processor;
+
+import io.netty.buffer.ByteBuf;
+import java.util.ArrayList;
+import java.util.List;
+import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.pulsar.broker.service.persistent.PersistentTopic;
+import org.apache.pulsar.common.api.proto.MessageMetadata;
+import org.apache.pulsar.common.protocol.Commands;
+
+@RequiredArgsConstructor
+@Slf4j
+public class CustomBatchProducer {
+
+    private final List<String> messages = new ArrayList<>();
+    private final PersistentTopic persistentTopic;
+    private final int batchingMaxMessages;
+
+    public void sendAsync(final String value) {
+        messages.add(value);
+        if (messages.size() >= batchingMaxMessages) {
+            flush();
+        }
+    }
+
+    public void flush() {
+        final ByteBuf buf = CustomBatchFormat.serialize(messages);
+        final ByteBuf headerAndPayload = Commands.serializeMetadataAndPayload(Commands.ChecksumType.None,
+                createCustomMetadata(), buf);
+        buf.release();
+        persistentTopic.publishMessage(headerAndPayload, (e, ledgerId, entryId) -> {
+            if (e == null) {
+                log.info("Send successfully to {} ({}, {})", persistentTopic.getName(), ledgerId, entryId);
+            } else {
+                log.error("Failed to send: {}", e.getMessage());
+            }
+        });
+        messages.clear();
+    }
+
+    private static MessageMetadata createCustomMetadata() {
+        final MessageMetadata messageMetadata = new MessageMetadata();
+        // Here are required fields
+        messageMetadata.setProducerName("");
+        messageMetadata.setSequenceId(0L);
+        messageMetadata.setPublishTime(0L);
+        // Add the property to identify the message format
+        messageMetadata.addProperty().setKey(CustomBatchFormat.KEY).setValue(CustomBatchFormat.VALUE);
+        return messageMetadata;
+    }
+}

pulsar-broker/src/test/java/org/apache/pulsar/client/processor/DefaultProcessorWithRefCnt.java

@@ -0,0 +1,47 @@
+/**
+ * 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.processor;
+
+import java.util.function.Consumer;
+import lombok.Getter;
+import org.apache.pulsar.client.api.MessagePayloadContext;
+import org.apache.pulsar.client.api.Message;
+import org.apache.pulsar.client.api.MessagePayload;
+import org.apache.pulsar.client.api.MessagePayloadProcessor;
+import org.apache.pulsar.client.api.Schema;
+import org.apache.pulsar.client.impl.MessagePayloadImpl;
+
+/**
+ * The processor for Pulsar format messages and maintains a total reference count.
+ *
+ * It's used to verify {@link MessagePayloadContext#getMessageAt} and {@link MessagePayloadContext#asSingleMessage} have release the
+ * ByteBuf successfully.
+ */
+public class DefaultProcessorWithRefCnt implements MessagePayloadProcessor {
+
+    @Getter
+    int totalRefCnt = 0;
+
+    @Override
+    public <T> void process(MessagePayload payload, MessagePayloadContext context, Schema<T> schema,
+                            Consumer<Message<T>> messageConsumer) throws Exception {
+        MessagePayloadProcessor.DEFAULT.process(payload, context, schema, messageConsumer);
+        totalRefCnt += ((MessagePayloadImpl) payload).getByteBuf().refCnt();
+    }
+}