[Event Hubs] Implement PartitionManager methods

sdk/eventhub/event-hubs/review/event-hubs.api.md

@@ -28,6 +28,22 @@ export interface BatchOptions {
     partitionKey?: string;
 }
 
+// @public
+export interface Checkpoint {
+    // (undocumented)
+    consumerGroupName: string;
+    // (undocumented)
+    eventHubName: string;
+    // (undocumented)
+    instanceId: string;
+    // (undocumented)
+    offset: number;
+    // (undocumented)
+    partitionId: string;
+    // (undocumented)
+    sequenceNumber: number;
+}
+
 export { DataTransformer }
 
 export { DefaultDataTransformer }
@@ -156,6 +172,16 @@ export class EventProcessor {
     stop(): Promise<void>;
 }
 
+// @public
+export class InMemoryPartitionManager implements PartitionManager {
+    // (undocumented)
+    claimOwnerships(partitionOwnerships: PartitionOwnership[]): Promise<PartitionOwnership[]>;
+    // (undocumented)
+    listOwnerships(eventHubName: string, consumerGroupName: string): Promise<PartitionOwnership[]>;
+    // (undocumented)
+    updateCheckpoint(checkpoint: Checkpoint): Promise<void>;
+}
+
 export { MessagingError }
 
 // @public
@@ -174,6 +200,28 @@ export interface PartitionContext {
     readonly partitionId: string;
 }
 
