[Event Hub] Improve flush functionality of buffered partition producer. (Azure#44904)

* Improve flush functionality of buffered parition producer.

* Remove EventDataBatchCarrier and use EventDataBatch.EMPTY instead

* address review feedback

Author: anuchandy

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

@@ -8,11 +8,15 @@
 
 ### Bugs Fixed
 
-- Fixes issue where `EventHubBufferedProducerClient` and `EventHubBufferedProducerAsyncClient` are unable to enqueue events when `SendOptions.getPartitionId()` is set. ([#44392](https://github.com/Azure/azure-sdk-for-java/pull/44392))
-- Fixes issue where `EventHubBufferedProducerClient` and `EventHubBufferedProducerAsyncClient` returns 0 after enqueueing events or calling `getBufferedEventCount()`. ([#44392](https://github.com/Azure/azure-sdk-for-java/pull/44392))
+- Fixes issue where `EventHubBufferedProducerClient` and `EventHubBufferedProducerAsyncClient` are unable to `flush`. ([#44904](https://github.com/Azure/azure-sdk-for-java/pull/44904))
 
 ## 5.20.2 (2025-03-24)
 
+### Bugs Fixed
+
+- Fixes issue where `EventHubBufferedProducerClient` and `EventHubBufferedProducerAsyncClient` are unable to enqueue events when `SendOptions.getPartitionId()` is set. ([#44392](https://github.com/Azure/azure-sdk-for-java/pull/44392))
+- Fixes issue where `EventHubBufferedProducerClient` and `EventHubBufferedProducerAsyncClient` returns 0 after enqueueing events or calling `getBufferedEventCount()`. ([#44392](https://github.com/Azure/azure-sdk-for-java/pull/44392))
+
 ### Other Changes
 
 #### Dependency Updates

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

@@ -239,10 +239,12 @@ public void onComplete() {
          * @param alwaysPublish {@code true} to always push batch downstream. {@code false}, otherwise.
          */
         private void updateOrPublishBatch(EventData eventData, boolean alwaysPublish) {
-            if (alwaysPublish) {
-                publishDownstream();
+            final boolean isFlush = isFlushSignal(eventData);
+            if (alwaysPublish || isFlush) {
+                publishDownstream(isFlush);
                 return;
-            } else if (eventData == null) {
+            }
+            if (eventData == null) {
                 // EventData will be null in the case when options.maxWaitTime() has elapsed  and we want to push the
                 // batch downstream.
                 return;
@@ -256,7 +258,7 @@ private void updateOrPublishBatch(EventData eventData, boolean alwaysPublish) {
                     return;
                 }
 
-                publishDownstream();
+                publishDownstream(false);
                 added = currentBatch.tryAdd(eventData);
             }
 
@@ -271,19 +273,24 @@ private void updateOrPublishBatch(EventData eventData, boolean alwaysPublish) {
         /**
          * Publishes batch downstream if there are events in the batch and updates it.
          */
-        private void publishDownstream() {
+        private void publishDownstream(boolean isFlush) {
             EventDataBatch previous = null;
-
             try {
                 synchronized (lock) {
                     previous = this.currentBatch;
-
                     if (previous == null) {
                         logger.warning("Batch should not be null, setting a new batch.");
-
                         this.currentBatch = batchSupplier.get();
+                        if (isFlush) {
+                            downstream.onNext(EventDataBatch.EMPTY);
+                        }
                         return;
                     } else if (previous.getEvents().isEmpty()) {
+                        if (isFlush) {
+                            // Even if the batch is empty, we'll push EMPTY on a flush signal.
+                            // This ensures any flush related flags set at the call site are reset.
+                            downstream.onNext(EventDataBatch.EMPTY);
+                        }
                         return;
                     }
 
@@ -332,5 +339,9 @@ private void publishDownstream() {
                 }
             }
         }
+
+        private static boolean isFlushSignal(EventData eventData) {
+            return eventData instanceof FlushSignal;
+        }
     }
 }

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

@@ -35,6 +35,8 @@
  *     producer.
  */
 public final class EventDataBatch {
+    static final EventDataBatch EMPTY
+        = new EventDataBatch(0, null, null, null, EventHubsProducerInstrumentation.NOOP_INSTANCE);
     private static final ClientLogger LOGGER = new ClientLogger(EventDataBatch.class);
     private final int maxMessageSize;
     private final String partitionKey;