addDecodeStage (Azure#22808)

* add decode stage

Co-authored-by: annie-mac <[email protected]>
Co-authored-by: annie-mac <[email protected]>

Author: 3 people

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestManager.java

@@ -55,6 +55,7 @@
 import javax.net.ssl.SSLException;
 import java.net.SocketAddress;
 import java.nio.channels.ClosedChannelException;
+import java.time.Instant;
 import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
@@ -730,8 +731,14 @@ private void messageReceived(final ChannelHandlerContext context, final RntbdRes
             return;
         }
 
+        requestRecord.stage(RntbdRequestRecord.Stage.DECODE_STARTED, response.getDecodeStartTime());
+
+        // When decode completed, it means sdk has received the full response from server
+        requestRecord.stage(
+                RntbdRequestRecord.Stage.RECEIVED,
+                response.getDecodeEndTime() != null ? response.getDecodeEndTime() : Instant.now());
+
         requestRecord.responseLength(response.getMessageLength());
-        requestRecord.stage(RntbdRequestRecord.Stage.RECEIVED);
 
         final HttpResponseStatus status = response.getStatus();
         final UUID activityId = response.getActivityId();

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdRequestRecord.java

@@ -63,6 +63,7 @@ public abstract class RntbdRequestRecord extends CompletableFuture<StoreResponse
     private volatile Instant timePipelined;
     private final Instant timeQueued;
     private volatile Instant timeSent;
+    private volatile Instant timeDecodeStarted;
     private volatile Instant timeReceived;
     private volatile boolean sendingRequestHasStarted;
     private volatile RntbdChannelAcquisitionTimeline channelAcquisitionTimeline;
@@ -115,8 +116,10 @@ public Stage stage() {
     }
 
     public RntbdRequestRecord stage(final Stage value) {
+        return this.stage(value, Instant.now());
+    }
 
-        final Instant time = Instant.now();
+    public RntbdRequestRecord stage(final Stage value, Instant time) {
 
         STAGE.updateAndGet(this, current -> {
 
@@ -143,9 +146,18 @@ public RntbdRequestRecord stage(final Stage value) {
                     }
                     this.timeSent = time;
                     break;
-                case RECEIVED:
+
+                case DECODE_STARTED:
                     if (current != Stage.SENT) {
-                        logger.debug("Expected transition from SENT to RECEIVED, not {} to RECEIVED", current);
+                        logger.debug("Expected transition from SENT to DECODE_STARTED, not {} to DECODE_STARTED", current);
+                        break;
+                    }
+                    this.timeDecodeStarted = time;
+                    break;
+
+                case RECEIVED:
+                    if (current != Stage.DECODE_STARTED) {
+                        logger.debug("Expected transition from DECODE_STARTED to RECEIVED, not {} to RECEIVED", current);
                         break;
                     }
                     this.timeReceived = time;
@@ -181,6 +193,8 @@ public Instant timeCreated() {
         return this.args.timeCreated();
     }
 
+    public Instant timeDecodeStarted() { return this.timeDecodeStarted; }
+
     public Instant timePipelined() {
         return this.timePipelined;
     }
@@ -272,6 +286,7 @@ public RequestTimeline takeTimelineSnapshot() {
         Instant timeChannelAcquisitionStarted = this.timeChannelAcquisitionStarted();
         Instant timePipelined = this.timePipelined();
         Instant timeSent = this.timeSent();
+        Instant timeDecodeStarted = this.timeDecodeStarted();
         Instant timeReceived = this.timeReceived();
         Instant timeCompleted = this.timeCompleted();
         Instant timeCompletedOrNow = timeCompleted == null ? now : timeCompleted;
@@ -286,7 +301,9 @@ public RequestTimeline takeTimelineSnapshot() {
             new RequestTimeline.Event("pipelined",
                 timePipelined, timeSent == null ? timeCompletedOrNow : timeSent),
             new RequestTimeline.Event("transitTime",
-                timeSent, timeReceived == null ? timeCompletedOrNow : timeReceived),
+                timeSent, timeDecodeStarted == null ? timeCompletedOrNow : timeDecodeStarted),
+            new RequestTimeline.Event("decodeTime",
+                timeDecodeStarted, timeReceived == null ? timeCompletedOrNow : timeReceived),
             new RequestTimeline.Event("received",
                 timeReceived, timeCompletedOrNow),
             new RequestTimeline.Event("completed",
@@ -307,7 +324,7 @@ public String toString() {
     // region Types
 
     public enum Stage {
-        QUEUED, CHANNEL_ACQUISITION_STARTED, PIPELINED, SENT, RECEIVED, COMPLETED
+        QUEUED, CHANNEL_ACQUISITION_STARTED, PIPELINED, SENT, DECODE_STARTED, RECEIVED, COMPLETED
     }
 
     static final class JsonSerializer extends StdSerializer<RntbdRequestRecord> {

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdResponse.java

@@ -21,6 +21,7 @@
 
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
+import java.time.Instant;
 import java.util.Map;
 import java.util.UUID;
 import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
@@ -55,6 +56,12 @@ public final class RntbdResponse implements ReferenceCounted {
     @JsonProperty
     private volatile int referenceCount;
 
+    @JsonIgnore
+    private Instant decodeStartTime;
+
+    @JsonIgnore
+    private Instant decodeEndTime;
+
     // endregion
 
     // region Constructors
@@ -135,6 +142,22 @@ public Long getTransportRequestId() {
         return this.getHeader(RntbdResponseHeader.TransportRequestID);
     }
 
+    public Instant getDecodeStartTime() {
+        return this.decodeStartTime;
+    }
+
+    public void setDecodeStartTime(Instant decodeStartTime) {
+        this.decodeStartTime = decodeStartTime;
+    }
+
+    public Instant getDecodeEndTime() {
+        return this.decodeEndTime;
+    }
+
+    public void setDecodeEndTime(Instant decodeEndTime) {
+        this.decodeEndTime = decodeEndTime;
+    }
+
     // endregion
 
     // region Methods

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/directconnectivity/rntbd/RntbdResponseDecoder.java

@@ -9,11 +9,14 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.time.Instant;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
 
 public final class RntbdResponseDecoder extends ByteToMessageDecoder {
 
     private static final Logger logger = LoggerFactory.getLogger(RntbdResponseDecoder.class);
+    private static final AtomicReference<Instant> decodeStartTime = new AtomicReference<>();
 
     /**
      * Deserialize from an input {@link ByteBuf} to an {@link RntbdResponse} instance.
@@ -27,11 +30,16 @@ public final class RntbdResponseDecoder extends ByteToMessageDecoder {
     @Override
     protected void decode(final ChannelHandlerContext context, final ByteBuf in, final List<Object> out) {
 
+        decodeStartTime.compareAndSet(null, Instant.now());
+
         if (RntbdFramer.canDecodeHead(in)) {
 
             final RntbdResponse response = RntbdResponse.decode(in);
 
             if (response != null) {
+                response.setDecodeEndTime(Instant.now());
+                response.setDecodeStartTime(decodeStartTime.getAndSet(null));
+
                 logger.debug("{} DECODE COMPLETE: {}", context.channel(), response);
                 in.discardReadBytes();
                 out.add(response.retain());

sdk/cosmos/azure-cosmos/src/test/java/com/azure/cosmos/CosmosDiagnosticsTest.java

@@ -1051,6 +1051,7 @@ private void validateTransportRequestTimelineDirect(String diagnostics) {
         assertThat(diagnostics).contains("\"eventName\":\"channelAcquisitionStarted\"");
         assertThat(diagnostics).contains("\"eventName\":\"pipelined\"");
         assertThat(diagnostics).contains("\"eventName\":\"transitTime\"");
+        assertThat(diagnostics).contains("\"eventName\":\"decodeTime");
         assertThat(diagnostics).contains("\"eventName\":\"received\"");
         assertThat(diagnostics).contains("\"eventName\":\"completed\"");
         assertThat(diagnostics).contains("\"startTimeUTC\"");