[EH] Keep alive fixes on producer (#33788)

* start adding through keep alive

* keep alive changes

* add tests

* consumer async

* update tests

* update test

* adding more tests

* add idle timeout test

* tests

* update test to use reconnect format

* async transport

Author: l0lawrence

sdk/eventhub/azure-eventhub/azure/eventhub/_client_base.py

@@ -311,7 +311,6 @@ def __init__(
             self._credential = EventhubAzureNamedKeyTokenCredential(credential)  # type: ignore
         else:
             self._credential = credential  # type: ignore
-        self._keep_alive = kwargs.get("keep_alive", 30)
         self._auto_reconnect = kwargs.get("auto_reconnect", True)
         self._auth_uri = f"sb://{self._address.hostname}{self._address.path}"
         self._config = Configuration(

sdk/eventhub/azure-eventhub/azure/eventhub/_consumer.py

@@ -83,7 +83,7 @@ def __init__(
         event_position = kwargs.get("event_position", None)
         prefetch = kwargs.get("prefetch", 300)
         owner_level = kwargs.get("owner_level", None)
-        keep_alive = kwargs.get("keep_alive", None)
+        keep_alive = kwargs.get("keep_alive", 30)
         auto_reconnect = kwargs.get("auto_reconnect", True)
         track_last_enqueued_event_properties = kwargs.get(
             "track_last_enqueued_event_properties", False

sdk/eventhub/azure-eventhub/azure/eventhub/_producer_client.py

@@ -197,6 +197,9 @@ def __init__(
             network_tracing=kwargs.get("logging_enable"),
             **kwargs
         )
+
+        self._keep_alive = kwargs.get("keep_alive", None)
+
         self._producers: Dict[str, Optional[EventHubProducer]] = {
             ALL_PARTITIONS: self._create_producer()
         }
@@ -380,6 +383,7 @@ def _create_producer(
             send_timeout=send_timeout,
             idle_timeout=self._idle_timeout,
             amqp_transport=self._amqp_transport,
+            keep_alive=self._keep_alive,
         )
         return handler
 

sdk/eventhub/azure-eventhub/azure/eventhub/aio/_consumer_async.py

@@ -82,7 +82,7 @@ def __init__(self, client: "EventHubConsumerClient", source: str, **kwargs: Any)
         event_position = kwargs.get("event_position", None)
         prefetch = kwargs.get("prefetch", 300)
         owner_level = kwargs.get("owner_level", None)
-        keep_alive = kwargs.get("keep_alive", None)
+        keep_alive = kwargs.get("keep_alive", 30)
         auto_reconnect = kwargs.get("auto_reconnect", True)
         track_last_enqueued_event_properties = kwargs.get(
             "track_last_enqueued_event_properties", False

sdk/eventhub/azure-eventhub/azure/eventhub/aio/_producer_client_async.py

@@ -184,6 +184,7 @@ def __init__(
             network_tracing=kwargs.pop("logging_enable", False),
             **kwargs
         )
+        self._keep_alive = kwargs.get("keep_alive", None)
         self._producers: Dict[str, Optional[EventHubProducer]] = {
             ALL_PARTITIONS: self._create_producer()
         }
@@ -365,6 +366,7 @@ def _create_producer(
             send_timeout=send_timeout,
             idle_timeout=self._idle_timeout,
             amqp_transport = self._amqp_transport,
+            keep_alive = self._keep_alive,
             **self._internal_kwargs
         )
         return handler

sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_send_async.py

@@ -34,6 +34,16 @@
 from azure.eventhub._pyamqp.authentication import SASTokenAuth
 from azure.eventhub._pyamqp.client import ReceiveClient
 from azure.eventhub._pyamqp.error import AMQPConnectionError
+from azure.eventhub._utils import transform_outbound_single_message
+try:
+    import uamqp
+    from uamqp import compat
+    from azure.eventhub.aio._transport._uamqp_transport_async import UamqpTransportAsync as UamqpTransport
+except (ModuleNotFoundError, ImportError):
+    uamqp = None
+    UamqpTransport = None
+
+from azure.eventhub.aio._transport._pyamqp_transport_async import PyamqpTransportAsync as PyamqpTransport
 
 
 @pytest.mark.liveTest
@@ -529,3 +539,49 @@ async def on_success(events, pid):
         assert sent_events[-1][1] == "0"
 
         assert not on_error.err
+
+@pytest.mark.parametrize("keep_alive", [None, 30, 60])
+@pytest.mark.liveTest
+@pytest.mark.asyncio
+def test_send_with_keep_alive_async(connstr_receivers, keep_alive, uamqp_transport):
+    connection_str, receivers = connstr_receivers
+    client = EventHubProducerClient.from_connection_string(connection_str, keep_alive=keep_alive, uamqp_transport=uamqp_transport)
+    assert client._producers["all-partitions"]._keep_alive == keep_alive
+
+
+@pytest.mark.parametrize("keep_alive", [None, 5, 30])
+@pytest.mark.liveTest
+@pytest.mark.asyncio
+async def test_send_long_wait_idle_timeout(connstr_receivers, keep_alive, uamqp_transport):
+    if uamqp_transport:
+        amqp_transport = UamqpTransport
+        retry_total = 3
+    else:
+        amqp_transport = PyamqpTransport
+        retry_total = 0
+    connection_str, receivers = connstr_receivers
+    client = EventHubProducerClient.from_connection_string(connection_str, keep_alive=keep_alive, idle_timeout=10, retry_total=retry_total, uamqp_transport=uamqp_transport)
+    sender = await client._create_producer(partition_id="0")
+    async with sender:
+        await sender._open_with_retry()
+        ed = EventData('data')
+        ed = transform_outbound_single_message(ed, EventData, amqp_transport.to_outgoing_amqp_message)
+        sender._unsent_events = [ed._message]
+        # hit idle timeout error
+        await asyncio.sleep(11)
+
+        if uamqp_transport:
+            sender._unsent_events[0].on_send_complete = sender._on_outcome
+            if keep_alive !=5:
+                with pytest.raises((uamqp.errors.ConnectionClose,
+                                        uamqp.errors.MessageHandlerError, OperationTimeoutError)):
+                    await sender._send_event_data()
+            else:
+                await sender._send_event_data()
+
+        if not uamqp_transport:
+            if keep_alive == 5:
+                await sender._send_event_data()
+            else:
+                with pytest.raises(AMQPConnectionError):
+                    await sender._send_event_data()

sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_send.py

@@ -34,6 +34,16 @@
 from azure.eventhub._pyamqp.authentication import SASTokenAuth
 from azure.eventhub._pyamqp import ReceiveClient
 from azure.eventhub._pyamqp.error import AMQPConnectionError
+from azure.eventhub._utils import transform_outbound_single_message
+try:
+    import uamqp
+    from uamqp import compat
+    from azure.eventhub._transport._uamqp_transport import UamqpTransport
+except (ModuleNotFoundError, ImportError):
+    uamqp = None
+    UamqpTransport = None
+
+from azure.eventhub._transport._pyamqp_transport import PyamqpTransport
 
 
 @pytest.mark.liveTest
@@ -557,3 +567,46 @@ def on_success(events, pid):
         assert sent_events[-1][1] == "0"
 
         assert not on_error.err
+
+@pytest.mark.parametrize("keep_alive", [None, 30, 60])
+@pytest.mark.liveTest
+def test_send_with_keep_alive(connstr_receivers, keep_alive, uamqp_transport):
+    connection_str, receivers = connstr_receivers
+    client = EventHubProducerClient.from_connection_string(connection_str, keep_alive=keep_alive, uamqp_transport=uamqp_transport)
+    assert client._producers["all-partitions"]._keep_alive == keep_alive
+
+@pytest.mark.parametrize("keep_alive", [None, 5, 30])
+@pytest.mark.liveTest
+def test_send_long_wait_idle_timeout(connstr_receivers, keep_alive, uamqp_transport):
+    if uamqp_transport:
+        amqp_transport = UamqpTransport
+        retry_total = 3
+    else:
+        amqp_transport = PyamqpTransport
+        retry_total = 0
+    connection_str, receivers = connstr_receivers
+    client = EventHubProducerClient.from_connection_string(connection_str, keep_alive=keep_alive, idle_timeout=10, retry_total=retry_total, uamqp_transport=uamqp_transport)
+    sender = client._create_producer(partition_id="0")
+    with sender:
+        sender._open_with_retry()
+        ed = EventData('data')
+        ed = transform_outbound_single_message(ed, EventData, amqp_transport.to_outgoing_amqp_message)
+        sender._unsent_events = [ed._message]
+        # hit idle timeout error
+        time.sleep(11)
+
+        if uamqp_transport:
+            sender._unsent_events[0].on_send_complete = sender._on_outcome
+            if keep_alive !=5:
+                with pytest.raises((uamqp.errors.ConnectionClose,
+                                        uamqp.errors.MessageHandlerError, OperationTimeoutError)):
+                    sender._send_event_data()
+            else:
+                sender._send_event_data()
+
+        if not uamqp_transport:
+            if keep_alive == 5:
+                sender._send_event_data()
+            else:
+                with pytest.raises(AMQPConnectionError):
+                    sender._send_event_data()