EventProcessor Load balancing

sdk/eventhub/azure-eventhubs/azure/eventhub/eventprocessor/_ownership_manager.py

@@ -3,7 +3,6 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # -----------------------------------------------------------------------------------
 
-from typing import List, Iterable, Any, Dict
 import time
 import random
 import math
@@ -34,8 +33,8 @@ async def claim_ownership(self):
         """Claims ownership for this EventProcessor
         1. Retrieves all partition ids of an event hub from azure event hub service
         2. Retrieves current ownership list via this EventProcessor's PartitionManager.
-        3. Searches claimable partitions for this EventProcessor. Refer to claim_ownership() for details.
-        4. Claims the ownership for the claimable partitions
+        3. Balances number of ownership. Refer to _balance_ownership() for details.
+        4. Claims the ownership for the balanced number of partitions.
 
         :return: List[Dict[Any]]
         """
@@ -53,19 +52,35 @@ async def _retrieve_partition_ids(self):
         self.all_parition_ids = await self.eventhub_client.get_partition_ids()
 
     async def _balance_ownership(self):
+        """Balances and claims ownership of partitions for this EventProcessor.
+        The balancing algorithm is:
+        1. Find partitions with inactive ownership and partitions that haven never been claimed before
+        2. Find the number of active owners, including this EventProcessor, for all partitions.
+        3. Calculate the average count of partitions that an owner should own.
+        (number of partitions // number of active owners)
+        4. Calculate the largest allowed count of partitions that an owner can own.
+        math.ceil(number of partitions / number of active owners). This should be equal or 1 greater than the average count
+        5. Adjust the number of partitions owned by this EventProcessor (owner)
+            a. if this EventProcessor owns more than largest allowed count, abandon one partition
+            b. if this EventProcessor owns less than average count, add one from the inactive or unclaimed partitions,
+                or steal one from another owner that has the largest number of ownership among all owners (EventProcessors)
+            c. Otherwise, no change to the ownership
+
+        The balancing algorithm adjust one partition at a time to gradually build the balanced ownership.
+        Ownership must be renewed to keep it active. So the returned result includes both existing ownership and
+        the newly adjusted ownership.
+        This method balances but doesn't claim ownership. The caller of this method tries to claim the result ownership
+        list. But it may not successfully claim all of them because of concurrency. Other EventProcessors may happen to
+        claim a partition at that time. Since balancing and claiming are run in infinite repeatedly,
+        it achieves balancing among all EventProcessors after some time of running.
+
+        :return: List[Dict[str, Any]], A list of ownership.
+        """
         ownership_list = await self.partition_manager.list_ownership(self.eventhub_client.eh_name, self.consumer_group_name)
         ownership_dict = dict((x["partition_id"], x) for x in ownership_list)  # put the list to dict for fast lookup
-        '''
-        now = time.time()
-        partition_ids_no_ownership = list(filter(lambda x: x not in ownership_dict, self.all_parition_ids))
-        inactive_ownership = filter(lambda x: x["last_modified_time"] + self.ownership_timeout < now, ownership_list)
-        claimable_partition_ids = partition_ids_no_ownership + [x["partition_id"] for x in inactive_ownership]
-        active_ownership = list(filter(lambda x: x["last_modified_time"] + self.ownership_timeout >= now, ownership_list))
-        active_ownership_count_group_by_owner = Counter([x["owner_id"] for x in active_ownership])
-        active_ownership_self = list(filter(lambda x: x["owner_id"] == self.owner_id, active_ownership))
-        '''
-        claimable_partition_ids = []
-        active_ownership_self = []
+
+        claimable_partition_ids = []  # partitions with inactive ownership and partitions that have never been claimed yet
+        active_ownership_self = []  # active ownership of this EventProcessor
         active_ownership_count_group_by_owner = Counter()
         for partition_id in self.all_parition_ids:
             ownership = ownership_dict.get(partition_id)
@@ -100,7 +115,7 @@ async def _balance_ownership(self):
                                                             {"partition_id": random_partition_id,
                                                              "eventhub_name": self.eventhub_client.eh_name,
                                                              "consumer_group_name": self.consumer_group_name,
-                                                             "owner_level": 0})
+                                                             "owner_level": 0})  # TODO: consider removing owner_level
                 random_chosen_to_claim["owner_id"] = self.owner_id
                 to_claim.append(random_chosen_to_claim)
             else:  # steal from another owner that has the most count

