All Downloads are FREE. Search and download functionalities are using the official Maven repository.

io.streamnative.pulsar.handlers.kop.utils.MessageMetadataUtils Maven / Gradle / Ivy

There is a newer version: 4.0.0.4
Show newest version
/**
 * 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.utils;

import io.jsonwebtoken.lang.Collections;
import io.netty.buffer.ByteBuf;
import io.streamnative.pulsar.handlers.kop.exceptions.MetadataCorruptedException;
import java.util.Iterator;
import java.util.concurrent.CompletableFuture;
import java.util.function.Predicate;
import javax.annotation.Nullable;
import javax.ws.rs.NotSupportedException;
import lombok.AllArgsConstructor;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.apache.bookkeeper.client.LedgerHandle;
import org.apache.bookkeeper.client.api.LedgerEntry;
import org.apache.bookkeeper.mledger.AsyncCallbacks;
import org.apache.bookkeeper.mledger.Entry;
import org.apache.bookkeeper.mledger.ManagedLedger;
import org.apache.bookkeeper.mledger.ManagedLedgerException;
import org.apache.bookkeeper.mledger.Position;
import org.apache.bookkeeper.mledger.impl.ManagedLedgerImpl;
import org.apache.bookkeeper.mledger.impl.PositionImpl;
import org.apache.pulsar.broker.intercept.ManagedLedgerInterceptorImpl;
import org.apache.pulsar.broker.service.BrokerService;
import org.apache.pulsar.client.api.RawMessage;
import org.apache.pulsar.client.impl.RawMessageImpl;
import org.apache.pulsar.common.api.proto.BrokerEntryMetadata;
import org.apache.pulsar.common.api.proto.MessageMetadata;
import org.apache.pulsar.common.intercept.AppendIndexMetadataInterceptor;
import org.apache.pulsar.common.intercept.BrokerEntryMetadataInterceptor;
import org.apache.pulsar.common.protocol.Commands;
import org.apache.pulsar.common.util.FutureUtil;
import org.apache.pulsar.compaction.CompactedTopicContext;
import org.apache.pulsar.compaction.PulsarTopicCompactionService;
import org.apache.pulsar.compaction.TopicCompactionService;

/**
 * Utils for Pulsar MessageId.
 */
@Slf4j
public class MessageMetadataUtils {

    public static boolean isBrokerIndexMetadataInterceptorConfigured(BrokerService brokerService) {
        if (Collections.isEmpty(brokerService.getBrokerEntryMetadataInterceptors())) {
            return false;
        }
        for (BrokerEntryMetadataInterceptor interceptor : brokerService.getBrokerEntryMetadataInterceptors()) {
            if (interceptor instanceof AppendIndexMetadataInterceptor) {
                return true;
            }
        }
        return false;
    }

    public static long getCurrentOffset(ManagedLedger managedLedger) {
        final var interceptor = managedLedger.getManagedLedgerInterceptor();
        if (interceptor != null) {
            return ((ManagedLedgerInterceptorImpl) interceptor).getIndex();
        } else {
            return -1L;
        }
    }

    public static long getHighWatermark(ManagedLedger managedLedger) {
        return getCurrentOffset(managedLedger) + 1;
    }

    public static long getLogEndOffset(ManagedLedger managedLedger) {
        return getCurrentOffset(managedLedger) + 1;
    }

    public static long getPublishTime(final ByteBuf byteBuf) throws MetadataCorruptedException {
        final int readerIndex = byteBuf.readerIndex();
        final MessageMetadata metadata = parseMessageMetadata(byteBuf);
        byteBuf.readerIndex(readerIndex);
        if (metadata.hasPublishTime()) {
            return metadata.getPublishTime();
        } else {
            throw new MetadataCorruptedException("Field 'publish_time' is not set");
        }
    }

    public static CompletableFuture getOffsetOfPosition(ManagedLedgerImpl managedLedger,
                                                              PositionImpl position,
                                                              boolean needCheckMore,
                                                              long timestamp,
                                                              boolean skipMessagesWithoutIndex) {
        final CompletableFuture future = new CompletableFuture<>();
        managedLedger.asyncReadEntry(position, new AsyncCallbacks.ReadEntryCallback() {
            @Override
            public void readEntryFailed(ManagedLedgerException exception, Object ctx) {
                if (exception instanceof ManagedLedgerException.NonRecoverableLedgerException) {
                    // The position doesn't exist, it usually happens when the rollover of managed ledger leads to
                    // the deletion of all expired ledgers. In this case, there's only one empty ledger in the managed
                    // ledger. So here we complete it with the latest offset.
                    future.complete(getLogEndOffset(managedLedger));
                } else {
                    future.completeExceptionally(exception);
                }
            }

            @Override
            public void readEntryComplete(Entry entry, Object ctx) {
                try {
                    if (needCheckMore) {
                        long offset = peekOffsetFromEntry(entry);
                        final long publishTime = getPublishTime(entry.getDataBuffer());
                        if (publishTime >= timestamp) {
                            future.complete(offset);
                        } else {
                            future.complete(offset + 1);
                        }
                    } else {
                        future.complete(peekBaseOffsetFromEntry(entry));
                    }
                } catch (MetadataCorruptedException.NoBrokerEntryMetadata e) {
                    if (skipMessagesWithoutIndex) {
                        log.warn("The entry {} doesn't have BrokerEntryMetadata, return 0 as the offset", position);
                        future.complete(0L);
                    } else {
                        future.completeExceptionally(e);
                    }
                } catch (MetadataCorruptedException e) {
                    future.completeExceptionally(e);
                } finally {
                    if (entry != null) {
                        entry.release();
                    }
                }

            }
        }, null);
        return future;
    }

