Arrow Fx Coroutines CompletableFuture

CompletableFuture.kt

import arrow.fx.coroutines.CancelToken
import arrow.fx.coroutines.SuspendConnection
import arrow.fx.coroutines.cancellable
import java.util.concurrent.CancellationException
import java.util.concurrent.CompletableFuture
import java.util.concurrent.CompletionException
import java.util.concurrent.CompletionStage
import java.util.concurrent.ExecutionException
import java.util.concurrent.Future
import kotlin.coroutines.Continuation
import kotlin.coroutines.CoroutineContext
import kotlin.coroutines.EmptyCoroutineContext
import kotlin.coroutines.startCoroutine

fun <T> future(
  context: CoroutineContext = EmptyCoroutineContext,
  block: suspend () -> T
): CompletableFuture<T> {
  val future = CompletableFuture<T>()
  val conn = SuspendConnection()

  future.whenComplete { _, _ ->
    // If the future completed externally, we should cancel the `block`
    conn::cancel.startCoroutine(Continuation(EmptyCoroutineContext) { })
  }

  block.startCoroutine(Continuation(context + conn) {
    it.fold(future::complete, future::completeExceptionally)
  })

  return future
}

suspend fun <A> CompletionStage<A>.await(): A =
  if (this is Future<*> && isDone()) {
    try {
      get() as A
    } catch (e: ExecutionException) {
      throw e.cause ?: e // unwrap original cause from ExecutionException
    }
  } else cancellable { callback ->


    whenComplete { res, exception ->
      if (exception == null) callback(Result.success(res))
      else when (exception) {
        is CancellationException -> Unit
        else -> callback(Result.failure((exception as? CompletionException)?.cause ?: exception))
      }
    }

    CancelToken { (this as? CompletableFuture<A>)?.cancel(false) }
  }