+// @public
+export interface PartitionOwnership {
+    // (undocumented)
+    consumerGroupName: string;
+    // (undocumented)
+    ETag?: string;
+    // (undocumented)
+    eventHubName: string;
+    // (undocumented)
+    instanceId: string;
+    // (undocumented)
+    lastModifiedTime?: number;
+    // (undocumented)
+    offset?: number;
+    // (undocumented)
+    ownerLevel: number;
+    // (undocumented)
+    partitionId: string;
+    // (undocumented)
+    sequenceNumber?: number;
+}
+
 // @public
 export interface PartitionProperties {
     beginningSequenceNumber: number;

sdk/eventhub/event-hubs/src/eventProcessor.ts

@@ -60,7 +60,7 @@ export interface PartitionProcessorFactory {
 export interface PartitionManager {
   listOwnerships(eventHubName: string, consumerGroupName: string): Promise<PartitionOwnership[]>;
   claimOwnerships(partitionOwnerships: PartitionOwnership[]): Promise<PartitionOwnership[]>;
-  createCheckpoint(checkpoint: Checkpoint): Promise<void>;
+  updateCheckpoint(checkpoint: Checkpoint): Promise<void>;
 }
 
 // Options passed when creating EventProcessor, everything is optional

sdk/eventhub/event-hubs/src/inMemoryPartitionManager.ts

@@ -0,0 +1,40 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { PartitionManager, PartitionOwnership } from "./eventProcessor";
+import { Checkpoint } from "./checkpointManager";
+import { generate_uuid } from "rhea-promise";
+
+/**
+ * A simple in-memory implementation of a `PartitionManager`
+ * @class
+ */
+export class InMemoryPartitionManager implements PartitionManager {
+  private _partitionOwnershipMap: Map<string, PartitionOwnership> = new Map();
+
+  async listOwnerships(
+    eventHubName: string,
+    consumerGroupName: string
+  ): Promise<PartitionOwnership[]> {
+    return Array.from(this._partitionOwnershipMap.values());
+  }
+
+  async claimOwnerships(partitionOwnerships: PartitionOwnership[]): Promise<PartitionOwnership[]> {
+    for (let partitionOwnership of partitionOwnerships) {
+      if (!this._partitionOwnershipMap.has(partitionOwnership.partitionId)) {
+        partitionOwnership.ETag = generate_uuid();
+        this._partitionOwnershipMap.set(partitionOwnership.partitionId, partitionOwnership);
+      }
+    }
+    return partitionOwnerships;
+  }
+
+  async updateCheckpoint(checkpoint: Checkpoint): Promise<void> {
+    let partitionOwnership = this._partitionOwnershipMap.get(checkpoint.partitionId);
+    if (partitionOwnership) {
+      partitionOwnership.sequenceNumber = checkpoint.sequenceNumber;
+      partitionOwnership.offset = checkpoint.offset;
+      partitionOwnership.ETag = generate_uuid();
+    }
+  }
+}

sdk/eventhub/event-hubs/src/index.ts

@@ -21,8 +21,10 @@ export { PartitionProperties, EventHubProperties } from "./managementClient";
 export { EventHubProducer } from "./sender";
 export { EventHubConsumer, EventIteratorOptions } from "./receiver";
 export { EventDataBatch } from "./eventDataBatch";
-export { EventProcessor } from "./eventProcessor";
+export { EventProcessor, PartitionOwnership } from "./eventProcessor";
 export { PartitionContext } from "./partitionContext";
+export { InMemoryPartitionManager} from "./inMemoryPartitionManager"
+export { Checkpoint } from "./checkpointManager";
 export {
   MessagingError,
   DataTransformer,

sdk/eventhub/event-hubs/test/eventProcessor.spec.ts

@@ -13,9 +13,13 @@ import {
   EventData,
   EventProcessor,
   PartitionContext,
-  delay
+  delay,
+  InMemoryPartitionManager,
+  PartitionOwnership,
+  Checkpoint
 } from "../src";
 import { EnvVarKeys, getEnvVars } from "./utils/testUtils";
+import { generate_uuid } from "rhea-promise";
 const env = getEnvVars();
 
 describe("Event Processor", function(): void {
@@ -96,4 +100,58 @@ describe("Event Processor", function(): void {
       isCloseCalled.should.equal(true);
     });
   });
+
+  describe("InMemory Partition Manager", function(): void {
+    const inMemoryPartitionManager = new InMemoryPartitionManager();
+
+    it("should claim ownership ", async function(): Promise<void> {
+      const partitionOwnership1: PartitionOwnership = {
+        eventHubName: "myEventHub",
+        consumerGroupName: EventHubClient.defaultConsumerGroupName,
+        instanceId: generate_uuid(),
+        partitionId: "0",
+        ownerLevel: 10
+      };
+      const partitionOwnership2: PartitionOwnership = {
+        eventHubName: "myEventHub",
+        consumerGroupName: EventHubClient.defaultConsumerGroupName,
+        instanceId: generate_uuid(),
+        partitionId: "1",
+        ownerLevel: 10
+      };
+      const partitionOwnership = await inMemoryPartitionManager.claimOwnerships([
+        partitionOwnership1,
+        partitionOwnership2
+      ]);
+      partitionOwnership.length.should.equals(2);
+    });
+
+    it("should get a list of ownerships ", async function(): Promise<void> {
+      const partitionOwnershipList = await inMemoryPartitionManager.listOwnerships(
+        "myEventHub",
+        EventHubClient.defaultConsumerGroupName
+      );
+      partitionOwnershipList.length.should.equals(2);
+    });
+
+    it("should update checkpoint ", async function(): Promise<void> {
+      const checkpoint: Checkpoint = {
+        eventHubName: "myEventHub",
+        consumerGroupName: EventHubClient.defaultConsumerGroupName,
+        instanceId: generate_uuid(),
+        partitionId: "0",
+        sequenceNumber: 10,
+        offset: 50
+      };
+
+      await inMemoryPartitionManager.updateCheckpoint(checkpoint);
+      const partitionOwnershipList = await inMemoryPartitionManager.listOwnerships(
+        "myEventHub",
+        EventHubClient.defaultConsumerGroupName
+      );
+      partitionOwnershipList[0].partitionId.should.equals("0");
+      partitionOwnershipList[0].sequenceNumber!.should.equals(10);
+      partitionOwnershipList[0].offset!.should.equals(50);
+    });
+  });
 }).timeout(90000);