dev.pellet.server.extension.NIO2CoroutineExtensions.kt Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of pellet-server Show documentation

An opinionated Kotlin web framework, with best-practices built-in

There is a newer version: 0.0.16

package dev.pellet.server.extension

import dev.pellet.server.buffer.PelletBuffer
import kotlinx.coroutines.CancellableContinuation
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.suspendCancellableCoroutine
import java.nio.ByteBuffer
import java.nio.channels.AsynchronousChannel
import java.nio.channels.AsynchronousCloseException
import java.nio.channels.AsynchronousServerSocketChannel
import java.nio.channels.AsynchronousSocketChannel
import java.nio.channels.CompletionHandler
import kotlin.coroutines.resumeWithException

// Inspired by https://github.com/Kotlin/kotlinx.coroutines/blob/87eaba8a287285d4c47f84c91df7671fcb58271f/integration/kotlinx-coroutines-nio/src/Nio.kt

suspend fun AsynchronousServerSocketChannel.awaitAccept(): AsynchronousSocketChannel {
    return suspendCancellableCoroutine { continuation ->
        this.accept(continuation, anyAsyncContinuationHandler())
        closeOnCancellation(continuation)
    }
}

suspend fun AsynchronousSocketChannel.awaitRead(
    buffer: PelletBuffer
): Int {
    return suspendCancellableCoroutine { continuation ->
        this.read(buffer.byteBuffer, continuation, anyAsyncContinuationHandler())
        closeOnCancellation(continuation)
    }
}

suspend fun AsynchronousSocketChannel.awaitWrite(
    buffer: ByteBuffer
): Int {
    return suspendCancellableCoroutine { continuation ->
        this.write(buffer, continuation, anyAsyncContinuationHandler())
        closeOnCancellation(continuation)
    }
}

private fun AsynchronousChannel.closeOnCancellation(
    continuation: CancellableContinuation<*>
) {
    continuation.invokeOnCancellation {
        // intentionally ignore exceptions on close
        runCatching {
            close()
        }
    }
}

@Suppress("UNCHECKED_CAST")
fun  anyAsyncContinuationHandler(): CompletionHandler> =
    AnyAsyncContinuationHandler as CompletionHandler>

object AnyAsyncContinuationHandler : CompletionHandler> {

    override fun completed(
        result: Any,
        attachment: CancellableContinuation
    ) {
        @OptIn(ExperimentalCoroutinesApi::class)
        attachment.resume(result, null)
    }

    override fun failed(
        exception: Throwable,
        attachment: CancellableContinuation
    ) {
        if (exception is AsynchronousCloseException && attachment.isCancelled) {
            return
        }
        attachment.resumeWithException(exception)
    }
}