(WIP) Adding support for geo-dr and offsets.

eng/versioning/version_client.txt

@@ -492,6 +492,12 @@ io.clientcore:http-stress;1.0.0-beta.1;1.0.0-beta.1
 # In the pom, the version update tag after the version should name the unreleased package and the dependency version:
 # <!-- {x-version-update;unreleased_com.azure:azure-core;dependency} -->
 
+unreleased_com.azure:azure-core-amqp;2.10.0-beta.1
+
+unreleased_com.azure:azure-messaging-servicebus;7.18.0-beta.2
+unreleased_com.azure:azure-messaging-eventhubs;5.20.0-beta.1
+unreleased_com.azure:azure-messaging-eventhubs-checkpointstore-blob;1.21.0-beta.1
+
 unreleased_com.azure:azure-monitor-opentelemetry-exporter;1.0.0-beta.31
 unreleased_com.azure:azure-monitor-opentelemetry-autoconfigure;1.0.0-beta.1
 

sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/AmqpMessageConstant.java

@@ -126,7 +126,11 @@ public enum AmqpMessageConstant {
     /**
      * The state of message.
      */
-    MESSAGE_STATE_ANNOTATION_NAME("x-opt-message-state");
+    MESSAGE_STATE_ANNOTATION_NAME("x-opt-message-state"),
+    /**
+     * The replication segment for the message.
+     */
+    REPLICATION_SEGMENT_ANNOTATION_NAME("x-opt-sequence-number-epoch");
 
     private static final Map<String, AmqpMessageConstant> RESERVED_CONSTANTS_MAP = new HashMap<>();
     private final String constant;

sdk/eventhubs/azure-messaging-eventhubs-checkpointstore-blob/CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## 1.21.0-beta.1 (Unreleased)
 
+- Added feature to persist/rehydrate replication segment field to/from Azure Storage Blobs.
+
 ### Features Added
 
 ### Breaking Changes

sdk/eventhubs/azure-messaging-eventhubs-checkpointstore-blob/src/main/java/com/azure/messaging/eventhubs/checkpointstore/blob/BlobCheckpointStore.java

@@ -41,7 +41,7 @@
  * @see EventProcessorClient
  */
 public class BlobCheckpointStore implements CheckpointStore {
-
+    static final String REPLICATION_SEGMENT = "replicationsegment";
     static final String SEQUENCE_NUMBER = "sequencenumber";
     static final String OFFSET = "offset";
     static final String OWNER_ID = "ownerid";
@@ -57,6 +57,7 @@ public class BlobCheckpointStore implements CheckpointStore {
     private static final String SEQUENCE_NUMBER_LOG_KEY = "sequenceNumber";
     private static final String BLOB_NAME_LOG_KEY = "blobName";
     private static final String OFFSET_LOG_KEY = "offset";
+    private static final String REPLICATION_SEGMENT_LOG_KEY = "replicationSegment";
 
     /**
      * An empty string.
@@ -137,10 +138,12 @@ private Mono<Checkpoint> convertToCheckpoint(BlobItem blobItem) {
                 .addKeyValue(BLOB_NAME_LOG_KEY, blobItem.getName())
                 .addKeyValue(SEQUENCE_NUMBER_LOG_KEY, metadata.get(SEQUENCE_NUMBER))
                 .addKeyValue(OFFSET_LOG_KEY, metadata.get(OFFSET))
+                .addKeyValue(REPLICATION_SEGMENT_LOG_KEY, metadata.get(REPLICATION_SEGMENT))
                 .log(Messages.CHECKPOINT_INFO);
 
             Long sequenceNumber = null;
             Long offset = null;
+            Integer replicationSegment = null;
             if (!CoreUtils.isNullOrEmpty(metadata.get(SEQUENCE_NUMBER))) {
                 sequenceNumber = Long.parseLong(metadata.get(SEQUENCE_NUMBER));
             }
@@ -149,13 +152,19 @@ private Mono<Checkpoint> convertToCheckpoint(BlobItem blobItem) {
                 offset = Long.parseLong(metadata.get(OFFSET));
             }
 
-            Checkpoint checkpoint = new Checkpoint().setFullyQualifiedNamespace(names[0])
+            if (!CoreUtils.isNullOrEmpty(metadata.get(REPLICATION_SEGMENT))) {
+                replicationSegment = Integer.parseInt(metadata.get(REPLICATION_SEGMENT));
+            }
+
+            Checkpoint checkpoint = new Checkpoint()
+                .setFullyQualifiedNamespace(names[0])
                 .setEventHubName(names[1])
                 .setConsumerGroup(names[2])
                 // names[3] is "checkpoint"
                 .setPartitionId(names[4])
                 .setSequenceNumber(sequenceNumber)
-                .setOffset(offset);
+                .setOffset(offset)
+                .setReplicationSegment(replicationSegment);
 
             return Mono.just(checkpoint);
         }
@@ -238,20 +247,29 @@ public Mono<Void> updateCheckpoint(Checkpoint checkpoint) {
                 "Both sequence number and offset cannot be null when updating a checkpoint")));
         }
 
-        String partitionId = checkpoint.getPartitionId();
-        String blobName = getBlobName(checkpoint.getFullyQualifiedNamespace(), checkpoint.getEventHubName(),
+        final String partitionId = checkpoint.getPartitionId();
+        final String blobName = getBlobName(checkpoint.getFullyQualifiedNamespace(), checkpoint.getEventHubName(),
             checkpoint.getConsumerGroup(), partitionId, CHECKPOINT_PATH);
+
         if (!blobClients.containsKey(blobName)) {
             blobClients.put(blobName, blobContainerAsyncClient.getBlobAsyncClient(blobName));
         }
 
-        Map<String, String> metadata = new HashMap<>();
-        String sequenceNumber
-            = checkpoint.getSequenceNumber() == null ? null : String.valueOf(checkpoint.getSequenceNumber());
-
-        String offset = checkpoint.getOffset() == null ? null : String.valueOf(checkpoint.getOffset());
+        final String sequenceNumber = checkpoint.getSequenceNumber() != null
+            ? String.valueOf(checkpoint.getSequenceNumber())
+            : null;
+        final String offset = checkpoint.getOffset() != null
+            ? String.valueOf(checkpoint.getOffset())
+            : null;
+        final String replicationSegment = checkpoint.getReplicationSegment() != null
+            ? String.valueOf(checkpoint.getReplicationSegment())
+            : null;
+
+        final Map<String, String> metadata = new HashMap<>();
         metadata.put(SEQUENCE_NUMBER, sequenceNumber);
         metadata.put(OFFSET, offset);
+        metadata.put(REPLICATION_SEGMENT, replicationSegment);
+
         BlobAsyncClient blobAsyncClient = blobClients.get(blobName);
 
         return blobAsyncClient.exists().flatMap(exists -> {

sdk/eventhubs/azure-messaging-eventhubs-checkpointstore-blob/src/test/java/com/azure/messaging/eventhubs/checkpointstore/blob/BlobCheckpointStoreTests.java

@@ -25,6 +25,7 @@
 import org.junit.jupiter.api.condition.DisabledOnJre;
 import org.junit.jupiter.api.condition.JRE;
 import org.mockito.ArgumentMatchers;
+import org.mockito.Captor;
 import org.mockito.Mock;
 import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
@@ -44,15 +45,18 @@
 import static com.azure.messaging.eventhubs.checkpointstore.blob.BlobCheckpointStore.OFFSET;
 import static com.azure.messaging.eventhubs.checkpointstore.blob.BlobCheckpointStore.OWNERSHIP_PATH;
 import static com.azure.messaging.eventhubs.checkpointstore.blob.BlobCheckpointStore.OWNER_ID;
+import static com.azure.messaging.eventhubs.checkpointstore.blob.BlobCheckpointStore.REPLICATION_SEGMENT;
 import static com.azure.messaging.eventhubs.checkpointstore.blob.BlobCheckpointStore.SEQUENCE_NUMBER;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyMap;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.ArgumentMatchers.isNull;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 /**
@@ -70,11 +74,14 @@ public class BlobCheckpointStoreTests {
     @Mock
     private BlobAsyncClient blobAsyncClient;
 
+    @Captor
+    private ArgumentCaptor<Map<String, String>> metadataArgumentCaptor;
+
     private AutoCloseable autoCloseable;
 
     @BeforeEach
     public void beforeEach() {
-        this.autoCloseable = MockitoAnnotations.openMocks(this);
+        autoCloseable = MockitoAnnotations.openMocks(this);
     }
 
     @AfterEach
@@ -164,11 +171,11 @@ public void testListCheckpoint() {
         final String checkpointPrefix = prefix + CHECKPOINT_PATH;
 
         final BlobCheckpointStore blobCheckpointStore = new BlobCheckpointStore(blobContainerAsyncClient);
-        final BlobItem blobItem = getCheckpointBlobItem("230", "1", checkpointPrefix + "0"); // valid blob
+        final BlobItem blobItem = getCheckpointBlobItem("230", "1", checkpointPrefix + "0", "35"); // valid blob
         final BlobItem blobItem2 = new BlobItem().setName(checkpointPrefix + "1"); // valid blob but not a valid checkpoint.
-        final BlobItem blobItem3 = getCheckpointBlobItem("233", "3", prefix + "1"); // invalid name
-        final PagedFlux<BlobItem> response
-            = new PagedFlux<>(() -> Mono.just(new PagedResponseBase<HttpHeaders, BlobItem>(null, 200, null,
+        final BlobItem blobItem3 = getCheckpointBlobItem("233", "3", prefix + "1", null); // invalid name
+        final PagedFlux<BlobItem> response = new PagedFlux<>(() -> Mono.just(
+            new PagedResponseBase<HttpHeaders, BlobItem>(null, 200, null,
                 Arrays.asList(blobItem, blobItem2, blobItem3), null, null)));
 
         when(blobContainerAsyncClient.listBlobs(any(ListBlobsOptions.class))).thenAnswer(invocation -> {
@@ -192,6 +199,7 @@ public void testListCheckpoint() {
                 assertEquals(consumerGroup, checkpoint.getConsumerGroup());
                 assertEquals(1L, checkpoint.getSequenceNumber());
                 assertEquals(230L, checkpoint.getOffset());
+                assertEquals(35, checkpoint.getReplicationSegment());
             })
             .verifyComplete();
     }
@@ -227,25 +235,41 @@ public void testUpdateCheckpoint() {
             .setEventHubName(eventHubName)
             .setConsumerGroup(consumerGroup)
             .setPartitionId(partitionId)
+            .setReplicationSegment(16)
             .setSequenceNumber(2L)
             .setOffset(100L);
 
-        final BlobItem blobItem = getCheckpointBlobItem("230", "1", blobName);
-        final PagedFlux<BlobItem> response
-            = new PagedFlux<>(() -> Mono.just(new PagedResponseBase<HttpHeaders, BlobItem>(null, 200, null,
+        final BlobItem blobItem = getCheckpointBlobItem("230", "1", blobName, "25");
+        final PagedFlux<BlobItem> response = new PagedFlux<>(() -> Mono.just(
+            new PagedResponseBase<HttpHeaders, BlobItem>(null, 200, null,
                 Collections.singletonList(blobItem), null, null)));
 
         when(blobContainerAsyncClient.getBlobAsyncClient(blobName)).thenReturn(blobAsyncClient);
         when(blobContainerAsyncClient.listBlobs(any(ListBlobsOptions.class))).thenReturn(response);
         when(blobAsyncClient.getBlockBlobAsyncClient()).thenReturn(blockBlobAsyncClient);
         when(blobAsyncClient.exists()).thenReturn(Mono.just(true));
-
-        when(blobAsyncClient.setMetadata(ArgumentMatchers.<Map<String, String>>any())).thenReturn(Mono.empty());
+        when(blobAsyncClient.setMetadata(ArgumentMatchers.any())).thenReturn(Mono.empty());
 
         final BlobCheckpointStore blobCheckpointStore = new BlobCheckpointStore(blobContainerAsyncClient);
 
         // Act & Assert
-        StepVerifier.create(blobCheckpointStore.updateCheckpoint(checkpoint)).verifyComplete();
+        StepVerifier.create(blobCheckpointStore.updateCheckpoint(checkpoint))
+            .verifyComplete();
+
+        verify(blobAsyncClient).setMetadata(metadataArgumentCaptor.capture());
+
+        final Map<String, String> actual = metadataArgumentCaptor.getValue();
+
+        assertEquals(3, actual.size());
+
+        assertTrue(actual.containsKey(REPLICATION_SEGMENT));
+        assertEquals(String.valueOf(checkpoint.getReplicationSegment()), actual.get(REPLICATION_SEGMENT));
+
+        assertTrue(actual.containsKey(OFFSET));
+        assertEquals(String.valueOf(checkpoint.getOffset()), actual.get(OFFSET));
+
+        assertTrue(actual.containsKey(SEQUENCE_NUMBER));
+        assertEquals(String.valueOf(checkpoint.getSequenceNumber()), actual.get(SEQUENCE_NUMBER));
     }
 
     /**
@@ -271,15 +295,16 @@ public void testUpdateCheckpointForNewPartition() {
             .setConsumerGroup("cg")
             .setPartitionId("0")
             .setSequenceNumber(2L)
-            .setOffset(100L);
+            .setOffset(100L)
+            .setReplicationSegment(15);
         final String legacyPrefix = getLegacyPrefix(checkpoint.getFullyQualifiedNamespace(),
             checkpoint.getEventHubName(), checkpoint.getConsumerGroup());
         final String blobName = legacyPrefix + CHECKPOINT_PATH + checkpoint.getPartitionId();
 
         HttpHeaders httpHeaders = new HttpHeaders();
         httpHeaders.add(HttpHeaderName.ETAG, "etag2");
 
-        BlobItem blobItem = getCheckpointBlobItem("230", "1", blobName);
+        BlobItem blobItem = getCheckpointBlobItem("230", "1", blobName, "20");
 
         PagedFlux<BlobItem> response
             = new PagedFlux<BlobItem>(() -> Mono.just(new PagedResponseBase<HttpHeaders, BlobItem>(null, 200, null,
@@ -291,12 +316,31 @@ public void testUpdateCheckpointForNewPartition() {
         when(blobAsyncClient.getBlockBlobAsyncClient()).thenReturn(blockBlobAsyncClient);
         when(blobAsyncClient.exists()).thenReturn(Mono.just(false));
 
-        when(blockBlobAsyncClient.uploadWithResponse(ArgumentMatchers.<Flux<ByteBuffer>>any(), eq(0L), isNull(),
-            anyMap(), isNull(), isNull(), isNull()))
-                .thenReturn(Mono.just(new ResponseBase<>(null, 200, httpHeaders, null, null)));
+        when(blockBlobAsyncClient.uploadWithResponse(ArgumentMatchers.any(), eq(0L),
+            isNull(), anyMap(), isNull(), isNull(), isNull()))
+            .thenReturn(Mono.just(new ResponseBase<>(null, 200, httpHeaders, null, null)));
 
         BlobCheckpointStore blobCheckpointStore = new BlobCheckpointStore(blobContainerAsyncClient);
-        StepVerifier.create(blobCheckpointStore.updateCheckpoint(checkpoint)).verifyComplete();
+
+        // Act & Assert
+        StepVerifier.create(blobCheckpointStore.updateCheckpoint(checkpoint))
+            .verifyComplete();
+
+        verify(blockBlobAsyncClient).uploadWithResponse(any(), eq(0L), isNull(),
+            metadataArgumentCaptor.capture(), isNull(), isNull(), isNull());
+
+        final Map<String, String> actual = metadataArgumentCaptor.getValue();
+
+        assertEquals(3, actual.size());
+
+        assertTrue(actual.containsKey(REPLICATION_SEGMENT));
+        assertEquals(String.valueOf(checkpoint.getReplicationSegment()), actual.get(REPLICATION_SEGMENT));
+
+        assertTrue(actual.containsKey(OFFSET));
+        assertEquals(String.valueOf(checkpoint.getOffset()), actual.get(OFFSET));
+
+        assertTrue(actual.containsKey(SEQUENCE_NUMBER));
+        assertEquals(String.valueOf(checkpoint.getSequenceNumber()), actual.get(SEQUENCE_NUMBER));
     }
 
     /**
@@ -443,11 +487,20 @@ private static BlobItem getOwnershipBlobItem(String owner, String etag, String b
         return new BlobItem().setName(blobName).setMetadata(metadata).setProperties(properties);
     }
 
-    private static BlobItem getCheckpointBlobItem(String offset, String sequenceNumber, String blobName) {
-        Map<String, String> metadata = new HashMap<>();
+    private static BlobItem getCheckpointBlobItem(String offset, String sequenceNumber, String blobName,
+        String replicationSegment) {
+
+        final Map<String, String> metadata = new HashMap<>();
         metadata.put(SEQUENCE_NUMBER, sequenceNumber);
         metadata.put(OFFSET, offset);
-        return new BlobItem().setName(blobName).setMetadata(metadata);
+
+        if (replicationSegment != null) {
+            metadata.put(REPLICATION_SEGMENT, replicationSegment);
+        }
+
+        return new BlobItem()
+            .setName(blobName)
+            .setMetadata(metadata);
     }
 
     private static String getLegacyPrefix(String fullyQualifiedNamespace, String eventHubName, String consumerGroup) {

sdk/eventhubs/azure-messaging-eventhubs/CHANGELOG.md

@@ -4,7 +4,12 @@
 
 ### Features Added
 
-Setting the v2 stack as the default. ([43725](https://github.com/Azure/azure-sdk-for-java/pull/43725))
+- Setting the v2 stack as the default. ([43725](https://github.com/Azure/azure-sdk-for-java/pull/43725)).
+- Added support geo-replication feature.
+- Added `Checkpoint.getReplicationSegment()`.
+- Added `getBeginningReplicationSegment()` and `getLastEnqueuedReplicationSegment()` to `PartitionProperties`.
+- Added `LastEnqueuedEventProperties.getReplicationSegment()`.
+- Added overloads to `EventPosition.fromSequenceNumber` that takes replication segment.
 
 ### Breaking Changes
 

sdk/eventhubs/azure-messaging-eventhubs/pom.xml

@@ -53,7 +53,7 @@
     <dependency>
       <groupId>com.azure</groupId>
       <artifactId>azure-core-amqp</artifactId>
-      <version>2.9.12</version> <!-- {x-version-update;com.azure:azure-core-amqp;dependency} -->
+      <version>2.10.0-beta.1</version> <!-- {x-version-update;unreleased_com.azure:azure-core-amqp;dependency} -->
     </dependency>
 
     <!-- Test dependencies -->

sdk/eventhubs/azure-messaging-eventhubs/src/main/java/com/azure/messaging/eventhubs/EventData.java

@@ -271,11 +271,24 @@ public EventData setBodyAsBinaryData(BinaryData binaryData) {
      *
      * @return The offset within the Event Hub partition of the received event. {@code null} if the {@link EventData}
      *     was not received from Event Hubs service.
+     * @deprecated This value is obsolete and should no longer be used. Please use {@link #getOffsetString()} instead.
      */
+    @Deprecated
     public Long getOffset() {
         return systemProperties.getOffset();
     }
 
+    /**
+     * Gets the offset of the event when it was received from the associated Event Hub partition. This is only present
+     * on a <b>received</b> {@link EventData}.
+     *
+     * @return The offset within the Event Hub partition of the received event. {@code null} if the {@link EventData}
+     *     was not received from Event Hubs service.
+     */
+    public String getOffsetString() {
+        return systemProperties.getOffsetString();
+    }
+
     /**
      * Gets the partition hashing key if it was set when originally publishing the event. If it exists, this value was
      * used to compute a hash to select a partition to send the message to. This is only present on a <b>received</b>
@@ -397,6 +410,15 @@ public EventData setMessageId(String messageId) {
         return this;
     }
 
+    /**
+     * Gets the replication segment for the event.
+     *
+     * @return The replication segment.  -1 or null if geo-disaster recovery is not enabled.
+     */
+    public Integer getReplicationSegment() {
+        return systemProperties.getReplicationSegment();
+    }
+
     /**
      * True if the object is an {@link EventData} and the binary contents of {@link #getBody()} are equal.
      */