Add min parameter to await

ktorio · May 13, 2024 · 31b0425 · 31b0425
1 parent 6aa8c78
commit 31b0425
Show file tree

Hide file tree

Showing 7 changed files with 26 additions and 18 deletions.
diff --git a/ktor-io/api/ktor-io.api b/ktor-io/api/ktor-io.api
@@ -7,7 +7,7 @@ public final class io/ktor/utils/io/ByteChannel : io/ktor/utils/io/BufferedByteW
 	public fun <init> ()V
 	public fun <init> (Z)V
 	public synthetic fun <init> (ZILkotlin/jvm/internal/DefaultConstructorMarker;)V
-	public fun awaitContent (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public fun awaitContent (ILkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public fun cancel (Ljava/lang/Throwable;)V
 	public fun close ()V
 	public fun flush (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
@@ -43,7 +43,7 @@ public final class io/ktor/utils/io/ByteChannelUtilsKt {
 
 public abstract interface class io/ktor/utils/io/ByteReadChannel {
 	public static final field Companion Lio/ktor/utils/io/ByteReadChannel$Companion;
-	public abstract fun awaitContent (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public abstract fun awaitContent (ILkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public abstract fun cancel (Ljava/lang/Throwable;)V
 	public abstract fun getClosedCause ()Ljava/lang/Throwable;
 	public abstract fun getReadBuffer ()Lkotlinx/io/Source;
@@ -54,6 +54,10 @@ public final class io/ktor/utils/io/ByteReadChannel$Companion {
 	public final fun getEmpty ()Lio/ktor/utils/io/ByteReadChannel;
 }
 
+public final class io/ktor/utils/io/ByteReadChannel$DefaultImpls {
+	public static synthetic fun awaitContent$default (Lio/ktor/utils/io/ByteReadChannel;ILkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
+}
+
 public final class io/ktor/utils/io/ByteReadChannelKt {
 	public static final fun cancel (Lio/ktor/utils/io/ByteReadChannel;)V
 }
@@ -85,6 +89,7 @@ public final class io/ktor/utils/io/ByteReadChannelOperationsKt {
 	public static final fun readShort (Lio/ktor/utils/io/ByteReadChannel;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun readUTF8Line (Lio/ktor/utils/io/ByteReadChannel;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public static final fun readUTF8LineTo (Lio/ktor/utils/io/ByteReadChannel;Ljava/lang/Appendable;ILkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public static synthetic fun readUTF8LineTo$default (Lio/ktor/utils/io/ByteReadChannel;Ljava/lang/Appendable;ILkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object;
 	public static final fun reader (Lkotlinx/coroutines/CoroutineScope;Lkotlin/coroutines/CoroutineContext;Lio/ktor/utils/io/ByteChannel;Lkotlin/jvm/functions/Function2;)Lio/ktor/utils/io/ReaderJob;
 	public static final fun reader (Lkotlinx/coroutines/CoroutineScope;Lkotlin/coroutines/CoroutineContext;ZLkotlin/jvm/functions/Function2;)Lio/ktor/utils/io/ReaderJob;
 	public static synthetic fun reader$default (Lkotlinx/coroutines/CoroutineScope;Lkotlin/coroutines/CoroutineContext;ZLkotlin/jvm/functions/Function2;ILjava/lang/Object;)Lio/ktor/utils/io/ReaderJob;
@@ -173,7 +178,7 @@ public abstract interface class io/ktor/utils/io/ChannelJob {
 
 public final class io/ktor/utils/io/CountedByteReadChannel : io/ktor/utils/io/ByteReadChannel {
 	public fun <init> (Lio/ktor/utils/io/ByteReadChannel;)V
-	public fun awaitContent (Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
+	public fun awaitContent (ILkotlin/coroutines/Continuation;)Ljava/lang/Object;
 	public fun cancel (Ljava/lang/Throwable;)V
 	public fun getClosedCause ()Ljava/lang/Throwable;
 	public final fun getDelegate ()Lio/ktor/utils/io/ByteReadChannel;

diff --git a/ktor-io/common/src/io/ktor/utils/io/ByteChannel.kt b/ktor-io/common/src/io/ktor/utils/io/ByteChannel.kt
@@ -58,11 +58,11 @@ public class ByteChannel(public val autoFlush: Boolean = false) : ByteReadChanne
         get() = (closedCause != null) || (isClosedForWrite && flushBufferSize == 0 && _readBuffer.exhausted())
 
     @OptIn(InternalAPI::class)
-    override suspend fun awaitContent(): Boolean {
+    override suspend fun awaitContent(min: Int): Boolean {
         rethrowCloseCauseIfNeeded()
 
         slot.sleepWhile {
-            flushBufferSize == 0 && _readBuffer.size == 0L && !isClosedForRead
+            flushBufferSize + _readBuffer.size < min && !isClosedForRead
         }
         if (_readBuffer.size < CHANNEL_MAX_SIZE) moveFlushToReadBuffer()
         return _closedCause.value == null

diff --git a/ktor-io/common/src/io/ktor/utils/io/ByteReadChannel.kt b/ktor-io/common/src/io/ktor/utils/io/ByteReadChannel.kt
@@ -22,12 +22,12 @@ public interface ByteReadChannel {
     public val readBuffer: Source
 
     /**
-     * Suspend the channel until it has bytes or gets closed. Throws exception if the channel was closed with an error.
-     * If there are bytes available in the channel, this function returns immediately.
+     * Suspend the channel until it has [min] bytes or gets closed. Throws exception if the channel was closed with an
+     * error. If there are bytes available in the channel, this function returns immediately.
      *
      * @return return `false` eof is reached, otherwise `true`.
      */
-    public suspend fun awaitContent(): Boolean
+    public suspend fun awaitContent(min: Int = 1): Boolean
 
     public fun cancel(cause: Throwable?)
 
@@ -41,7 +41,7 @@ public interface ByteReadChannel {
             @InternalAPI
             override val readBuffer: Source = Buffer()
 
-            override suspend fun awaitContent(): Boolean = false
+            override suspend fun awaitContent(min: Int): Boolean = false
 
             override fun cancel(cause: Throwable?) {
             }

diff --git a/ktor-io/common/src/io/ktor/utils/io/CountedByteReadChannel.kt b/ktor-io/common/src/io/ktor/utils/io/CountedByteReadChannel.kt
@@ -33,9 +33,9 @@ public class CountedByteReadChannel(public val delegate: ByteReadChannel) : Byte
     override val readBuffer: Source
         get() = delegate.readBuffer
 
-    override suspend fun awaitContent(): Boolean {
+    override suspend fun awaitContent(min: Int): Boolean {
         val before = delegate.availableForRead
-        val result = delegate.awaitContent()
+        val result = delegate.awaitContent(min)
         count += (initial - before).toLong()
         initial = delegate.availableForRead
         return result

diff --git a/ktor-io/common/src/io/ktor/utils/io/SourceByteReadChannel.kt b/ktor-io/common/src/io/ktor/utils/io/SourceByteReadChannel.kt
@@ -28,7 +28,7 @@ internal class SourceByteReadChannel(private val source: Source) : ByteReadChann
             return source
         }
 
-    override suspend fun awaitContent(): Boolean {
+    override suspend fun awaitContent(min: Int): Boolean {
         closedCause?.let { throw it }
         return false
     }

diff --git a/ktor-io/jvm/src/io/ktor/utils/io/LookAheadSession.kt b/ktor-io/jvm/src/io/ktor/utils/io/LookAheadSession.kt
@@ -36,7 +36,7 @@ public class LookAheadSuspendSession(private val channel: ByteReadChannel) {
     @OptIn(InternalAPI::class)
     public suspend fun awaitAtLeast(min: Int): Boolean {
         if (channel.readBuffer.remaining >= min) return true
-        channel.awaitContent()
+        channel.awaitContent(min)
         return channel.readBuffer.remaining >= min
     }
 

diff --git a/ktor-io/jvm/src/io/ktor/utils/io/jvm/javaio/Reading.kt b/ktor-io/jvm/src/io/ktor/utils/io/jvm/javaio/Reading.kt
@@ -54,14 +54,17 @@ internal class RawSourceChannel(
     override val readBuffer: Source
         get() = buffer
 
-    override suspend fun awaitContent(): Boolean {
+    override suspend fun awaitContent(min: Int): Boolean {
         if (closedToken != null) return true
 
         withContext(coroutineContext) {
-            val result = try {
-                source.readAtMostTo(buffer, Long.MAX_VALUE)
-            } catch (cause: EOFException) {
-                -1L
+            var result = 0L
+            while (buffer.size < min && result >= 0) {
+                result = try {
+                    source.readAtMostTo(buffer, Long.MAX_VALUE)
+                } catch (cause: EOFException) {
+                    -1L
+                }
             }
 
             if (result == -1L) {