Fix code review issues

sdk/eventhub/azure-eventhubs/azure/eventhub/aio/eventprocessor/__init__.py

@@ -7,6 +7,7 @@
 from .partition_processor import PartitionProcessor, CloseReason
 from .partition_manager import PartitionManager, OwnershipLostError
 from .partition_context import PartitionContext
+from .sample_partition_manager import SamplePartitionManager
 
 __all__ = [
     'CloseReason',
@@ -15,4 +16,5 @@
     'PartitionManager',
     'OwnershipLostError',
     'PartitionContext',
+    'SamplePartitionManager',
 ]

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

@@ -24,7 +24,7 @@ class OwnershipManager(object):
     """
     def __init__(
             self, eventhub_client: EventHubClient, consumer_group_name: str, owner_id: str,
-            partition_manager: PartitionManager, ownership_timeout: int
+            partition_manager: PartitionManager, ownership_timeout: float
     ):
         self.cached_parition_ids = []  # type: List[str]
         self.eventhub_client = eventhub_client
@@ -86,7 +86,7 @@ async def _balance_ownership(self, all_partition_ids):
             self.eventhub_name, self.consumer_group_name
         )
         now = time.time()
-        ownership_dict = dict((x["partition_id"], x) for x in ownership_list)  # put the list to dict for fast lookup
+        ownership_dict = {x["partition_id"]: x for x in ownership_list}  # put the list to dict for fast lookup
         not_owned_partition_ids = [pid for pid in all_partition_ids if pid not in ownership_dict]
         timed_out_partition_ids = [ownership["partition_id"] for ownership in ownership_list
                                    if ownership["last_modified_time"] + self.ownership_timeout < now]
@@ -100,6 +100,8 @@ async def _balance_ownership(self, all_partition_ids):
 
         # calculate expected count per owner
         all_partition_count = len(all_partition_ids)
+        # owners_count is the number of active owners. If self.owner_id is not yet among the active owners,
+        # then plus 1 to include self. This will make owners_count >= 1.
         owners_count = len(active_ownership_by_owner) + \
                        (0 if self.owner_id in active_ownership_by_owner else 1)
         expected_count_per_owner = all_partition_count // owners_count
@@ -109,7 +111,6 @@ async def _balance_ownership(self, all_partition_ids):
         to_claim = active_ownership_self
         if len(active_ownership_self) > most_count_allowed_per_owner:  # needs to abandon a partition
             to_claim.pop()  # abandon one partition if owned too many
-            # TODO: Release an ownership immediately so other EventProcessors won't need to wait it to timeout
         elif len(active_ownership_self) < expected_count_per_owner:
             # Either claims an inactive partition, or steals from other owners
             if claimable_partition_ids:  # claim an inactive partition if there is

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

@@ -56,7 +56,9 @@ async def process_events(self, events, partition_context):
     def __init__(
             self, eventhub_client: EventHubClient, consumer_group_name: str,
             partition_processor_factory: Callable[..., PartitionProcessor],
-            partition_manager: PartitionManager, **kwargs
+            partition_manager: PartitionManager, *,
+            initial_event_position: EventPosition = EventPosition("-1"), polling_interval: float = 10.0
+
     ):
         """
         Instantiate an EventProcessor.
@@ -84,14 +86,9 @@ def __init__(
         self._eventhub_name = eventhub_client.eh_name
         self._partition_processor_factory = partition_processor_factory
         self._partition_manager = partition_manager
-        self._initial_event_position = kwargs.get("initial_event_position", "-1")
-        # TODO: initial position provider will be a callable
-        #  so users can create initial event position for every partition
-        self._max_batch_size = eventhub_client.config.max_batch_size
-        self._receive_timeout = eventhub_client.config.receive_timeout
-        self._polling_interval = kwargs.get("polling_interval", 10)
+        self._initial_event_position = initial_event_position  # will be replaced by reset event position in preview 4
+        self._polling_interval = polling_interval
         self._ownership_timeout = self._polling_interval * 2
-        # TODO: Team haven't decided if this is a separate argument
         self._tasks = {}  # type: Dict[str, asyncio.Task]
         self._id = str(uuid.uuid4())
         self._running = False
@@ -122,18 +119,19 @@ async def start(self):
                 try:
                     claimed_ownership_list = await ownership_manager.claim_ownership()
                 except Exception as err:
-                    log.exception("An exception occurred during balancing and claiming ownership for eventhub %r "
-                                  "consumer group %r. Retrying after %r seconds",
-                                  self._eventhub_name, self._consumer_group_name, self._polling_interval, exc_info=err)
+                    log.warning("An exception (%r) occurred during balancing and claiming ownership for eventhub %r "
+                                "consumer group %r. Retrying after %r seconds",
+                                err, self._eventhub_name, self._consumer_group_name, self._polling_interval)
                     await asyncio.sleep(self._polling_interval)
                     continue
+
+                to_cancel_list = self._tasks.keys()
                 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()
+                    log.info("EventProcessor %r hasn't claimed an ownership. It keeps claiming.", self._id)
                 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)
