io.streamnative.pulsar.handlers.kop.KafkaPayloadProcessor Maven / Gradle / Ivy

Go to download
/**
 * Copyright (c) 2019 - 2024 StreamNative, Inc.. All Rights Reserved.
 */
/**
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package io.streamnative.pulsar.handlers.kop;

import io.netty.buffer.ByteBuf;
import io.netty.util.concurrent.FastThreadLocal;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.util.function.Consumer;
import org.apache.kafka.common.header.Header;
import org.apache.kafka.common.record.MemoryRecords;
import org.apache.kafka.common.record.Record;
import org.apache.kafka.common.record.RecordBatch;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.apache.pulsar.client.api.Message;
import org.apache.pulsar.client.api.MessagePayload;
import org.apache.pulsar.client.api.MessagePayloadContext;
import org.apache.pulsar.client.api.MessagePayloadProcessor;
import org.apache.pulsar.client.api.Schema;
import org.apache.pulsar.client.impl.MessagePayloadImpl;
import org.apache.pulsar.client.impl.MessagePayloadUtils;
import org.apache.pulsar.common.allocator.PulsarByteBufAllocator;
import org.apache.pulsar.common.api.proto.SingleMessageMetadata;

/**
 * Process Kafka messages so that Pulsar consumer can recognize.
 */
public class KafkaPayloadProcessor implements MessagePayloadProcessor {

    static final String ENTRY_ORIGINAL_BASEOFFSET_KEY = "entry.originalOffset";
    private static final StringDeserializer deserializer = new StringDeserializer();
    private static final FastThreadLocal LOCAL_SINGLE_MESSAGE_METADATA =
            new FastThreadLocal() {
                @Override
                protected SingleMessageMetadata initialValue() throws Exception {
                    return new SingleMessageMetadata();
                }
            };

    @Override
    public  void process(MessagePayload payload,
                            MessagePayloadContext context,
                            Schema schema,
                            Consumer> messageConsumer) throws Exception {
        if (!isKafkaFormat(context)) {
            DEFAULT.process(payload, context, schema, messageConsumer);
            return;
        }

        final ByteBuf buf = MessagePayloadUtils.convertToByteBuf(payload);
        try {
            final MemoryRecords records = MemoryRecords.readableRecords(buf.nioBuffer());
            final RecordBatch firstBatch = records.firstBatch();
            if (firstBatch == null) {
                return;
            }
            final int numMessages = context.getNumMessages();
            long baseOffset = baseOffset(firstBatch, context);
            for (RecordBatch batch : records.batches()) {
                if (batch.isControlBatch()) {
                    continue;
                }
                // TODO: Currently KoP doesn't support multi batches in an entry so it works well at this moment. After
                //  we supported multi batches in future, the following code should be changed. See
                //  https://github.com/streamnative/kop/issues/537 for details.
                for (Record record : batch) {
                    final MessagePayload singlePayload = newByteBufFromRecord(record);
                    try {
                        int batchIndex = (int) (record.offset() - baseOffset);
                        messageConsumer.accept(
                                context.getMessageAt(batchIndex, numMessages, singlePayload, true, schema));
                    } finally {
                        singlePayload.release();
                    }
                }
            }
        } finally {
            buf.release();
        }
    }

    private static boolean isKafkaFormat(final MessagePayloadContext context) {
        final String value = context.getProperty("entry.format");
        return value != null && value.equalsIgnoreCase("kafka");
    }

    private static byte[] bufferToBytes(final ByteBuffer buffer) {
        final byte[] data = new byte[buffer.remaining()];
        buffer.get(data);
        return data;
    }

    private MessagePayload newByteBufFromRecord(final Record record) {
        final SingleMessageMetadata singleMessageMetadata = LOCAL_SINGLE_MESSAGE_METADATA.get();
        singleMessageMetadata.clear();
        if (record.hasKey()) {
            final byte[] data = bufferToBytes(record.key());
            // It's okay to pass a null topic because it's not used in StringDeserializer
            singleMessageMetadata.setPartitionKey(deserializer.deserialize(null, data));
            singleMessageMetadata.setOrderingKey(data);
        }

        final ByteBuffer valueBuffer;
        if (record.hasValue()) {
            valueBuffer = record.value();
            singleMessageMetadata.setPayloadSize(record.valueSize());
        } else {
            valueBuffer = null;
            singleMessageMetadata.setNullValue(true);
            singleMessageMetadata.setPayloadSize(0);
        }

        for (Header header : record.headers()) {
            singleMessageMetadata.addProperty()
                    .setKey(header.key())
                    .setValue(new String(header.value(), StandardCharsets.UTF_8));
        }

        final ByteBuf buf = PulsarByteBufAllocator.DEFAULT.buffer(
                4 + singleMessageMetadata.getSerializedSize() + record.valueSize());
        buf.writeInt(singleMessageMetadata.getSerializedSize());
        singleMessageMetadata.writeTo(buf);
        if (valueBuffer != null) {
            buf.writeBytes(valueBuffer);
        }
        return MessagePayloadImpl.create(buf);
    }

    private static long baseOffset(final RecordBatch recordBatch, final MessagePayloadContext context) {
        if (recordBatch.magic() >= RecordBatch.MAGIC_VALUE_V2) {
            return recordBatch.baseOffset();
        }

        String property = context.getProperty(ENTRY_ORIGINAL_BASEOFFSET_KEY);
        return property != null ? Long.parseLong(property) : recordBatch.baseOffset();
    }
}