diff --git a/managed-ledger/src/main/java/org/apache/bookkeeper/mledger/impl/LedgerMetadataUtils.java b/managed-ledger/src/main/java/org/apache/bookkeeper/mledger/impl/LedgerMetadataUtils.java index 399811e4b3b17..8571a36584e2b 100644 --- a/managed-ledger/src/main/java/org/apache/bookkeeper/mledger/impl/LedgerMetadataUtils.java +++ b/managed-ledger/src/main/java/org/apache/bookkeeper/mledger/impl/LedgerMetadataUtils.java @@ -39,12 +39,17 @@ public final class LedgerMetadataUtils { "compacted-ledger".getBytes(StandardCharsets.UTF_8); private static final byte[] METADATA_PROPERTY_COMPONENT_SCHEMA = "schema".getBytes(StandardCharsets.UTF_8); + private static final byte[] METADATA_PROPERTY_COMPONENT_DELAYED_INDEX_BUCKET = + "delayed-index-bucket".getBytes(StandardCharsets.UTF_8); + private static final String METADATA_PROPERTY_MANAGED_LEDGER_NAME = "pulsar/managed-ledger"; private static final String METADATA_PROPERTY_CURSOR_NAME = "pulsar/cursor"; private static final String METADATA_PROPERTY_COMPACTEDTOPIC = "pulsar/compactedTopic"; private static final String METADATA_PROPERTY_COMPACTEDTO = "pulsar/compactedTo"; private static final String METADATA_PROPERTY_SCHEMAID = "pulsar/schemaId"; + private static final String METADATA_PROPERTY_DELAYED_INDEX_BUCKETID = "pulsar/delayedIndexBucketId"; + /** * Build base metadata for every ManagedLedger. * @@ -100,6 +105,20 @@ public static Map buildMetadataForSchema(String schemaId) { ); } + /** + * Build additional metadata for a delayed message index bucket. + * + * @param bucketKey key of the delayed message bucket + * @return an immutable map which describes the schema + */ + public static Map buildMetadataForDelayedIndexBucket(String bucketKey) { + return Map.of( + METADATA_PROPERTY_APPLICATION, METADATA_PROPERTY_APPLICATION_PULSAR, + METADATA_PROPERTY_COMPONENT, METADATA_PROPERTY_COMPONENT_DELAYED_INDEX_BUCKET, + METADATA_PROPERTY_DELAYED_INDEX_BUCKETID, bucketKey.getBytes(StandardCharsets.UTF_8) + ); + } + /** * Build additional metadata for the placement policy config. * diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/BookkeeperBucketSnapshotStorage.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/BookkeeperBucketSnapshotStorage.java new file mode 100644 index 0000000000000..1cadc6d98e268 --- /dev/null +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/BookkeeperBucketSnapshotStorage.java @@ -0,0 +1,247 @@ +/* + * 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.delayed.bucket; + +import com.google.protobuf.InvalidProtocolBufferException; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Enumeration; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; +import javax.validation.constraints.NotNull; +import lombok.extern.slf4j.Slf4j; +import org.apache.bookkeeper.client.BKException; +import org.apache.bookkeeper.client.BookKeeper; +import org.apache.bookkeeper.client.LedgerEntry; +import org.apache.bookkeeper.client.LedgerHandle; +import org.apache.bookkeeper.mledger.impl.LedgerMetadataUtils; +import org.apache.pulsar.broker.PulsarService; +import org.apache.pulsar.broker.ServiceConfiguration; +import org.apache.pulsar.broker.delayed.proto.DelayedMessageIndexBucketSnapshotFormat.SnapshotMetadata; +import org.apache.pulsar.broker.delayed.proto.DelayedMessageIndexBucketSnapshotFormat.SnapshotSegment; +import org.apache.pulsar.common.util.FutureUtil; + +@Slf4j +public class BookkeeperBucketSnapshotStorage implements BucketSnapshotStorage { + + private static final byte[] LedgerPassword = "".getBytes(); + + private final PulsarService pulsar; + private final ServiceConfiguration config; + private BookKeeper bookKeeper; + + public BookkeeperBucketSnapshotStorage(PulsarService pulsar) { + this.pulsar = pulsar; + this.config = pulsar.getConfig(); + } + + @Override + public CompletableFuture createBucketSnapshot(SnapshotMetadata snapshotMetadata, + List bucketSnapshotSegments, + String bucketKey) { + return createLedger(bucketKey) + .thenCompose(ledgerHandle -> addEntry(ledgerHandle, snapshotMetadata.toByteArray()) + .thenCompose(__ -> addSnapshotSegments(ledgerHandle, bucketSnapshotSegments)) + .thenCompose(__ -> closeLedger(ledgerHandle)) + .thenApply(__ -> ledgerHandle.getId())); + } + + @Override + public CompletableFuture getBucketSnapshotMetadata(long bucketId) { + return openLedger(bucketId).thenCompose( + ledgerHandle -> getLedgerEntryThenCloseLedger(ledgerHandle, 0, 0). + thenApply(entryEnumeration -> parseSnapshotMetadataEntry(entryEnumeration.nextElement()))); + } + + @Override + public CompletableFuture> getBucketSnapshotSegment(long bucketId, long firstSegmentEntryId, + long lastSegmentEntryId) { + return openLedger(bucketId).thenCompose( + ledgerHandle -> getLedgerEntryThenCloseLedger(ledgerHandle, firstSegmentEntryId, + lastSegmentEntryId).thenApply(this::parseSnapshotSegmentEntries)); + } + + @Override + public CompletableFuture getBucketSnapshotLength(long bucketId) { + return openLedger(bucketId).thenApply(ledgerHandle -> { + long length = ledgerHandle.getLength(); + closeLedger(ledgerHandle); + return length; + }); + } + + @Override + public CompletableFuture deleteBucketSnapshot(long bucketId) { + return deleteLedger(bucketId); + } + + @Override + public void start() throws Exception { + this.bookKeeper = pulsar.getBookKeeperClientFactory().create( + pulsar.getConfiguration(), + pulsar.getLocalMetadataStore(), + pulsar.getIoEventLoopGroup(), + Optional.empty(), + null + ); + } + + @Override + public void close() throws Exception { + if (bookKeeper != null) { + bookKeeper.close(); + } + } + + private CompletableFuture addSnapshotSegments(LedgerHandle ledgerHandle, + List bucketSnapshotSegments) { + List> addFutures = new ArrayList<>(); + for (SnapshotSegment bucketSnapshotSegment : bucketSnapshotSegments) { + addFutures.add(addEntry(ledgerHandle, bucketSnapshotSegment.toByteArray())); + } + + return FutureUtil.waitForAll(addFutures); + } + + private SnapshotMetadata parseSnapshotMetadataEntry(LedgerEntry ledgerEntry) { + try { + return SnapshotMetadata.parseFrom(ledgerEntry.getEntry()); + } catch (InvalidProtocolBufferException e) { + throw new BucketSnapshotSerializationException(e); + } + } + + private List parseSnapshotSegmentEntries(Enumeration entryEnumeration) { + List snapshotMetadataList = new ArrayList<>(); + try { + while (entryEnumeration.hasMoreElements()) { + LedgerEntry ledgerEntry = entryEnumeration.nextElement(); + snapshotMetadataList.add(SnapshotSegment.parseFrom(ledgerEntry.getEntry())); + } + return snapshotMetadataList; + } catch (IOException e) { + throw new BucketSnapshotSerializationException(e); + } + } + + @NotNull + private CompletableFuture createLedger(String bucketKey) { + CompletableFuture future = new CompletableFuture<>(); + Map metadata = LedgerMetadataUtils.buildMetadataForDelayedIndexBucket(bucketKey); + bookKeeper.asyncCreateLedger( + config.getManagedLedgerDefaultEnsembleSize(), + config.getManagedLedgerDefaultWriteQuorum(), + config.getManagedLedgerDefaultAckQuorum(), + BookKeeper.DigestType.fromApiDigestType(config.getManagedLedgerDigestType()), + LedgerPassword, + (rc, handle, ctx) -> { + if (rc != BKException.Code.OK) { + future.completeExceptionally(bkException("Failed to create ledger", rc, -1)); + } else { + future.complete(handle); + } + }, null, metadata); + return future; + } + + private CompletableFuture openLedger(Long ledgerId) { + final CompletableFuture future = new CompletableFuture<>(); + bookKeeper.asyncOpenLedger( + ledgerId, + BookKeeper.DigestType.fromApiDigestType(config.getManagedLedgerDigestType()), + LedgerPassword, + (rc, handle, ctx) -> { + if (rc != BKException.Code.OK) { + future.completeExceptionally(bkException("Failed to open ledger", rc, ledgerId)); + } else { + future.complete(handle); + } + }, null + ); + return future; + } + + private CompletableFuture closeLedger(LedgerHandle ledgerHandle) { + CompletableFuture future = new CompletableFuture<>(); + ledgerHandle.asyncClose((rc, handle, ctx) -> { + if (rc != BKException.Code.OK) { + log.warn("Failed to close a Ledger Handle: {}", ledgerHandle.getId()); + future.completeExceptionally(bkException("Failed to close ledger", rc, ledgerHandle.getId())); + } else { + future.complete(null); + } + }, null); + return future; + } + + private CompletableFuture addEntry(LedgerHandle ledgerHandle, byte[] data) { + final CompletableFuture future = new CompletableFuture<>(); + ledgerHandle.asyncAddEntry(data, + (rc, handle, entryId, ctx) -> { + if (rc != BKException.Code.OK) { + future.completeExceptionally(bkException("Failed to add entry", rc, ledgerHandle.getId())); + } else { + future.complete(null); + } + }, null + ); + + return future.whenComplete((__, ex) -> { + if (ex != null) { + deleteLedger(ledgerHandle.getId()); + } + }); + } + + CompletableFuture> getLedgerEntryThenCloseLedger(LedgerHandle ledger, + long firstEntryId, long lastEntryId) { + final CompletableFuture> future = new CompletableFuture<>(); + ledger.asyncReadEntries(firstEntryId, lastEntryId, + (rc, handle, entries, ctx) -> { + if (rc != BKException.Code.OK) { + future.completeExceptionally(bkException("Failed to read entry", rc, ledger.getId())); + } else { + future.complete(entries); + } + closeLedger(handle); + }, null + ); + return future; + } + + private CompletableFuture deleteLedger(long ledgerId) { + CompletableFuture future = new CompletableFuture<>(); + bookKeeper.asyncDeleteLedger(ledgerId, (int rc, Object cnx) -> { + if (rc != BKException.Code.OK) { + future.completeExceptionally(bkException("Failed to delete ledger", rc, ledgerId)); + } else { + future.complete(null); + } + }, null); + return future; + } + + private static BucketSnapshotPersistenceException bkException(String operation, int rc, long ledgerId) { + String message = BKException.getMessage(rc) + + " - ledger=" + ledgerId + " - operation=" + operation; + return new BucketSnapshotPersistenceException(message); + } +} diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/Bucket.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/Bucket.java index fbd6d765705d4..2a7ee4d196a09 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/Bucket.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/Bucket.java @@ -123,13 +123,13 @@ long getAndUpdateBucketId() { } CompletableFuture asyncSaveBucketSnapshot( - ImmutableBucket bucketState, DelayedMessageIndexBucketSnapshotFormat.SnapshotMetadata snapshotMetadata, + ImmutableBucket bucket, DelayedMessageIndexBucketSnapshotFormat.SnapshotMetadata snapshotMetadata, List bucketSnapshotSegments) { - - return bucketSnapshotStorage.createBucketSnapshot(snapshotMetadata, bucketSnapshotSegments) + final String bucketKey = bucket.bucketKey(); + return bucketSnapshotStorage.createBucketSnapshot(snapshotMetadata, bucketSnapshotSegments, bucketKey) .thenCompose(newBucketId -> { - bucketState.setBucketId(newBucketId); - String bucketKey = bucketState.bucketKey(); + bucket.setBucketId(newBucketId); + return putBucketKeyId(bucketKey, newBucketId).exceptionally(ex -> { log.warn("Failed to record bucketId to cursor property, bucketKey: {}", bucketKey); return null; diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/BucketSnapshotPersistenceException.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/BucketSnapshotPersistenceException.java new file mode 100644 index 0000000000000..210bbb9165eec --- /dev/null +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/BucketSnapshotPersistenceException.java @@ -0,0 +1,32 @@ +/* + * 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.delayed.bucket; + +import org.apache.pulsar.broker.service.BrokerServiceException; + +public class BucketSnapshotPersistenceException extends BrokerServiceException.PersistenceException { + + public BucketSnapshotPersistenceException(Throwable t) { + super(t); + } + + public BucketSnapshotPersistenceException(String msg) { + super(msg); + } +} diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/BucketSnapshotSerializationException.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/BucketSnapshotSerializationException.java new file mode 100644 index 0000000000000..463574e730a64 --- /dev/null +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/BucketSnapshotSerializationException.java @@ -0,0 +1,30 @@ +/* + * 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.delayed.bucket; + +public class BucketSnapshotSerializationException extends RuntimeException { + + public BucketSnapshotSerializationException(String message) { + super(message); + } + + public BucketSnapshotSerializationException(Throwable t) { + super(t); + } +} diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/BucketSnapshotStorage.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/BucketSnapshotStorage.java index 3ab4ce1ad2792..c6941e289f1ac 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/BucketSnapshotStorage.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/BucketSnapshotStorage.java @@ -30,10 +30,12 @@ public interface BucketSnapshotStorage { * * @param snapshotMetadata the metadata of snapshot * @param bucketSnapshotSegments the list of snapshot segments + * @param bucketKey the key of bucket is used to generate custom storage metadata * @return the future with bucketId(ledgerId). */ CompletableFuture createBucketSnapshot(SnapshotMetadata snapshotMetadata, - List bucketSnapshotSegments); + List bucketSnapshotSegments, + String bucketKey); /** * Get delayed message index bucket snapshot metadata. @@ -47,8 +49,8 @@ CompletableFuture createBucketSnapshot(SnapshotMetadata snapshotMetadata, * Get a sequence of delayed message index bucket snapshot segments. * * @param bucketId the bucketId of snapshot - * @param firstSegmentEntryId entryId of first segment of sequence - * @param lastSegmentEntryId entryId of last segment of sequence + * @param firstSegmentEntryId entryId of first segment of sequence (include) + * @param lastSegmentEntryId entryId of last segment of sequence (include) * @return the future with snapshot segment */ CompletableFuture> getBucketSnapshotSegment(long bucketId, long firstSegmentEntryId, diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/BrokerServiceException.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/BrokerServiceException.java index eb073762f232f..fd3a391bca34a 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/BrokerServiceException.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/BrokerServiceException.java @@ -88,6 +88,10 @@ public static class PersistenceException extends BrokerServiceException { public PersistenceException(Throwable t) { super(t); } + + public PersistenceException(String msg) { + super(msg); + } } public static class TopicTerminatedException extends BrokerServiceException { diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/delayed/BookkeeperBucketSnapshotStorageTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/delayed/BookkeeperBucketSnapshotStorageTest.java new file mode 100644 index 0000000000000..1effe756ff2ab --- /dev/null +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/delayed/BookkeeperBucketSnapshotStorageTest.java @@ -0,0 +1,201 @@ +/* + * 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.delayed; + +import com.google.protobuf.ByteString; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import org.apache.pulsar.broker.auth.MockedPulsarServiceBaseTest; +import org.apache.pulsar.broker.delayed.bucket.BookkeeperBucketSnapshotStorage; +import org.apache.pulsar.broker.delayed.proto.DelayedMessageIndexBucketSnapshotFormat; +import org.testng.Assert; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +public class BookkeeperBucketSnapshotStorageTest extends MockedPulsarServiceBaseTest { + + private BookkeeperBucketSnapshotStorage bucketSnapshotStorage; + + @BeforeClass + @Override + protected void setup() throws Exception { + super.internalSetup(); + bucketSnapshotStorage = new BookkeeperBucketSnapshotStorage(pulsar); + bucketSnapshotStorage.start(); + } + + @AfterClass + @Override + protected void cleanup() throws Exception { + super.internalCleanup(); + bucketSnapshotStorage.close(); + } + + @Test + public void testCreateSnapshot() throws ExecutionException, InterruptedException { + DelayedMessageIndexBucketSnapshotFormat.SnapshotMetadata snapshotMetadata = + DelayedMessageIndexBucketSnapshotFormat.SnapshotMetadata.newBuilder().build(); + List bucketSnapshotSegments = new ArrayList<>(); + CompletableFuture future = + bucketSnapshotStorage.createBucketSnapshot(snapshotMetadata, + bucketSnapshotSegments, UUID.randomUUID().toString()); + Long bucketId = future.get(); + Assert.assertNotNull(bucketId); + } + + @Test + public void testGetSnapshot() throws ExecutionException, InterruptedException { + DelayedMessageIndexBucketSnapshotFormat.SnapshotSegmentMetadata segmentMetadata = + DelayedMessageIndexBucketSnapshotFormat.SnapshotSegmentMetadata.newBuilder() + .setMaxScheduleTimestamp(System.currentTimeMillis()) + .putDelayedIndexBitMap(100L, ByteString.copyFrom(new byte[1])).build(); + + DelayedMessageIndexBucketSnapshotFormat.SnapshotMetadata snapshotMetadata = + DelayedMessageIndexBucketSnapshotFormat.SnapshotMetadata.newBuilder() + .addMetadataList(segmentMetadata) + .build(); + List bucketSnapshotSegments = new ArrayList<>(); + + long timeMillis = System.currentTimeMillis(); + DelayedMessageIndexBucketSnapshotFormat.DelayedIndex delayedIndex = + DelayedMessageIndexBucketSnapshotFormat.DelayedIndex.newBuilder().setLedgerId(100L).setEntryId(10L) + .setTimestamp(timeMillis).build(); + DelayedMessageIndexBucketSnapshotFormat.SnapshotSegment snapshotSegment = + DelayedMessageIndexBucketSnapshotFormat.SnapshotSegment.newBuilder().addIndexes(delayedIndex).build(); + bucketSnapshotSegments.add(snapshotSegment); + bucketSnapshotSegments.add(snapshotSegment); + + CompletableFuture future = + bucketSnapshotStorage.createBucketSnapshot(snapshotMetadata, + bucketSnapshotSegments, UUID.randomUUID().toString()); + Long bucketId = future.get(); + Assert.assertNotNull(bucketId); + + CompletableFuture> bucketSnapshotSegment = + bucketSnapshotStorage.getBucketSnapshotSegment(bucketId, 1, 3); + + List snapshotSegments = bucketSnapshotSegment.get(); + Assert.assertEquals(2, snapshotSegments.size()); + for (DelayedMessageIndexBucketSnapshotFormat.SnapshotSegment segment : snapshotSegments) { + for (DelayedMessageIndexBucketSnapshotFormat.DelayedIndex index : segment.getIndexesList()) { + Assert.assertEquals(100L, index.getLedgerId()); + Assert.assertEquals(10L, index.getEntryId()); + Assert.assertEquals(timeMillis, index.getTimestamp()); + } + } + } + + @Test + public void testGetSnapshotMetadata() throws ExecutionException, InterruptedException { + long timeMillis = System.currentTimeMillis(); + + Map map = new HashMap<>(); + map.put(100L, ByteString.copyFrom("test1", StandardCharsets.UTF_8)); + map.put(200L, ByteString.copyFrom("test2", StandardCharsets.UTF_8)); + + DelayedMessageIndexBucketSnapshotFormat.SnapshotSegmentMetadata segmentMetadata = + DelayedMessageIndexBucketSnapshotFormat.SnapshotSegmentMetadata.newBuilder() + .setMaxScheduleTimestamp(timeMillis) + .putAllDelayedIndexBitMap(map).build(); + + DelayedMessageIndexBucketSnapshotFormat.SnapshotMetadata snapshotMetadata = + DelayedMessageIndexBucketSnapshotFormat.SnapshotMetadata.newBuilder() + .addMetadataList(segmentMetadata) + .build(); + List bucketSnapshotSegments = new ArrayList<>(); + + CompletableFuture future = + bucketSnapshotStorage.createBucketSnapshot(snapshotMetadata, + bucketSnapshotSegments, UUID.randomUUID().toString()); + Long bucketId = future.get(); + Assert.assertNotNull(bucketId); + + DelayedMessageIndexBucketSnapshotFormat.SnapshotMetadata bucketSnapshotMetadata = + bucketSnapshotStorage.getBucketSnapshotMetadata(bucketId).get(); + + DelayedMessageIndexBucketSnapshotFormat.SnapshotSegmentMetadata metadata = + bucketSnapshotMetadata.getMetadataList(0); + + Assert.assertEquals(timeMillis, metadata.getMaxScheduleTimestamp()); + Assert.assertEquals("test1", metadata.getDelayedIndexBitMapMap().get(100L).toStringUtf8()); + Assert.assertEquals("test2", metadata.getDelayedIndexBitMapMap().get(200L).toStringUtf8()); + } + + @Test + public void testDeleteSnapshot() throws ExecutionException, InterruptedException { + DelayedMessageIndexBucketSnapshotFormat.SnapshotMetadata snapshotMetadata = + DelayedMessageIndexBucketSnapshotFormat.SnapshotMetadata.newBuilder().build(); + List bucketSnapshotSegments = new ArrayList<>(); + CompletableFuture future = + bucketSnapshotStorage.createBucketSnapshot(snapshotMetadata, + bucketSnapshotSegments, UUID.randomUUID().toString()); + Long bucketId = future.get(); + Assert.assertNotNull(bucketId); + + bucketSnapshotStorage.deleteBucketSnapshot(bucketId).get(); + + try { + bucketSnapshotStorage.getBucketSnapshotMetadata(bucketId).get(); + Assert.fail("Should fail"); + } catch (Exception e) { + Assert.assertTrue(e.getCause().getMessage().contains("No such ledger exists")); + } + } + + @Test + public void testGetBucketSnapshotLength() throws ExecutionException, InterruptedException { + DelayedMessageIndexBucketSnapshotFormat.SnapshotSegmentMetadata segmentMetadata = + DelayedMessageIndexBucketSnapshotFormat.SnapshotSegmentMetadata.newBuilder() + .setMaxScheduleTimestamp(System.currentTimeMillis()) + .putDelayedIndexBitMap(100L, ByteString.copyFrom(new byte[1])).build(); + + DelayedMessageIndexBucketSnapshotFormat.SnapshotMetadata snapshotMetadata = + DelayedMessageIndexBucketSnapshotFormat.SnapshotMetadata.newBuilder() + .addMetadataList(segmentMetadata) + .build(); + List bucketSnapshotSegments = new ArrayList<>(); + + long timeMillis = System.currentTimeMillis(); + DelayedMessageIndexBucketSnapshotFormat.DelayedIndex delayedIndex = + DelayedMessageIndexBucketSnapshotFormat.DelayedIndex.newBuilder().setLedgerId(100L).setEntryId(10L) + .setTimestamp(timeMillis).build(); + DelayedMessageIndexBucketSnapshotFormat.SnapshotSegment snapshotSegment = + DelayedMessageIndexBucketSnapshotFormat.SnapshotSegment.newBuilder().addIndexes(delayedIndex).build(); + bucketSnapshotSegments.add(snapshotSegment); + bucketSnapshotSegments.add(snapshotSegment); + + CompletableFuture future = + bucketSnapshotStorage.createBucketSnapshot(snapshotMetadata, + bucketSnapshotSegments, UUID.randomUUID().toString()); + Long bucketId = future.get(); + Assert.assertNotNull(bucketId); + + Long bucketSnapshotLength = bucketSnapshotStorage.getBucketSnapshotLength(bucketId).get(); + System.out.println(bucketSnapshotLength); + Assert.assertTrue(bucketSnapshotLength > 0L); + } + +} diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/delayed/MockBucketSnapshotStorage.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/delayed/MockBucketSnapshotStorage.java index 89831a1d5e771..9b2fbda4195da 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/delayed/MockBucketSnapshotStorage.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/delayed/MockBucketSnapshotStorage.java @@ -55,7 +55,7 @@ public MockBucketSnapshotStorage() { @Override public CompletableFuture createBucketSnapshot( - SnapshotMetadata snapshotMetadata, List bucketSnapshotSegments) { + SnapshotMetadata snapshotMetadata, List bucketSnapshotSegments, String bucketKey) { return CompletableFuture.supplyAsync(() -> { long bucketId = maxBucketId.getAndIncrement(); List entries = new ArrayList<>();