Orient DataStream around InputStream

DataStream now is written around the idea of using an underyling
InputStream more often than a Publisher, so default implementations
were updated (now default provided for asPublisher, and asInputStream
no longer has a default).

Also closes InputStreams in a couple places that we could have leaked.

Addresses smithy-lang#913

Author: mtdowling

aws/aws-sigv4/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigV4Signer.java

@@ -130,7 +130,7 @@ public HttpRequest sign(HttpRequest request, AwsCredentialsIdentity identity, Co
     }
 
     private String getPayloadHash(DataStream dataStream) {
-        return hexHash(dataStream.waitForByteBuffer());
+        return hexHash(dataStream.asByteBuffer());
     }
 
     private String hexHash(ByteBuffer bytes) {

aws/client/aws-client-awsjson/src/it/java/software/amazon/smithy/java/client/aws/jsonprotocols/AwsJson1ProtocolTests.java

@@ -39,7 +39,7 @@ public void requestTest(DataStream expected, DataStream actual) {
         if (expected.contentLength() != 0) {
             // Use the node parser to strip out white space.
             expectedJson = Node.printJson(
-                    Node.parse(new String(ByteBufferUtils.getBytes(expected.waitForByteBuffer()),
+                    Node.parse(new String(ByteBufferUtils.getBytes(expected.asByteBuffer()),
                             StandardCharsets.UTF_8)));
         }
         assertEquals(expectedJson, new StringBuildingSubscriber(actual).getResult());

aws/client/aws-client-awsjson/src/main/java/software/amazon/smithy/java/aws/client/awsjson/AwsJsonProtocol.java

@@ -101,7 +101,7 @@ public <I extends SerializableStruct, O extends SerializableStruct> O deserializ
             return codec.deserializeShape(EMPTY_PAYLOAD, builder);
         }
 
-        var bytes = content.waitForByteBuffer();
+        var bytes = content.asByteBuffer();
         return codec.deserializeShape(bytes, builder);
     }
 }

aws/client/aws-client-restjson/src/it/java/software/amazon/smithy/java/client/aws/restjson/RestJson1ProtocolTests.java

@@ -41,7 +41,7 @@ public void requestTest(DataStream expected, DataStream actual) {
         var actualStr = new StringBuildingSubscriber(actual).getResult();
         if (expected.contentLength() != 0) {
             var expectedStr = new String(
-                    ByteBufferUtils.getBytes(expected.waitForByteBuffer()),
+                    ByteBufferUtils.getBytes(expected.asByteBuffer()),
                     StandardCharsets.UTF_8);
             if ("application/json".equals(expected.contentType())) {
                 var expectedNode = Node.parse(expectedStr);

aws/client/aws-client-restxml/src/it/java/software/amazon/smithy/java/aws/client/restxml/RestXmlProtocolTests.java

@@ -41,8 +41,8 @@ public class RestXmlProtocolTests {
             })
     public void requestTest(DataStream expected, DataStream actual) {
         if (expected.contentLength() != 0) {
-            var a = new String(ByteBufferUtils.getBytes(actual.waitForByteBuffer()), StandardCharsets.UTF_8);
-            var b = new String(ByteBufferUtils.getBytes(expected.waitForByteBuffer()), StandardCharsets.UTF_8);
+            var a = new String(ByteBufferUtils.getBytes(actual.asByteBuffer()), StandardCharsets.UTF_8);
+            var b = new String(ByteBufferUtils.getBytes(expected.asByteBuffer()), StandardCharsets.UTF_8);
             if ("application/xml".equals(expected.contentType())) {
                 if (!XMLComparator.compareXMLStrings(a, b)) {
                     // Do this comparison to see what is different.

aws/integrations/aws-lambda-endpoint/src/main/java/software/amazon/smithy/java/aws/integrations/lambda/LambdaEndpoint.java

@@ -157,7 +157,7 @@ private static ProxyResponse getResponse(HttpResponse httpResponse, boolean shou
 
         DataStream val = httpResponse.getSerializedValue();
         if (val != null) {
-            ByteBuffer buf = val.waitForByteBuffer();
+            ByteBuffer buf = val.asByteBuffer();
             String body;
             // TODO: handle base64 encoding better
             if (shouldBase64Encode) {

aws/server/aws-server-restjson/src/it/java/software/amazon/smithy/java/aws/server/restjson/AwsRestJson1ProtocolTests.java

@@ -85,9 +85,9 @@ public void responseTest(DataStream expected, DataStream actual) {
                 .isTrue()
                 .isSameAs(actual.hasKnownLength());
 
-        String actualJson = new String(ByteBufferUtils.getBytes(actual.waitForByteBuffer()), StandardCharsets.UTF_8);
+        String actualJson = new String(ByteBufferUtils.getBytes(actual.asByteBuffer()), StandardCharsets.UTF_8);
         String expectedJson = new String(
-                ByteBufferUtils.getBytes(expected.waitForByteBuffer()),
+                ByteBufferUtils.getBytes(expected.asByteBuffer()),
                 StandardCharsets.UTF_8);
         if (expected.contentLength() == 0) {
             assertThat(actualJson).isIn("", "{}");

client/client-http/src/main/java/software/amazon/smithy/java/client/http/HttpErrorDeserializer.java

@@ -138,7 +138,7 @@ public ModeledException createError(
                 ShapeBuilder<ModeledException> builder
         ) {
             try {
-                ByteBuffer bytes = createDataStream(response).waitForByteBuffer();
+                ByteBuffer bytes = createDataStream(response).asByteBuffer();
                 return codec.deserializeShape(bytes, builder);
             } catch (Exception e) {
                 throw new RuntimeException("Failed to deserialize error", e);
@@ -248,7 +248,7 @@ private static CallException makeErrorFromPayload(
         try {
             // Read the payload into a JSON document so we can efficiently find __type and then directly
             // deserialize the document into the identified builder.
-            ByteBuffer buffer = content.waitForByteBuffer();
+            ByteBuffer buffer = content.asByteBuffer();
 
             if (buffer.remaining() > 0) {
                 var document = codec.createDeserializer(buffer).readDocument();

client/client-http/src/main/java/software/amazon/smithy/java/client/http/JavaHttpClientTransport.java

@@ -5,6 +5,9 @@
 
 package software.amazon.smithy.java.client.http;
 
+import static java.net.http.HttpRequest.BodyPublishers;
+import static java.net.http.HttpResponse.BodyHandlers;
+
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.URI;
@@ -110,16 +113,15 @@ private java.net.http.HttpRequest createJavaRequest(Context context, HttpRequest
         java.net.http.HttpRequest.BodyPublisher bodyPublisher;
         if (request.body().hasKnownLength()) {
             if (request.body().contentLength() == 0) {
-                bodyPublisher = java.net.http.HttpRequest.BodyPublishers.noBody();
+                bodyPublisher = BodyPublishers.noBody();
             } else {
-                bodyPublisher = java.net.http.HttpRequest.BodyPublishers.ofByteArray(
-                        ByteBufferUtils.getBytes(request.body().waitForByteBuffer()));
+                bodyPublisher = BodyPublishers.ofByteArray(ByteBufferUtils.getBytes(request.body().asByteBuffer()));
             }
         } else {
-            bodyPublisher = java.net.http.HttpRequest.BodyPublishers.fromPublisher(request.body());
+            bodyPublisher = BodyPublishers.fromPublisher(request.body());
         }
 
-        java.net.http.HttpRequest.Builder httpRequestBuilder = java.net.http.HttpRequest.newBuilder()
+        var httpRequestBuilder = java.net.http.HttpRequest.newBuilder()
                 .version(smithyToHttpVersion(request.httpVersion()))
                 .method(request.method(), bodyPublisher)
                 .uri(request.uri());
@@ -141,13 +143,24 @@ private java.net.http.HttpRequest createJavaRequest(Context context, HttpRequest
     }
 
     private HttpResponse sendRequest(java.net.http.HttpRequest request) {
+        java.net.http.HttpResponse<InputStream> res = null;
         try {
-            var res = client.send(request, java.net.http.HttpResponse.BodyHandlers.ofInputStream());
+            res = client.send(request, BodyHandlers.ofInputStream());
             return createSmithyResponse(res);
         } catch (IOException | InterruptedException | RuntimeException e) {
+            // Close the response body stream if we got a response but failed to process it
+            if (res != null) {
+                try {
+                    res.body().close();
+                } catch (IOException closeException) {
+                    LOGGER.trace("Failed to close response body after error", closeException);
+                }
+            }
+
             if (e instanceof HttpConnectTimeoutException) {
                 throw new ConnectTimeoutException(e);
             }
+
             // The client pipeline also does this remapping, but to adhere to the required contract of
             // ClientTransport, we remap here too if needed.
             throw ClientTransport.remapExceptions(e);

client/client-http/src/main/java/software/amazon/smithy/java/client/http/plugins/HttpChecksumPlugin.java

@@ -50,7 +50,7 @@ private static HttpRequest processRequest(RequestHook<?, ?, HttpRequest> hook) {
         static HttpRequest addContentMd5Header(HttpRequest request) {
             var body = request.body();
             if (body != null) {
-                var buffer = body.waitForByteBuffer();
+                var buffer = body.asByteBuffer();
                 var bytes = ByteBufferUtils.getBytes(buffer);
                 try {
                     byte[] hash = MessageDigest.getInstance("MD5").digest(bytes);