io.streamnative.pulsar.handlers.kop.storage.ProducerStateManagerSnapshotSerDes Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of pulsar-protocol-handler-kafka Show documentation
Kafka on Pulsar implemented using Pulsar Protocol Handler
There is a newer version: 3.3.1.5
Show newest version
/**
 * Copyright (c) 2019 - 2024 StreamNative, Inc.. All Rights Reserved.
 */
package io.streamnative.pulsar.handlers.kop.storage;

import java.io.ByteArrayInputStream;
import java.io.DataInputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.TreeMap;
import lombok.extern.slf4j.Slf4j;
import org.apache.pulsar.metadata.api.MetadataSerde;
import org.apache.pulsar.metadata.api.Stat;

/**
 * Note: to keep the compatibility, it should not be used in {@link ProducerStateManagerSnapshotBufferPartition}.
 */
@Slf4j
public class ProducerStateManagerSnapshotSerDes implements MetadataSerde {

    private static final byte[] MAGIC_NULL = { 0x7F };
    private static final byte[] MAGIC_V0 = { 0x00 };

    @Override
    public byte[] serialize(String path, ProducerStateManagerSnapshot snapshot) throws IOException {
        if (snapshot == null) {
            return Arrays.copyOf(MAGIC_NULL, MAGIC_NULL.length);
        }

        int bufferSize = 1/* magic */ + 4/* topicPartition */ + 4/* topicUUID */ + 8/* offset */
            + (4 + snapshot.producers().size() * ProducerStateEntry.SIZE) // producers
            + (4 + snapshot.ongoingTxns().size() * TxnMetadata.SIZE) // ongoingTxns
            + (4 + snapshot.abortedIndexList().size() * AbortedTxn.SIZE); // abortedIndexList
        final var topicPartitionBytes = snapshot.topicPartition().getBytes(StandardCharsets.UTF_8);
        bufferSize += topicPartitionBytes.length;
        final var topicUUIDBytes = snapshot.topicUUID().getBytes(StandardCharsets.UTF_8);
        bufferSize += topicUUIDBytes.length;
        if (bufferSize < 0) {
            throw new IOException("ProducerStateManagerSnapshot buffer overflow: " + bufferSize);
        }
        final var stream = new ByteArrayStream(bufferSize);
        stream.write(MAGIC_V0);
        stream.writeInt(topicPartitionBytes.length);
        stream.write(topicPartitionBytes);
        stream.writeInt(topicUUIDBytes.length);
        stream.write(topicUUIDBytes);
        stream.writeLong(snapshot.offset());
        stream.writeInt(snapshot.producers().size());
        for (final var producerStateEntry : snapshot.producers().values()) {
            producerStateEntry.writeTo(stream);
        }
        stream.writeInt(snapshot.ongoingTxns().size());
        for (final var txn : snapshot.ongoingTxns().values()) {
            txn.writeTo(stream);
        }
        stream.writeInt(snapshot.abortedIndexList().size());
        for (final var abortedTxn : snapshot.abortedIndexList()) {
            abortedTxn.writeTo(stream);
        }
        return stream.getWriteBuffer();
    }

    @Override
    public ProducerStateManagerSnapshot deserialize(String path, byte[] content, Stat stat) throws IOException {
        try (final var stream = new DataInputStream(new ByteArrayInputStream(content))) {
            final var magic = new byte[1];
            if (stream.readNBytes(magic, 0, 1) < 1) {
                throw new IOException("Failed to read 1 byte");
            }
            if (Arrays.equals(magic, MAGIC_NULL)) {
                return null;
            } else if (!Arrays.equals(magic, MAGIC_V0)) {
                // For the invalid magic byte, we just log it and return null.
                // Otherwise, it will throw an exception and the snapshot will not able to be written.
                log.error("Only v0 format is supported yet");
                return null;
            }
            final var topicPartition = readString(stream);
            final var topicUUID = readString(stream);
            final var offset = stream.readLong();
            final var producers = new HashMap();
            int n = stream.readInt();
            for (int i = 0; i < n; i++) {
                final var producer = ProducerStateEntry.readFrom(stream);
                producers.put(producer.producerId, producer);
            }
            n = stream.readInt();
            final var ongoingTxns = new TreeMap();
            for (int i = 0; i < n; i++) {
                final var txn = TxnMetadata.readFrom(stream);
                ongoingTxns.put(txn.firstOffset, txn);
            }
            n = stream.readInt();
            final var abortedTxnList = new ArrayList();
            for (int i = 0; i < n; i++) {
                abortedTxnList.add(AbortedTxn.readFrom(stream));
            }
            return new ProducerStateManagerSnapshot(topicPartition, topicUUID, offset, producers, ongoingTxns,
                abortedTxnList);
        }
    }

    private static String readString(DataInputStream stream) throws IOException {
        final var bytes = new byte[stream.readInt()];
        final var n = stream.readNBytes(bytes, 0, bytes.length);
        if (n < bytes.length) {
            throw new IOException("Failed to read " + bytes.length + " bytes");
        }
        return new String(bytes, StandardCharsets.UTF_8);
    }
}