diff --git a/runtime/runtime-core/api/runtime-core.api b/runtime/runtime-core/api/runtime-core.api index 0e8f85de3a..e488735251 100644 --- a/runtime/runtime-core/api/runtime-core.api +++ b/runtime/runtime-core/api/runtime-core.api @@ -1097,6 +1097,7 @@ public abstract interface class aws/smithy/kotlin/runtime/io/SdkSource : java/io public final class aws/smithy/kotlin/runtime/io/SdkSourceKt { public static final fun readFully (Laws/smithy/kotlin/runtime/io/SdkSource;Laws/smithy/kotlin/runtime/io/SdkBuffer;J)V + public static final fun readRemaining (Laws/smithy/kotlin/runtime/io/SdkSource;Laws/smithy/kotlin/runtime/io/SdkBuffer;)J public static final fun readToByteArray (Laws/smithy/kotlin/runtime/io/SdkSource;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; public static final fun toSdkByteReadChannel (Laws/smithy/kotlin/runtime/io/SdkSource;Lkotlinx/coroutines/CoroutineScope;)Laws/smithy/kotlin/runtime/io/SdkByteReadChannel; public static synthetic fun toSdkByteReadChannel$default (Laws/smithy/kotlin/runtime/io/SdkSource;Lkotlinx/coroutines/CoroutineScope;ILjava/lang/Object;)Laws/smithy/kotlin/runtime/io/SdkByteReadChannel; diff --git a/runtime/runtime-core/common/src/aws/smithy/kotlin/runtime/io/SdkSource.kt b/runtime/runtime-core/common/src/aws/smithy/kotlin/runtime/io/SdkSource.kt index aefb75dc08..061e2afaf2 100644 --- a/runtime/runtime-core/common/src/aws/smithy/kotlin/runtime/io/SdkSource.kt +++ b/runtime/runtime-core/common/src/aws/smithy/kotlin/runtime/io/SdkSource.kt @@ -110,3 +110,22 @@ public fun SdkSource.readFully(sink: SdkBuffer, byteCount: Long) { totalBytesRead += rc } } + +/** + * **Caution** Reads the entire contents of the source into [sink]. + * This function will read until the source is exhausted and no bytes remain + * + * @param sink the buffer that data read from the source will be appended to + */ +@InternalApi +public fun SdkSource.readRemaining(sink: SdkBuffer): Long { + var totalBytesRead: Long = 0 + var bytesRead = read(sink, Long.MAX_VALUE) + + while (bytesRead != -1L) { + totalBytesRead += bytesRead + bytesRead = read(sink, Long.MAX_VALUE) + } + + return totalBytesRead +} diff --git a/runtime/runtime-core/common/test/aws/smithy/kotlin/runtime/io/SdkSourceTest.kt b/runtime/runtime-core/common/test/aws/smithy/kotlin/runtime/io/SdkSourceTest.kt new file mode 100644 index 0000000000..14146676f7 --- /dev/null +++ b/runtime/runtime-core/common/test/aws/smithy/kotlin/runtime/io/SdkSourceTest.kt @@ -0,0 +1,57 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package aws.smithy.kotlin.runtime.io + +import kotlin.test.Test +import kotlin.test.assertEquals + +class SdkSourceTest { + @Test + fun readRemaining() { + val data = "Hello world" + val dataLength = data.length.toLong() + val readCycles = 100 + val totalDataLength = dataLength * readCycles + + // Manual and readRemaining + val source = createTestSource(data, dataLength, readCycles) + val buffer = SdkBuffer() + val manualReads = 10 + repeat(manualReads) { + source.read(buffer, dataLength) + } + var readByReadRemaining = source.readRemaining(buffer) + assertEquals(readByReadRemaining, totalDataLength - manualReads * dataLength) + assertEquals(buffer.size, totalDataLength) + + // Only readRemaining + buffer.skip(totalDataLength) + readByReadRemaining = createTestSource(data, dataLength, readCycles).readRemaining(buffer) + assertEquals(readByReadRemaining, totalDataLength) + assertEquals(buffer.size, totalDataLength) + } +} + +private fun createTestSource( + data: String, + dataLength: Long, + readCycles: Int, +) = + object : SdkSource { + var remainingReadCycles = readCycles + + override fun read(sink: SdkBuffer, limit: Long): Long { + if (remainingReadCycles == 0) { + return -1L + } + + sink.writeUtf8(data) + remainingReadCycles-- + return dataLength + } + + override fun close() {} + }