    public static long peekOffsetFromEntry(Entry entry) throws MetadataCorruptedException {
        return peekOffset(entry.getDataBuffer(), entry.getPosition());
    }

    private static long peekOffset(ByteBuf buf, @Nullable Position position)
            throws MetadataCorruptedException {
        try {
            final BrokerEntryMetadata brokerEntryMetadata =
                    Commands.peekBrokerEntryMetadataIfExist(buf); // might throw IllegalArgumentException
            if (brokerEntryMetadata == null) {
                throw new MetadataCorruptedException.NoBrokerEntryMetadata();
            }
            return brokerEntryMetadata.getIndex(); // might throw IllegalStateException
        } catch (IllegalArgumentException | IllegalStateException e) {
            // This exception could be thrown by both peekBrokerEntryMetadataIfExist or null check
            throw new MetadataCorruptedException(
                    "Failed to peekOffsetFromEntry for " + position + ": " + e.getMessage());
        }
    }

    public static long peekBaseOffsetFromEntry(Entry entry) throws MetadataCorruptedException {
        return peekBaseOffset(entry.getDataBuffer(), entry.getPosition());
    }

    private static long peekBaseOffset(ByteBuf buf, @Nullable Position position)
            throws MetadataCorruptedException {
        MessageMetadata metadata = Commands.peekMessageMetadata(buf, null, 0);

        if (metadata == null) {
            throw new MetadataCorruptedException("Failed to peekMessageMetadata for " + position);
        }

        return peekBaseOffset(buf, position, metadata.getNumMessagesInBatch());
    }

    private static long peekBaseOffset(ByteBuf buf, @Nullable Position position, int numMessages)
            throws MetadataCorruptedException {
        return peekOffset(buf, position) - (numMessages - 1);
    }

    public static long peekBaseOffset(ByteBuf buf, int numMessages) throws MetadataCorruptedException {
        return peekBaseOffset(buf, null, numMessages);
    }

    public static MessageMetadata parseMessageMetadata(ByteBuf buf) throws MetadataCorruptedException {
        try {
            return Commands.parseMessageMetadata(buf);
        } catch (IllegalArgumentException e) {
            throw new MetadataCorruptedException(e.getMessage());
        }
    }

    public static CompletableFuture asyncFindPosition(final ManagedLedger managedLedger,
                                                                final long offset,
                                                                final boolean skipMessagesWithoutIndex) {
        return managedLedger.asyncFindPosition(new FindEntryByOffset(managedLedger.getName(),
                offset, skipMessagesWithoutIndex));
    }

    public static CompletableFuture asyncGetCompactedLedger(TopicCompactionService compactionService) {
        if (!(compactionService instanceof PulsarTopicCompactionService)) {
            return FutureUtil.failedFuture(
                    new NotSupportedException("Not support topic compactionService service class: "
                            + compactionService.getClass()));
        }

        var compactedTopic = ((PulsarTopicCompactionService) compactionService).getCompactedTopic();
        var compactedTopicContextFuture = compactedTopic.getCompactedTopicContextFuture();
        if (compactedTopicContextFuture == null) {
            return CompletableFuture.completedFuture(null);
        }

        return compactedTopicContextFuture.thenApply(CompactedTopicContext::getLedger);
    }

    public static CompletableFuture asyncFindPositionByCompactLedger(final LedgerHandle lh,
                                                                               final String managedLedgerName,
                                                                               final long offset,
                                                                               final boolean skipMessagesWithoutIndex) {
        if (lh.getLastAddConfirmed() < 0) {
            return CompletableFuture.completedFuture(null);
        }

        CompletableFuture promise = new CompletableFuture<>();

        Predicate predicate = rawEntryMetadata -> {
            BrokerEntryMetadata brokerEntryMetadata = rawEntryMetadata.getBrokerEntryMetadata();
            if (brokerEntryMetadata == null) {
                return skipMessagesWithoutIndex;
            } else {
                return rawEntryMetadata.getBrokerEntryMetadata().getIndex() >= offset;
            }
        };

        findFirstEntryMetadataLoop(predicate, 0L, lh.getLastAddConfirmed(), promise, null, lh);

        return promise.thenApply(rawEntryMetadata -> {
            if (rawEntryMetadata == null) {
                return null;
            }
            return rawEntryMetadata.getPosition();
        });
    }