sdk/eventhub/azure-eventhubs/azure/eventhub/eventprocessor/event_processor.py

@@ -7,15 +7,13 @@
 import uuid
 import asyncio
 import logging
-from enum import Enum
 
 from azure.eventhub import EventPosition, EventHubError
 from azure.eventhub.aio import EventHubClient
 from .checkpoint_manager import CheckpointManager
 from .partition_manager import PartitionManager
 from ._ownership_manager import OwnershipManager
-from .partition_processor import PartitionProcessor, CloseReason
-from .utils import get_running_loop
+from .partition_processor import CloseReason, PartitionProcessor
 
 log = logging.getLogger(__name__)
 
@@ -29,20 +27,21 @@ class EventProcessor(object):
 
     It provides the user a convenient way to receive events from multiple partitions and save checkpoints.
     If multiple EventProcessors are running for an event hub, they will automatically balance load.
-    This load balancing won't be available until preview 3.
 
     Example:
         .. code-block:: python
 
-            class MyPartitionProcessor(PartitionProcessor):
-                async def process_events(self, events):
-                    if events:
-                        # do something sync or async to process the events
-                        await self._checkpoint_manager.update_checkpoint(events[-1].offset, events[-1].sequence_number)
-
             import asyncio
             from azure.eventhub.aio import EventHubClient
             from azure.eventhub.eventprocessor import EventProcessor, PartitionProcessor, Sqlite3PartitionManager
+
+            class MyPartitionProcessor(object):
+                async def process_events(self, events, checkpoint_manager):
+                    if events:
+                        # do something sync or async to process the events
+                        await checkpoint_manager.update_checkpoint(events[-1].offset, events[-1].sequence_number)
+
+
             client = EventHubClient.from_connection_string("<your connection string>", receive_timeout=5, retry_total=3)
             partition_manager = Sqlite3PartitionManager()
             try:
@@ -55,7 +54,7 @@ async def process_events(self, events):
 
     """
     def __init__(self, eventhub_client: EventHubClient, consumer_group_name: str,
-                 partition_processor_factory,
+                 partition_processor_factory: Callable[..., PartitionProcessor],
                  partition_manager: PartitionManager, **kwargs):
         """
         Instantiate an EventProcessor.
