Report prefetched seqno instead of consumer lag (Azure#30731)

* Report prefetched seqno instead of consumer lag

Author: Liudmila Molkova

sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/implementation/AmqpMetricsProvider.java

@@ -11,20 +11,18 @@
 import com.azure.core.util.logging.ClientLogger;
 import com.azure.core.util.metrics.DoubleHistogram;
 import com.azure.core.util.metrics.LongCounter;
+import com.azure.core.util.metrics.LongGauge;
 import com.azure.core.util.metrics.Meter;
 import com.azure.core.util.metrics.MeterProvider;
 import org.apache.qpid.proton.amqp.Symbol;
 import org.apache.qpid.proton.amqp.transport.DeliveryState;
 import org.apache.qpid.proton.amqp.transport.ErrorCondition;
-import org.apache.qpid.proton.message.Message;
 
 import java.time.Instant;
-import java.util.Date;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
-
-import static com.azure.core.amqp.AmqpMessageConstant.ENQUEUED_TIME_UTC_ANNOTATION_NAME;
+import java.util.function.Supplier;
 
 /**
  * Helper class responsible for efficient reporting metrics in AMQP core. It's efficient and safe to use when there is no
@@ -34,7 +32,6 @@ public class AmqpMetricsProvider {
     public static final String STATUS_CODE_KEY = "amqpStatusCode";
     public static final String MANAGEMENT_OPERATION_KEY = "amqpOperation";
     private static final ClientLogger LOGGER = new ClientLogger(AmqpMetricsProvider.class);
-    private static final Symbol ENQUEUED_TIME_ANNOTATION = Symbol.valueOf(ENQUEUED_TIME_UTC_ANNOTATION_NAME.getValue());
 
     private static final String AZURE_CORE_AMQP_PROPERTIES_NAME = "azure-core.properties";
     private static final String AZURE_CORE_AMQP_PROPERTIES_VERSION_KEY = "version";
@@ -43,6 +40,9 @@ public class AmqpMetricsProvider {
         .getProperties(AZURE_CORE_AMQP_PROPERTIES_NAME)
         .getOrDefault(AZURE_CORE_AMQP_PROPERTIES_VERSION_KEY, null);
 
+    private static final AutoCloseable NOOP_CLOSEABLE = () -> {
+    };
+
     // all delivery state + 1 for `null` - we'll treat it as an error - no delivery was received
     private static final int DELIVERY_STATES_COUNT = DeliveryState.DeliveryStateType.values().length + 1;
 
@@ -59,7 +59,7 @@ public class AmqpMetricsProvider {
     private LongCounter sessionErrors = null;
     private LongCounter linkErrors = null;
     private LongCounter transportErrors = null;
-    private DoubleHistogram receivedLag = null;
+    private LongGauge prefetchedSequenceNumber = null;
     private LongCounter addCredits = null;
 
     /**
@@ -131,7 +131,7 @@ public AmqpMetricsProvider(Meter meter, String namespace, String entityPath) {
             this.linkErrors = this.meter.createLongCounter("messaging.az.amqp.client.link.errors", "AMQP link errors", "errors");
             this.transportErrors = this.meter.createLongCounter("messaging.az.amqp.client.transport.errors", "AMQP session errors", "errors");
             this.addCredits = this.meter.createLongCounter("messaging.az.amqp.consumer.credits.requested", "Number of requested credits", "credits");
-            this.receivedLag = this.meter.createDoubleHistogram("messaging.az.amqp.consumer.lag", "Approximate lag between time message was received and time it was enqueued on the broker.", "sec");
+            this.prefetchedSequenceNumber = this.meter.createLongGauge("messaging.az.amqp.prefetch.sequence_number", "Last prefetched sequence number", "seqNo");
         }
     }
 
@@ -140,12 +140,20 @@ public static AmqpMetricsProvider noop() {
     }
 
     /**
-     * Checks if record delivers is enabled (for micro-optimizations).
+     * Checks if send delivery metric is enabled (for micro-optimizations).
      */
     public boolean isSendDeliveryEnabled() {
         return isEnabled && sendDuration.isEnabled();
     }
 
+    /**
+     * Checks if prefetched sequence number is enabled (for micro-optimizations).
+     */
+    public boolean isPrefetchedSequenceNumberEnabled() {
+        return isEnabled && prefetchedSequenceNumber.isEnabled();
+    }
+
+
     /**
      * Records duration of AMQP send call.
      */
@@ -178,28 +186,14 @@ public void recordConnectionClosed(ErrorCondition condition) {
     }
 
     /**
-     * Records the message was received.
+     * Creates gauge subscription to report latest sequence number value.
      */
-    public void recordReceivedMessage(Message message) {
-        if (!isEnabled || !receivedLag.isEnabled()
-            || message == null
-            || message.getMessageAnnotations() == null
-            || message.getBody() == null) {
-            return;
+    public AutoCloseable trackPrefetchSequenceNumber(Supplier<Long> valueSupplier) {
+        if (!isEnabled || !prefetchedSequenceNumber.isEnabled()) {
+            return NOOP_CLOSEABLE;
         }
 
-        Map<Symbol, Object> properties = message.getMessageAnnotations().getValue();
-        Object enqueuedTimeDate = properties != null ? properties.get(ENQUEUED_TIME_ANNOTATION) : null;
-        if (enqueuedTimeDate instanceof Date) {
-            Instant enqueuedTime = ((Date) enqueuedTimeDate).toInstant();
-            long deltaMs = Instant.now().toEpochMilli() - enqueuedTime.toEpochMilli();
-            if (deltaMs < 0) {
-                deltaMs = 0;
-            }
-            receivedLag.record(deltaMs / 1000d, commonAttributes, Context.NONE);
-        } else if (enqueuedTimeDate != null) {
-            LOGGER.verbose("Received message has unexpected `x-opt-enqueued-time` annotation value - `{}`. Ignoring it.", enqueuedTimeDate);
-        }
+        return prefetchedSequenceNumber.registerCallback(valueSupplier, commonAttributes);
     }
 
     /**

sdk/core/azure-core-amqp/src/main/java/com/azure/core/amqp/implementation/ReactorReceiver.java

@@ -31,9 +31,11 @@
 import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
 
+import static com.azure.core.amqp.AmqpMessageConstant.SEQUENCE_NUMBER_ANNOTATION_NAME;
 import static com.azure.core.amqp.implementation.AmqpLoggingUtils.addErrorCondition;
 import static com.azure.core.amqp.implementation.AmqpLoggingUtils.addSignalTypeAndResult;
 import static com.azure.core.amqp.implementation.AmqpLoggingUtils.createContextWithConnectionId;
@@ -45,6 +47,7 @@
  * Handles receiving events from Event Hubs service and translating them to proton-j messages.
  */
 public class ReactorReceiver implements AmqpReceiveLink, AsyncCloseable, AutoCloseable {
+    private static final Symbol SEQUENCE_NUMBER_ANNOTATION = Symbol.valueOf(SEQUENCE_NUMBER_ANNOTATION_NAME.getValue());
     private final String entityPath;
     private final Receiver receiver;
     private final ReceiveLinkHandler handler;
@@ -65,6 +68,8 @@ public class ReactorReceiver implements AmqpReceiveLink, AsyncCloseable, AutoClo
 
     private final AtomicReference<Supplier<Integer>> creditSupplier = new AtomicReference<>();
     private final AmqpMetricsProvider metricsProvider;
+    private final AtomicLong lastSequenceNumber = new AtomicLong();
+    private final AutoCloseable trackPrefetchSeqNoSubscription;
 
     @Deprecated
     protected ReactorReceiver(AmqpConnection amqpConnection, String entityPath, Receiver receiver,
@@ -83,6 +88,7 @@ protected ReactorReceiver(AmqpConnection amqpConnection, String entityPath, Rece
         this.tokenManager = tokenManager;
         this.dispatcher = dispatcher;
         this.metricsProvider = metricsProvider;
+        this.trackPrefetchSeqNoSubscription = this.metricsProvider.trackPrefetchSequenceNumber(lastSequenceNumber::get);
 
         Map<String, Object> loggingContext = createContextWithConnectionId(handler.getConnectionId());
         loggingContext.put(LINK_NAME_KEY, this.handler.getLinkName());
@@ -104,8 +110,14 @@ protected ReactorReceiver(AmqpConnection amqpConnection, String entityPath, Rece
                                     "Cannot decode delivery when ReactorReceiver instance is closed."));
                                 return;
                             }
+
                             final Message message = decodeDelivery(delivery);
-                            metricsProvider.recordReceivedMessage(message);
+                            if (metricsProvider.isPrefetchedSequenceNumberEnabled()) {
+                                Long seqNo = getSequenceNumber(message);
+                                if (seqNo != null) {
+                                    lastSequenceNumber.set(seqNo);
+                                }
+                            }
 
                             final int creditsLeft = receiver.getRemoteCredit();
 
@@ -122,13 +134,13 @@ protected ReactorReceiver(AmqpConnection amqpConnection, String entityPath, Rece
                                     .addKeyValue("credits", credits)
                                     .log("Adding credits.");
                                 receiver.flow(credits);
-                                metricsProvider.recordAddCredits(credits);
                             } else {
                                 logger.atVerbose()
                                     .addKeyValue("credits", credits)
                                     .log("There are no credits to add.");
                             }
 
+                            metricsProvider.recordAddCredits(credits == null ? 0 : credits);
                             sink.success(message);
                         });
                     } catch (IOException | RejectedExecutionException e) {
@@ -448,6 +460,29 @@ private void completeClose() {
 
         handler.close();
         receiver.free();
+        try {
+            trackPrefetchSeqNoSubscription.close();
+        } catch (Exception e) {
+            logger.verbose("Error closing metrics subscription.", e);
+        }
+    }
+
+    private Long getSequenceNumber(Message message) {
+        if (message == null || message.getMessageAnnotations() == null || message.getBody() == null) {
+            return null;
+        }
+
+        Map<Symbol, Object> properties = message.getMessageAnnotations().getValue();
+        Object seqNo = properties != null ? properties.get(SEQUENCE_NUMBER_ANNOTATION) : null;
+        if (seqNo instanceof Integer) {
+            return ((Integer) seqNo).longValue();
+        } else if (seqNo instanceof Long) {
+            return (Long) seqNo;
+        } else if (seqNo != null) {
+            logger.verbose("Received message has unexpected `x-opt-sequence-number` annotation value - `{}`. Ignoring it.", seqNo);
+        }
+
+        return null;
     }
 
     @Override

sdk/core/azure-core-amqp/src/test/java/com/azure/core/amqp/implementation/AmqpMetricsProviderTest.java

@@ -5,6 +5,7 @@
 
 import com.azure.core.amqp.exception.AmqpResponseCode;
 import com.azure.core.test.utils.metrics.TestCounter;
+import com.azure.core.test.utils.metrics.TestGauge;
 import com.azure.core.test.utils.metrics.TestHistogram;
 import com.azure.core.test.utils.metrics.TestMeasurement;
 import com.azure.core.test.utils.metrics.TestMeter;
@@ -22,12 +23,13 @@
 
 import java.time.Instant;
 import java.util.Collections;
-import java.util.Date;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicLong;
 
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public class AmqpMetricsProviderTest {
@@ -56,7 +58,7 @@ public void disabledMeter() {
         assertDoesNotThrow(() -> provider.recordRequestResponseDuration(0, "foo", null));
         assertDoesNotThrow(() -> provider.recordSend(0, DeliveryState.DeliveryStateType.Accepted));
         assertDoesNotThrow(() -> provider.recordSend(0, null));
-        assertDoesNotThrow(() -> provider.recordReceivedMessage(Proton.message()));
+        assertDoesNotThrow(() -> provider.trackPrefetchSequenceNumber(() -> Long.valueOf(1)));
         assertDoesNotThrow(() -> provider.recordHandlerError(AmqpMetricsProvider.ErrorSource.TRANSPORT, new ErrorCondition(TIMEOUT_SYMBOL, "")));
         assertDoesNotThrow(() -> provider.recordHandlerError(null, new ErrorCondition(TIMEOUT_SYMBOL, "")));
 
@@ -77,7 +79,7 @@ public void defaultMeter() {
         assertDoesNotThrow(() -> provider.recordRequestResponseDuration(0, "foo", null));
         assertDoesNotThrow(() -> provider.recordSend(0, DeliveryState.DeliveryStateType.Accepted));
         assertDoesNotThrow(() -> provider.recordSend(0, null));
-        assertDoesNotThrow(() -> provider.recordReceivedMessage(Proton.message()));
+        assertDoesNotThrow(() -> provider.trackPrefetchSequenceNumber(() -> Long.valueOf(1)));
         assertDoesNotThrow(() -> provider.recordHandlerError(AmqpMetricsProvider.ErrorSource.LINK, new ErrorCondition(TIMEOUT_SYMBOL, "")));
         assertDoesNotThrow(() -> provider.recordHandlerError(null, new ErrorCondition(TIMEOUT_SYMBOL, "")));
     }
@@ -196,36 +198,46 @@ public void initCloseConnection() {
     }
 
     @Test
-    public void receivedMessage() {
+    @SuppressWarnings("try")
+    public void receivedMessage() throws Exception {
         TestMeter meter = new TestMeter();
         AmqpMetricsProvider provider = new AmqpMetricsProvider(meter, NAMESPACE, ENTITY_PATH);
 
-        Instant now = Instant.now();
-        MessageAnnotations futureEnqueuedDate = new MessageAnnotations(Collections.singletonMap(Symbol.valueOf("x-opt-enqueued-time"), Date.from(now.plusSeconds(100))));
+        AtomicLong seqNo = new AtomicLong(0);
 
-        provider.recordReceivedMessage(createMessage(Date.from(now.plusSeconds(100))));
-        provider.recordReceivedMessage(createMessage(Date.from(now.minusSeconds(100))));
+        TestGauge gauge;
+        try (AutoCloseable subscription = provider.trackPrefetchSequenceNumber(seqNo::get)) {
+            assertTrue(meter.getGauges().containsKey("messaging.az.amqp.prefetch.sequence_number"));
+            gauge = meter.getGauges().get("messaging.az.amqp.prefetch.sequence_number");
+            assertEquals(1, gauge.getSubscriptions().size());
+            assertSame(subscription, gauge.getSubscriptions().get(0));
+            TestGauge.Subscription testSubscription = (TestGauge.Subscription) subscription;
 
-        provider.recordReceivedMessage(Proton.message());
-        provider.recordReceivedMessage(createMessage("not a date"));
+            testSubscription.measure();
+            seqNo.set(1);
+            seqNo.set(2);
+            testSubscription.measure();
+            seqNo.set(42);
+            testSubscription.measure();
 
-        assertTrue(meter.getHistograms().containsKey("messaging.az.amqp.consumer.lag"));
-        TestHistogram histogram = meter.getHistograms().get("messaging.az.amqp.consumer.lag");
+            assertEquals(3, testSubscription.getMeasurements().size());
+            TestMeasurement<Long> measurement1 = testSubscription.getMeasurements().get(0);
 
-        assertEquals(2, histogram.getMeasurements().size());
-        TestMeasurement<Double> measurement1 = histogram.getMeasurements().get(0);
+            assertEquals(0, measurement1.getValue());
+            assertEquals(Context.NONE, measurement1.getContext());
+            assertCommonAttributes(measurement1.getAttributes(), NAMESPACE, ENTITY_NAME, ENTITY_PATH);
 
-        // negative delta becomes 0
-        assertEquals(0, measurement1.getValue());
-        assertEquals(Context.NONE, measurement1.getContext());
-        assertCommonAttributes(measurement1.getAttributes(), NAMESPACE, ENTITY_NAME, ENTITY_PATH);
+            assertEquals(2, testSubscription.getMeasurements().get(1).getValue());
+            assertEquals(42, testSubscription.getMeasurements().get(2).getValue());
+        }
 
-        assertEquals(100, histogram.getMeasurements().get(1).getValue(), 10);
+        seqNo.set(-1);
+        assertEquals(0, gauge.getSubscriptions().size());
     }
 
-    private Message createMessage(Object enqueuedTime) {
+    private Message createMessage(Object seqNo) {
         Message msg = Proton.message();
-        MessageAnnotations annotations = new MessageAnnotations(Collections.singletonMap(Symbol.valueOf("x-opt-enqueued-time"), enqueuedTime));
+        MessageAnnotations annotations = new MessageAnnotations(Collections.singletonMap(Symbol.valueOf("x-opt-sequence-number"), seqNo));
         msg.setMessageAnnotations(annotations);
         msg.setBody(new AmqpValue("body"));
 

sdk/core/azure-core-amqp/src/test/java/com/azure/core/amqp/implementation/ReactorReceiverTest.java

@@ -14,6 +14,7 @@
 import com.azure.core.amqp.exception.AmqpException;
 import com.azure.core.amqp.exception.AmqpResponseCode;
 import com.azure.core.amqp.implementation.handler.ReceiveLinkHandler;
+import com.azure.core.test.utils.metrics.TestGauge;
 import com.azure.core.test.utils.metrics.TestMeasurement;
 import com.azure.core.test.utils.metrics.TestMeter;
 import com.azure.core.util.Context;
@@ -45,7 +46,6 @@
 
 import java.io.IOException;
 import java.time.Duration;
-import java.time.Instant;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CountDownLatch;
@@ -766,13 +766,13 @@ void receiveMetricsAreReportedWithMessageIsReceived() throws IOException {
         // Arrange
         // This message was copied from one that was received.
         final byte[] messageBytes = new byte[] { 0, 83, 114, -63, 73, 6, -93, 21, 120, 45, 111, 112, 116, 45, 115, 101,
-            113, 117, 101, 110, 99, 101, 45, 110, 117, 109, 98, 101, 114, 85, 0, -93, 12, 120, 45, 111, 112, 116, 45,
+            113, 117, 101, 110, 99, 101, 45, 110, 117, 109, 98, 101, 114, 84, 42, -93, 12, 120, 45, 111, 112, 116, 45,
             111, 102, 102, 115, 101, 116, -95, 1, 48, -93, 19, 120, 45, 111, 112, 116, 45, 101, 110, 113, 117, 101, 117,
             101, 100, 45, 116, 105, 109, 101, -125, 0, 0, 1, 112, -54, 124, -41, 90, 0, 83, 117, -96, 12, 80, 111, 115,
             105, 116, 105, 111, 110, 53, 58, 32, 48};
 
         // change if changing message above
-        long messageEnqueuedTime = 1583945144154L;
+        long sequenceNumber = 42;
 
         final Link link = mock(Link.class);
         final Delivery delivery = mock(Delivery.class);
@@ -807,6 +807,9 @@ void receiveMetricsAreReportedWithMessageIsReceived() throws IOException {
         ReactorReceiver reactorReceiverWithMetrics = new ReactorReceiver(amqpConnection, "name/and/partition", receiver, receiverHandler, tokenManager,
             reactorDispatcher, retryOptions, metricsProvider);
 
+        TestGauge sequenceNumberMetric = meter.getGauges().get("messaging.az.amqp.prefetch.sequence_number");
+        TestGauge.Subscription subscription = sequenceNumberMetric.getSubscriptions().get(0);
+
         reactorReceiverWithMetrics.setEmptyCreditListener(creditSupplier);
 
         doAnswer(invocationOnMock -> {
@@ -822,10 +825,11 @@ void receiveMetricsAreReportedWithMessageIsReceived() throws IOException {
             .verify(VERIFY_TIMEOUT);
 
         // Assert
-        List<TestMeasurement<Double>> receivedLag = meter.getHistograms().get("messaging.az.amqp.consumer.lag").getMeasurements();
-        assertEquals(1, receivedLag.size());
-        TestMeasurement<Double> measurement = receivedLag.get(0);
-        assertEquals((Instant.now().toEpochMilli() - messageEnqueuedTime) / 1000d, measurement.getValue(), 100);
+        subscription.measure();
+        List<TestMeasurement<Long>> seqNumbers = subscription.getMeasurements();
+        assertEquals(1, seqNumbers.size());
+        TestMeasurement<Long> measurement = seqNumbers.get(0);
+        assertEquals(sequenceNumber, measurement.getValue());
         assertEquals(Context.NONE, measurement.getContext());
         assertEquals("namespace", measurement.getAttributes().get(ClientConstants.HOSTNAME_KEY));
         assertEquals("name", measurement.getAttributes().get(ClientConstants.ENTITY_NAME_KEY));