@@ -153,7 +151,7 @@ async def stop(self):
             _, task = self._tasks.popitem()
             task.cancel()
         log.info("EventProcessor %r has been cancelled", self._id)
-        await asyncio.sleep(2)  # give some time to finish after cancelled
+        await asyncio.sleep(2)  # give some time to finish after cancelled.
 
     def _cancel_tasks_for_partitions(self, to_cancel_partitions):
         for partition_id in to_cancel_partitions:
@@ -181,23 +179,22 @@ async def _receive(self, ownership):
             owner_id,
             self._partition_manager
         )
-        partition_processor.eventhub_name = ownership
         partition_consumer = self._eventhub_client.create_consumer(
             consumer_group_name,
             partition_id,
-            EventPosition(ownership.get("offset", self._initial_event_position))
+            EventPosition(ownership.get("offset", self._initial_event_position.value))
         )
 
         async def process_error(err):
-            log.error(
+            log.warning(
                 "PartitionProcessor of EventProcessor instance %r of eventhub %r partition %r consumer group %r"
                 " has met an error. The exception is %r.",
                 owner_id, eventhub_name, partition_id, consumer_group_name, err
             )
             try:
                 await partition_processor.process_error(err, partition_context)
             except Exception as err_again:  # pylint:disable=broad-except
-                log.error(
+                log.warning(
                     "PartitionProcessor of EventProcessor instance %r of eventhub %r partition %r consumer group %r"
                     " has another error during running process_error(). The exception is %r.",
                     owner_id, eventhub_name, partition_id, consumer_group_name, err_again
@@ -212,7 +209,7 @@ async def close(reason):
             try:
                 await partition_processor.close(reason, partition_context)
             except Exception as err:  # pylint:disable=broad-except
-                log.error(
+                log.warning(
                     "PartitionProcessor of EventProcessor instance %r of eventhub %r partition %r consumer group %r"
                     " has an error during running close(). The exception is %r.",
                     owner_id, eventhub_name, partition_id, consumer_group_name, err
@@ -222,7 +219,7 @@ async def close(reason):
             while True:
                 try:
                     await partition_processor.initialize(partition_context)
-                    events = await partition_consumer.receive(timeout=self._receive_timeout)
+                    events = await partition_consumer.receive()
                     await partition_processor.process_events(events, partition_context)
                 except asyncio.CancelledError:
                     log.info(
@@ -237,21 +234,18 @@ async def close(reason):
                         await close(CloseReason.SHUTDOWN)
                     else:
                         await close(CloseReason.OWNERSHIP_LOST)
-                    # TODO: release the ownership immediately via partition manager in preview 4
                     break
                 except EventHubError as eh_err:
                     await process_error(eh_err)
                     await close(CloseReason.EVENTHUB_EXCEPTION)
                     # An EventProcessor will pick up this partition again after the ownership is released
-                    # TODO: release the ownership immediately via partition manager in preview 4
                     break
                 except OwnershipLostError:
                     await close(CloseReason.OWNERSHIP_LOST)
                     break
                 except Exception as other_error:  # pylint:disable=broad-except
                     await process_error(other_error)
                     await close(CloseReason.PROCESS_EVENTS_ERROR)
-                    # TODO: release the ownership immediately via partition manager in preview 4
                     break
         finally:
             await partition_consumer.close()

sdk/eventhub/azure-eventhubs/azure/eventhub/aio/eventprocessor/partition_manager.py

@@ -35,14 +35,13 @@ async def list_ownership(self, eventhub_name: str, consumer_group_name: str) ->
                 etag
         """
 
-
     @abstractmethod
-    async def claim_ownership(self, partitions: Iterable[Dict[str, Any]]) -> Iterable[Dict[str, Any]]:
+    async def claim_ownership(self, ownership_list: Iterable[Dict[str, Any]]) -> Iterable[Dict[str, Any]]:
         """
         Tries to claim a list of specified ownership.
 
-        :param partitions: Iterable of dictionaries containing all the ownership to claim.
-        :type partitions: Iterable of dict
+        :param ownership_list: Iterable of dictionaries containing all the ownership to claim.
+        :type ownership_list: Iterable of dict
         :return: Iterable of dictionaries containing the following partition ownership information:
                 eventhub_name
                 consumer_group_name
@@ -55,7 +54,6 @@ async def claim_ownership(self, partitions: Iterable[Dict[str, Any]]) -> Iterabl
                 etag
         """
 
-
     @abstractmethod
     async def update_checkpoint(self, eventhub_name, consumer_group_name, partition_id, owner_id,
                                 offset, sequence_number) -> None:

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

@@ -3,7 +3,7 @@
 import os
 from azure.eventhub.aio import EventHubClient
 from azure.eventhub.aio.eventprocessor import EventProcessor, PartitionProcessor
-from azure.eventhub.aio.eventprocessor.partitionmanagers import SamplePartitionManager
+from azure.eventhub.aio.eventprocessor import SamplePartitionManager
 
 RECEIVE_TIMEOUT = 5  # timeout in seconds for a receiving operation. 0 or None means no timeout
 RETRY_TOTAL = 3  # max number of retries for receive operations within the receive timeout. Actual number of retries clould be less if RECEIVE_TIMEOUT is too small