AsyncChannelExt.kt

import java.net.SocketAddress
import java.nio.ByteBuffer
import java.nio.channels.AsynchronousByteChannel
import java.nio.channels.AsynchronousSocketChannel
import java.nio.channels.CompletionHandler
import kotlin.coroutines.Continuation
import kotlin.coroutines.resume
import kotlin.coroutines.resumeWithException
import kotlinx.coroutines.suspendCancellableCoroutine

class Completer<V> : CompletionHandler<V, Continuation<V>> {
  override fun completed(result: V, cont: Continuation<V>) = cont.resume(result)
  override fun failed(exc: Throwable, cont: Continuation<V>) = cont.resumeWithException(exc)
}

private val voidCompleter = Completer<Void>()
private val intCompleter = Completer<Int>()

suspend fun AsynchronousSocketChannel.connectTo(remote: SocketAddress) {
  suspendCancellableCoroutine { continuation ->
    continuation.invokeOnCancellation { close() }
    runCatching { connect(remote, continuation, voidCompleter) }
      .onFailure { continuation.resumeWithException(it) }
  }
}

suspend fun AsynchronousByteChannel.readInto(buffer: ByteBuffer): Int {
  return suspendCancellableCoroutine { continuation ->
    continuation.invokeOnCancellation { close() }
    read(buffer, continuation, intCompleter)
  }
}

suspend fun AsynchronousByteChannel.writeFrom(buffer: ByteBuffer): Int {
  return suspendCancellableCoroutine { continuation ->
    continuation.invokeOnCancellation { close() }
    write(buffer, continuation, intCompleter)
  }
}

suspend fun AsynchronousByteChannel.readInto(buffer: ByteBuffer, maxBytes: Int): Int {
  var totalBytes = 0
  var bytesRead: Int
  buffer.clear()
  buffer.limit(maxBytes)
  do {
    bytesRead = readInto(buffer)
    if (bytesRead < 0) {
      break
    }
    totalBytes += bytesRead
  } while (totalBytes < maxBytes)
  buffer.flip()
  return totalBytes
}