    public static CompletableFuture asyncFindOffsetByTimestampFromCompactedLeger(
            final LedgerHandle lh,
            final long timestamp,
            boolean skipMessagesWithoutIndex) {
        if (lh.getLastAddConfirmed() < 0) {
            return CompletableFuture.completedFuture(null);
        }

        CompletableFuture promise = new CompletableFuture<>();

        Predicate predicate = rawEntryMetadata -> {
            return rawEntryMetadata.getMessageMetadata().getPublishTime() >= timestamp;
        };

        findFirstEntryMetadataLoop(predicate, 0L, lh.getLastAddConfirmed(), promise, null, lh);

        return promise.thenApply(rawEntryMetadata -> {
            if (rawEntryMetadata == null) {
                return null;
            }

            if (rawEntryMetadata.getBrokerEntryMetadata() == null) {
                if (skipMessagesWithoutIndex) {
                    return 0L;
                } else {
                    throw new RuntimeException(new MetadataCorruptedException.NoBrokerEntryMetadata());
                }
            }

            return rawEntryMetadata.getBrokerEntryMetadata().getIndex();
        });
    }

    private static void findFirstEntryMetadataLoop(final Predicate predicate,
                                                   final long start, final long end,
                                                   final CompletableFuture promise,
                                                   final RawEntryMetadata lastMatchRawEntryMetadata,
                                                   final LedgerHandle lh) {
        if (start > end) {
            promise.complete(lastMatchRawEntryMetadata);
            return;
        }

        long mid = (start + end) / 2;
        readRawEntryMetadata(lh, mid).thenAccept(rawEntryMetadata -> {
            if (predicate.test(rawEntryMetadata)) {
                findFirstEntryMetadataLoop(predicate, start, mid - 1, promise, rawEntryMetadata, lh);
            } else {
                findFirstEntryMetadataLoop(predicate, mid + 1, end, promise, lastMatchRawEntryMetadata, lh);
            }
        }).exceptionally(ex -> {
            promise.completeExceptionally(ex);
            return null;
        });
    }

    private static CompletableFuture readRawEntryMetadata(LedgerHandle lh, long index) {
        return lh.readAsync(index, index).thenApply(ledgerEntries -> {
            try (ledgerEntries) {
                Iterator iterator = ledgerEntries.iterator();
                LedgerEntry ledgerEntry = iterator.next();
                ByteBuf buf = ledgerEntry.getEntryBuffer();
                try (RawMessage m = RawMessageImpl.deserializeFrom(buf)) {
                    return RawEntryMetadata.parseFrom(m.getMessageIdData().getLedgerId(),
                            m.getMessageIdData().getEntryId(),
                            m.getHeadersAndPayload());
                }
            }
        });
    }

    @AllArgsConstructor
    private static class FindEntryByOffset implements Predicate {
        private final String name;
        private final long offset;
        private final boolean skipMessagesWithoutIndex;

        @Override
        public boolean test(Entry entry) {
            if (entry == null) {
                // `entry` should not be null, add the null check here to fix the spotbugs check
                return false;
            }
            try {
                return peekOffsetFromEntry(entry) < offset;
            } catch (MetadataCorruptedException.NoBrokerEntryMetadata ignored) {
                // When skipMessagesWithoutIndex is false, just return false to stop finding the position. Otherwise,
                // we assume the messages without BrokerEntryMetadata are produced by KoP < 2.8.0 that doesn't
                // support BrokerEntryMetadata. In this case, these messages should be older than any message produced
                // by KoP with BrokerEntryMetadata enabled.
                return skipMessagesWithoutIndex;
            } catch (MetadataCorruptedException e) {
                log.error("[{}] Entry {} is corrupted: {}",
                        name, entry.getPosition(), e.getMessage());
                return false;
            } finally {
                entry.release();
            }
        }

        @Override
        public String toString() {
            return "FindEntryByOffset{ " + offset + "}";
        }
    }

    @Getter
    private static class RawEntryMetadata {
        private Position position;
        private BrokerEntryMetadata brokerEntryMetadata;
        private MessageMetadata messageMetadata;

        static RawEntryMetadata parseFrom(long ledgerId, long entryId, ByteBuf headersAndPayload) {
            RawEntryMetadata rawEntryMetadata = new RawEntryMetadata();
            rawEntryMetadata.position = PositionImpl.get(ledgerId, entryId);
            rawEntryMetadata.brokerEntryMetadata = Commands.parseBrokerEntryMetadataIfExist(headersAndPayload);
            rawEntryMetadata.messageMetadata = new MessageMetadata();
            Commands.parseMessageMetadata(headersAndPayload, rawEntryMetadata.messageMetadata);
            return rawEntryMetadata;
        }
    }

}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy