Miscellaneous Improvement to azure-core (Azure#27638)

Miscellaneous Improvement to azure-core

Author: alzimmermsft

sdk/core/azure-core/src/main/java/com/azure/core/annotation/ReturnValueWireType.java

@@ -4,7 +4,6 @@
 package com.azure.core.annotation;
 
 import com.azure.core.http.rest.Page;
-import com.azure.core.implementation.UnixTime;
 import com.azure.core.util.Base64Url;
 import com.azure.core.util.DateTimeRfc1123;
 
@@ -22,7 +21,6 @@
  * <ol>
  *     <li>{@link Base64Url}</li>
  *     <li>{@link DateTimeRfc1123}</li>
- *     <li>{@link UnixTime}</li>
  *     <li>{@link Page}</li>
  *     <li>{@link List List&lt;T&gt;} where {@code T} can be one of the four values above.</li>
  * </ol>

sdk/core/azure-core/src/main/java/com/azure/core/http/HttpPipeline.java

@@ -28,7 +28,7 @@ public final class HttpPipeline {
      * HttpPipeline#send(HttpPipelineCallContext)} and it's response.
      *
      * @param httpClient the http client to write request to wire and receive response from wire.
-     * @param pipelinePolicies pipeline policies in the order they need to applied, a copy of this array will be made
+     * @param pipelinePolicies pipeline policies in the order they need to be applied, a copy of this array will be made
      * hence changing the original array after the creation of pipeline will not  mutate the pipeline
      */
     HttpPipeline(HttpClient httpClient, List<HttpPipelinePolicy> pipelinePolicies) {
@@ -41,7 +41,7 @@ public final class HttpPipeline {
     /**
      * Get the policy at the passed index in the pipeline.
      *
-     * @param index index of the the policy to retrieve.
+     * @param index index of the policy to retrieve.
      * @return the policy stored at that index.
      */
     public HttpPipelinePolicy getPolicy(final int index) {

sdk/core/azure-core/src/main/java/com/azure/core/http/policy/AddDatePolicy.java

@@ -6,26 +6,32 @@
 import com.azure.core.http.HttpPipelineCallContext;
 import com.azure.core.http.HttpPipelineNextPolicy;
 import com.azure.core.http.HttpResponse;
+import com.azure.core.util.DateTimeRfc1123;
 import reactor.core.publisher.Mono;
 
 import java.time.OffsetDateTime;
-import java.time.ZoneId;
+import java.time.ZoneOffset;
 import java.time.format.DateTimeFormatter;
 import java.util.Locale;
 
 /**
  * The pipeline policy that adds a "Date" header in RFC 1123 format when sending an HTTP request.
  */
 public class AddDatePolicy implements HttpPipelinePolicy {
-    private final DateTimeFormatter format = DateTimeFormatter
+    private static final DateTimeFormatter FORMATTER = DateTimeFormatter
             .ofPattern("EEE, dd MMM yyyy HH:mm:ss 'GMT'")
-            .withZone(ZoneId.of("UTC"))
+            .withZone(ZoneOffset.UTC)
             .withLocale(Locale.US);
 
     @Override
     public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) {
         return Mono.defer(() -> {
-            context.getHttpRequest().getHeaders().set("Date", format.format(OffsetDateTime.now()));
+            OffsetDateTime now = OffsetDateTime.now();
+            try {
+                context.getHttpRequest().getHeaders().set("Date", DateTimeRfc1123.toRfc1123String(now));
+            } catch (IllegalArgumentException ignored) {
+                context.getHttpRequest().getHeaders().set("Date", FORMATTER.format(now));
+            }
             return next.process();
         });
     }

sdk/core/azure-core/src/main/java/com/azure/core/http/policy/AzureKeyCredentialPolicy.java

@@ -19,7 +19,7 @@
  * an exception will be thrown to prevent leaking the key.
  */
 public final class AzureKeyCredentialPolicy implements HttpPipelinePolicy {
-    // AzureKeyCredentailPolicy can be a commonly used policy, use a static logger.
+    // AzureKeyCredentialPolicy can be a commonly used policy, use a static logger.
     private static final ClientLogger LOGGER = new ClientLogger(AzureKeyCredentialPolicy.class);
 
     private final String name;

sdk/core/azure-core/src/main/java/com/azure/core/http/policy/FixedDelayOptions.java

@@ -17,7 +17,7 @@ public class FixedDelayOptions {
     private final Duration delay;
 
     /**
-     * Creates an instance of {@link maxRetries}.
+     * Creates an instance of {@link FixedDelayOptions}.
      *
      * @param maxRetries The max number of retry attempts that can be made.
      * @param delay The fixed delay duration between retry attempts.

sdk/core/azure-core/src/main/java/com/azure/core/http/policy/HttpLoggingPolicy.java

@@ -10,22 +10,25 @@
 import com.azure.core.http.HttpPipelineNextPolicy;
 import com.azure.core.http.HttpRequest;
 import com.azure.core.http.HttpResponse;
+import com.azure.core.implementation.AccessibleByteArrayOutputStream;
+import com.azure.core.implementation.ImplUtils;
 import com.azure.core.implementation.http.HttpPipelineCallContextHelper;
 import com.azure.core.implementation.jackson.ObjectMapperShim;
 import com.azure.core.util.Context;
 import com.azure.core.util.CoreUtils;
+import com.azure.core.util.FluxUtil;
 import com.azure.core.util.UrlBuilder;
 import com.azure.core.util.logging.ClientLogger;
 import com.azure.core.util.logging.LogLevel;
+import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
-import java.io.ByteArrayOutputStream;
 import java.io.IOException;
-import java.io.UnsupportedEncodingException;
+import java.io.UncheckedIOException;
 import java.net.URL;
 import java.nio.ByteBuffer;
-import java.nio.channels.Channels;
-import java.nio.channels.WritableByteChannel;
+import java.nio.charset.Charset;
+import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.util.Collections;
 import java.util.Locale;
@@ -122,14 +125,14 @@ public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineN
     private HttpRequestLoggingContext getRequestLoggingOptions(HttpPipelineCallContext callContext) {
         return new HttpRequestLoggingContext(callContext.getHttpRequest(),
             HttpPipelineCallContextHelper.getContext(callContext),
-            getRequestRetryCount(HttpPipelineCallContextHelper.getContext(callContext), LOGGER));
+            getRequestRetryCount(HttpPipelineCallContextHelper.getContext(callContext)));
     }
 
     private HttpResponseLoggingContext getResponseLoggingOptions(HttpResponse httpResponse, long startNs,
         HttpPipelineCallContext callContext) {
         return new HttpResponseLoggingContext(httpResponse, Duration.ofNanos(System.nanoTime() - startNs),
             HttpPipelineCallContextHelper.getContext(callContext),
-            getRequestRetryCount(HttpPipelineCallContextHelper.getContext(callContext), LOGGER));
+            getRequestRetryCount(HttpPipelineCallContextHelper.getContext(callContext)));
     }
 
     private final class DefaultHttpRequestLogger implements HttpRequestLogger {
@@ -164,7 +167,8 @@ public Mono<Void> logRequest(ClientLogger logger, HttpRequestLoggingContext logg
             }
 
             if (!httpLogDetailLevel.shouldLogBody()) {
-                return logAndReturn(logger, logLevel, requestLogMessage, null);
+                logMessage(logger, logLevel, requestLogMessage);
+                return Mono.empty();
             }
 
             if (request.getBody() == null) {
@@ -174,35 +178,40 @@ public Mono<Void> logRequest(ClientLogger logger, HttpRequestLoggingContext logg
                     .append(request.getHttpMethod())
                     .append(System.lineSeparator());
 
-                return logAndReturn(logger, logLevel, requestLogMessage, null);
+                logMessage(logger, logLevel, requestLogMessage);
+                return Mono.empty();
             }
 
             String contentType = request.getHeaders().getValue("Content-Type");
             long contentLength = getContentLength(logger, request.getHeaders());
 
             if (shouldBodyBeLogged(contentType, contentLength)) {
-                ByteArrayOutputStream outputStream = new ByteArrayOutputStream((int) contentLength);
-                WritableByteChannel bodyContentChannel = Channels.newChannel(outputStream);
+                AccessibleByteArrayOutputStream stream = new AccessibleByteArrayOutputStream((int) contentLength);
 
                 // Add non-mutating operators to the data stream.
                 request.setBody(
                     request.getBody()
-                        .flatMap(byteBuffer -> writeBufferToBodyStream(bodyContentChannel, byteBuffer))
+                        .doOnNext(byteBuffer -> {
+                            try {
+                                ImplUtils.writeByteBufferToStream(byteBuffer.duplicate(), stream);
+                            } catch (IOException ex) {
+                                throw LOGGER.logExceptionAsError(new UncheckedIOException(ex));
+                            }
+                        })
                         .doFinally(ignored -> {
                             requestLogMessage.append(contentLength)
                                 .append("-byte body:")
                                 .append(System.lineSeparator())
                                 .append(prettyPrintIfNeeded(logger, prettyPrintBody, contentType,
-                                    convertStreamToString(outputStream, logger)))
+                                    new String(stream.toByteArray(), 0, stream.count(), StandardCharsets.UTF_8)))
                                 .append(System.lineSeparator())
                                 .append("--> END ")
                                 .append(request.getHttpMethod())
                                 .append(System.lineSeparator());
 
-                            logAndReturn(logger, logLevel, requestLogMessage, null);
+                            logMessage(logger, logLevel, requestLogMessage);
                         }));
 
-                return Mono.empty();
             } else {
                 requestLogMessage.append(contentLength)
                     .append("-byte body: (content not logged)")
@@ -211,8 +220,10 @@ public Mono<Void> logRequest(ClientLogger logger, HttpRequestLoggingContext logg
                     .append(request.getHttpMethod())
                     .append(System.lineSeparator());
 
-                return logAndReturn(logger, logLevel, requestLogMessage, null);
+                logMessage(logger, logLevel, requestLogMessage);
             }
+
+            return Mono.empty();
         }
     }
 
@@ -251,40 +262,28 @@ public Mono<HttpResponse> logResponse(ClientLogger logger, HttpResponseLoggingCo
 
             if (!httpLogDetailLevel.shouldLogBody()) {
                 responseLogMessage.append("<-- END HTTP");
-                return logAndReturn(logger, logLevel, responseLogMessage, response);
+                logMessage(logger, logLevel, responseLogMessage);
+                return Mono.justOrEmpty(response);
             }
 
             String contentTypeHeader = response.getHeaderValue("Content-Type");
             long contentLength = getContentLength(logger, response.getHeaders());
 
             if (shouldBodyBeLogged(contentTypeHeader, contentLength)) {
-                HttpResponse bufferedResponse = response.buffer();
-                ByteArrayOutputStream outputStream = new ByteArrayOutputStream((int) contentLength);
-                WritableByteChannel bodyContentChannel = Channels.newChannel(outputStream);
-                return bufferedResponse.getBody()
-                    .flatMap(byteBuffer -> writeBufferToBodyStream(bodyContentChannel, byteBuffer))
-                    .doFinally(ignored -> {
-                        responseLogMessage.append("Response body:")
-                            .append(System.lineSeparator())
-                            .append(prettyPrintIfNeeded(logger, prettyPrintBody, contentTypeHeader,
-                                convertStreamToString(outputStream, logger)))
-                            .append(System.lineSeparator())
-                            .append("<-- END HTTP");
-
-                        logAndReturn(logger, logLevel, responseLogMessage, response);
-                    }).then(Mono.just(bufferedResponse));
+                return Mono.just(new LoggingHttpResponse(response, responseLogMessage, logger, logLevel,
+                    (int) contentLength, contentTypeHeader, prettyPrintBody));
             } else {
                 responseLogMessage.append("(body content not logged)")
                     .append(System.lineSeparator())
                     .append("<-- END HTTP");
 
-                return logAndReturn(logger, logLevel, responseLogMessage, response);
+                logMessage(logger, logLevel, responseLogMessage);
+                return Mono.just(response);
             }
         }
     }
 
-    private static <T> Mono<T> logAndReturn(ClientLogger logger, LogLevel logLevel, StringBuilder logMessageBuilder,
-        T data) {
+    private static void logMessage(ClientLogger logger, LogLevel logLevel, StringBuilder logMessageBuilder) {
         switch (logLevel) {
             case VERBOSE:
                 logger.verbose(logMessageBuilder.toString());
@@ -305,8 +304,6 @@ private static <T> Mono<T> logAndReturn(ClientLogger logger, LogLevel logLevel,
             default:
                 break;
         }
-
-        return Mono.justOrEmpty(data);
     }
 
     /*
@@ -450,36 +447,13 @@ private static boolean shouldBodyBeLogged(String contentTypeHeader, long content
             && contentLength < MAX_BODY_LOG_SIZE;
     }
 
-    /*
-     * Helper function which converts a ByteArrayOutputStream to a String without duplicating the internal buffer.
-     */
-    private static String convertStreamToString(ByteArrayOutputStream stream, ClientLogger logger) {
-        try {
-            return stream.toString("UTF-8");
-        } catch (UnsupportedEncodingException ex) {
-            throw logger.logExceptionAsError(new RuntimeException(ex));
-        }
-    }
-
-    /*
-     * Helper function which writes body ByteBuffers into the body message channel.
-     */
-    private static Mono<ByteBuffer> writeBufferToBodyStream(WritableByteChannel channel, ByteBuffer byteBuffer) {
-        try {
-            channel.write(byteBuffer.duplicate());
-            return Mono.just(byteBuffer);
-        } catch (IOException ex) {
-            return Mono.error(ex);
-        }
-    }
-
     /*
      * Gets the request retry count to include in logging.
      *
      * If there is no value set or it isn't a valid number null will be returned indicating that retry count won't be
      * logged.
      */
-    private static Integer getRequestRetryCount(Context context, ClientLogger logger) {
+    private static Integer getRequestRetryCount(Context context) {
         Object rawRetryCount = context.getData(RETRY_COUNT_CONTEXT).orElse(null);
         if (rawRetryCount == null) {
             return null;
@@ -488,7 +462,7 @@ private static Integer getRequestRetryCount(Context context, ClientLogger logger
         try {
             return Integer.valueOf(rawRetryCount.toString());
         } catch (NumberFormatException ex) {
-            logger.warning("Could not parse the request retry count: '{}'.", rawRetryCount);
+            LOGGER.warning("Could not parse the request retry count: '{}'.", rawRetryCount);
             return null;
         }
     }
@@ -503,4 +477,81 @@ private static ClientLogger getOrCreateMethodLogger(String methodName) {
 
         return CALLER_METHOD_LOGGER_CACHE.computeIfAbsent(methodName, ClientLogger::new);
     }
+
+    private static final class LoggingHttpResponse extends HttpResponse {
+        private final HttpResponse actualResponse;
+        private final StringBuilder responseLogMessage;
+        private final int contentLength;
+        private final ClientLogger logger;
+        private final boolean prettyPrintBody;
+        private final String contentTypeHeader;
+        private final LogLevel logLevel;
+
+        private LoggingHttpResponse(HttpResponse actualResponse, StringBuilder responseLogMessage,
+            ClientLogger logger, LogLevel logLevel, int contentLength, String contentTypeHeader,
+            boolean prettyPrintBody) {
+            super(actualResponse.getRequest());
+            this.actualResponse = actualResponse;
+            this.responseLogMessage = responseLogMessage;
+            this.logger = logger;
+            this.logLevel = logLevel;
+            this.contentLength = contentLength;
+            this.contentTypeHeader = contentTypeHeader;
+            this.prettyPrintBody = prettyPrintBody;
+        }
+
+        @Override
+        public int getStatusCode() {
+            return actualResponse.getStatusCode();
+        }
+
+        @Override
+        public String getHeaderValue(String name) {
+            return actualResponse.getHeaderValue(name);
+        }
+
+        @Override
+        public HttpHeaders getHeaders() {
+            return actualResponse.getHeaders();
+        }
+
+        @Override
+        public Flux<ByteBuffer> getBody() {
+            AccessibleByteArrayOutputStream stream = new AccessibleByteArrayOutputStream(contentLength);
+
+            return actualResponse.getBody()
+                .doOnNext(byteBuffer -> {
+                    try {
+                        ImplUtils.writeByteBufferToStream(byteBuffer.duplicate(), stream);
+                    } catch (IOException ex) {
+                        throw LOGGER.logExceptionAsError(new UncheckedIOException(ex));
+                    }
+                })
+                .doFinally(ignored -> {
+                    responseLogMessage.append("Response body:")
+                        .append(System.lineSeparator())
+                        .append(prettyPrintIfNeeded(logger, prettyPrintBody, contentTypeHeader,
+                            new String(stream.toByteArray(), 0, stream.count(), StandardCharsets.UTF_8)))
+                        .append(System.lineSeparator())
+                        .append("<-- END HTTP");
+
+                    logMessage(logger, logLevel, responseLogMessage);
+                });
+        }
+
+        @Override
+        public Mono<byte[]> getBodyAsByteArray() {
+            return FluxUtil.collectBytesFromNetworkResponse(getBody(), actualResponse.getHeaders());
+        }
+
+        @Override
+        public Mono<String> getBodyAsString() {
+            return getBodyAsByteArray().map(String::new);
+        }
+
+        @Override
+        public Mono<String> getBodyAsString(Charset charset) {
+            return getBodyAsByteArray().map(bytes -> new String(bytes, charset));
+        }
+    }
 }