@@ -73,6 +72,8 @@ def __init__(self, eventhub_client: EventHubClient, consumer_group_name: str,
         :type partition_manager: Class implementing the ~azure.eventhub.eventprocessor.PartitionManager.
         :param initial_event_position: The offset to start a partition consumer if the partition has no checkpoint yet.
         :type initial_event_position: int or str
+        :param polling_interval: The interval between any two pollings of balancing and claiming
+        :type float
 
         """
 
@@ -98,9 +99,12 @@ def __repr__(self):
     async def start(self):
         """Start the EventProcessor.
 
-        1. retrieve the partition ids from eventhubs.
-        2. claim partition ownership of these partitions.
-        3. repeatedly call EvenHubConsumer.receive() to retrieve events and call user defined PartitionProcessor.process_events().
+        1. Calls the OwnershipManager to keep claiming and balancing ownership of partitions in an
+        infinitely loop until self.stop() is called.
+        2. Cancels tasks for partitions that are no longer owned by this EventProcessor
+        3. Creates tasks for partitions that are newly claimed by this EventProcessor
+        4. Keeps tasks running for partitions that haven't changed ownership
+        5. Each task repeatedly calls EvenHubConsumer.receive() to retrieve events and call user defined partition processor
 
         :return: None
 
@@ -111,22 +115,23 @@ async def start(self):
             self._running = True
             while self._running:
                 claimed_ownership_list = await ownership_manager.claim_ownership()
-                claimed_partition_ids = [x["partition_id"] for x in claimed_ownership_list]
-                to_cancel_list = self._tasks.keys() - claimed_partition_ids
-                if to_cancel_list:
-                    self._cancel_tasks_for_partitions(to_cancel_list)
-                    log.info("EventProcesor %r has cancelled partitions %r", self._id, to_cancel_list)
-
-                if claimed_partition_ids:
+                if claimed_ownership_list:
+                    claimed_partition_ids = [x["partition_id"] for x in claimed_ownership_list]
+                    to_cancel_list = self._tasks.keys() - claimed_partition_ids
                     self._create_tasks_for_claimed_ownership(claimed_ownership_list)
                 else:
                     log.warning("EventProcessor %r hasn't claimed an ownership. It keeps claiming.", self._id)
+                    to_cancel_list = self._tasks.keys()
+                if to_cancel_list:
+                    self._cancel_tasks_for_partitions(to_cancel_list)
+                    log.info("EventProcesor %r has cancelled partitions %r", self._id, to_cancel_list)
                 await asyncio.sleep(self._polling_interval)
 
     async def stop(self):
-        """Stop all the partition consumer
+        """Stop claiming ownership and all the partition consumers owned by this EventProcessor
 
-        This method cancels tasks that are running EventHubConsumer.receive() for the partitions owned by this EventProcessor.
+        This method stops claiming ownership of owned partitions and cancels tasks that are running
+        EventHubConsumer.receive() for the partitions owned by this EventProcessor.
 
         :return: None
 
@@ -152,6 +157,13 @@ def _create_tasks_for_claimed_ownership(self, to_claim_ownership_list):
 
     async def _receive(self, ownership):
         log.info("start ownership, %r", ownership)
+        partition_processor = self._partition_processor_factory()
+        if not hasattr(partition_processor, "process_events"):
+            log.error(
+                "Fatal error: a partition processor should at least have method process_events(events, checkpoint_manager). EventProcessor will stop.")
+            await self.stop()
+            raise TypeError("Partition processor must has method process_events(events, checkpoint_manager")
+
         partition_consumer = self._eventhub_client.create_consumer(ownership["consumer_group_name"],
                                                                    ownership["partition_id"],
                                                                    EventPosition(ownership.get("offset", self._initial_event_position))
@@ -161,8 +173,6 @@ async def _receive(self, ownership):
                                                ownership["consumer_group_name"],
                                                ownership["owner_id"],
                                                self._partition_manager)
-        partition_processor = self._partition_processor_factory()
-
         async def initialize():
             if hasattr(partition_processor, "initialize"):
                 await partition_processor.initialize(checkpoint_manager)
@@ -192,6 +202,7 @@ async def close(close_reason):
                     )
                     await process_error(cancelled_error)
                     await close(CloseReason.SHUTDOWN)
+                    # TODO: release the ownership immediately via partition manager
                     break
                 except EventHubError as eh_err:
                     reason = CloseReason.LEASE_LOST if eh_err.error == "link:stolen" else CloseReason.EVENTHUB_EXCEPTION
@@ -205,11 +216,10 @@ async def close(close_reason):
                         eh_err
                     )
                     await process_error(eh_err)
-                    await close(reason)
+                    await close(reason)  # An EventProcessor will pick up this partition again after the ownership is released
+                    # TODO: release the ownership immediately via partition manager
                     break
                 except Exception as exp:
                     log.warning(exp)
-                    # TODO: will review whether to break and close partition processor after user's code has an exception
-            # TODO: try to inform other EventProcessors to take the partition when this partition is closed in preview 3?
         finally:
             await partition_consumer.close()

sdk/eventhub/azure-eventhubs/azure/eventhub/eventprocessor/partition_processor.py

@@ -23,7 +23,7 @@ class PartitionProcessor(Protocol):
     implementing this abstract class will be created for every partition the associated ~azure.eventhub.eventprocessor.EventProcessor owns.
 
     """
-    async def initialize(self, checkpoint_manager: CheckpointManager):
+    async def initialize(self):
         pass
 
     async def close(self, reason, checkpoint_manager: CheckpointManager):
@@ -45,7 +45,7 @@ async def process_events(self, events: List[EventData], checkpoint_manager: Chec
         :type events: list[~azure.eventhub.common.EventData]
 
         """
-        pass
+        raise NotImplementedError
 
     async def process_error(self, error, checkpoint_manager: CheckpointManager):
         """Called when an error happens

sdk/eventhub/azure-eventhubs/azure/eventhub/eventprocessor/sqlite3_partition_manager.py

@@ -121,9 +121,16 @@ async def update_checkpoint(self, eventhub_name, consumer_group_name, partition_
             offset, sequence_number):
         cursor = self.conn.cursor()
         try:
-            cursor.execute("update " + _check_table_name(self.ownership_table) + " set offset=?, sequence_number=? where eventhub_name=? and consumer_group_name=? and partition_id=?",
-                           (offset, sequence_number, eventhub_name, consumer_group_name, partition_id))
-            self.conn.commit()
+            cursor.execute("select owner_id from " + _check_table_name(self.ownership_table) + " where eventhub_name=? and consumer_group_name=? and partition_id=?",
+                           (eventhub_name, consumer_group_name, partition_id))
+            cursor_fetch = cursor.fetchall()
+            if cursor_fetch and owner_id == cursor_fetch[0][0]:
+                cursor.execute("update " + _check_table_name(self.ownership_table) + " set offset=?, sequence_number=? where eventhub_name=? and consumer_group_name=? and partition_id=?",
+                               (offset, sequence_number, eventhub_name, consumer_group_name, partition_id))
+                self.conn.commit()
+            else:
+                logger.info("EventProcessor couldn't checkpoint to partition %r because it no longer has the ownership", partition_id)
+
         finally:
             cursor.close()
 

sdk/eventhub/azure-eventhubs/examples/eventprocessor/event_processor_example.py

@@ -3,7 +3,6 @@
 import os
 from azure.eventhub.aio import EventHubClient
 from azure.eventhub.eventprocessor import EventProcessor
-from azure.eventhub.eventprocessor import PartitionProcessor
 from azure.eventhub.eventprocessor import Sqlite3PartitionManager
 
 RECEIVE_TIMEOUT = 5  # timeout in seconds for a receiving operation. 0 or None means no timeout
@@ -18,33 +17,23 @@ async def do_operation(event):
     print(event)
 
 
-class MyPartitionProcessor(PartitionProcessor):
-    def __init__(self, checkpoint_manager):
-        super(MyPartitionProcessor, self).__init__(checkpoint_manager)
-
-    async def process_events(self, events):
+class MyPartitionProcessor(object):
+    async def process_events(self, events, checkpoint_manager):
         if events:
             await asyncio.gather(*[do_operation(event) for event in events])
-            await self._checkpoint_manager.update_checkpoint(events[-1].offset, events[-1].sequence_number)
-
-
-def partition_processor_factory(checkpoint_manager):
-    return MyPartitionProcessor(checkpoint_manager)
-
-
-async def run_awhile(duration):
-    client = EventHubClient.from_connection_string(CONNECTION_STR, receive_timeout=RECEIVE_TIMEOUT,
-                                                   retry_total=RETRY_TOTAL)
-    partition_manager = Sqlite3PartitionManager()
-    event_processor = EventProcessor(client, "$default", MyPartitionProcessor, partition_manager)
-    try:
-        asyncio.ensure_future(event_processor.start())
-        await asyncio.sleep(duration)
-        await event_processor.stop()
-    finally:
-        await partition_manager.close()
+            await checkpoint_manager.update_checkpoint(events[-1].offset, events[-1].sequence_number)
+        else:
+            print("empty events received", "partition:", checkpoint_manager.partition_id)
 
 
 if __name__ == '__main__':
     loop = asyncio.get_event_loop()
-    loop.run_until_complete(run_awhile(60))
+    client = EventHubClient.from_connection_string(CONNECTION_STR, receive_timeout=RECEIVE_TIMEOUT, retry_total=RETRY_TOTAL)
+    partition_manager = Sqlite3PartitionManager(db_filename="eventprocessor_test_db")
+    event_processor = EventProcessor(client, "$default", MyPartitionProcessor, partition_manager, polling_interval=1)
+    try:
+        loop.run_until_complete(event_processor.start())
+    except KeyboardInterrupt:
+        loop.run_until_complete(event_processor.stop())
+    finally:
+        loop.stop()