commonMain.arrow.fx.coroutines.Race3.kt Maven / Gradle / Ivy

Go to download
@file:OptIn(ExperimentalContracts::class)

package arrow.fx.coroutines

import arrow.core.nonFatalOrThrow
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Deferred
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.async
import kotlinx.coroutines.cancelAndJoin
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.selects.select
import kotlin.contracts.ExperimentalContracts
import kotlin.contracts.InvocationKind
import kotlin.contracts.contract
import kotlin.coroutines.ContinuationInterceptor
import kotlin.coroutines.CoroutineContext
import kotlin.coroutines.EmptyCoroutineContext

public sealed class Race3 {
  public data class First(val winner: A) : Race3()
  public data class Second(val winner: B) : Race3()
  public data class Third(val winner: C) : Race3()

  public inline fun  fold(
    ifA: (A) -> D,
    ifB: (B) -> D,
    ifC: (C) -> D
  ): D {
    contract {
      callsInPlace(ifA, InvocationKind.AT_MOST_ONCE)
      callsInPlace(ifB, InvocationKind.AT_MOST_ONCE)
      callsInPlace(ifC, InvocationKind.AT_MOST_ONCE)
    }
    return when (this) {
      is First -> ifA(winner)
      is Second -> ifB(winner)
      is Third -> ifC(winner)
    }
  }
}

/**
 * Races the participants [fa], [fb] & [fc] in parallel on the [Dispatchers.Default].
 * The winner of the race cancels the other participants.
 * Cancelling the operation cancels all participants.
 *
 * @see raceN for the same function that can race on any [CoroutineContext].
 */
public suspend inline fun  raceN(
  crossinline fa: suspend CoroutineScope.() -> A,
  crossinline fb: suspend CoroutineScope.() -> B,
  crossinline fc: suspend CoroutineScope.() -> C
): Race3 {
  contract {
    callsInPlace(fa, InvocationKind.AT_MOST_ONCE)
    callsInPlace(fb, InvocationKind.AT_MOST_ONCE)
    callsInPlace(fc, InvocationKind.AT_MOST_ONCE)
  }
  return raceN(Dispatchers.Default, fa, fb, fc)
}

/**
 * Races the participants [fa], [fb] & [fc] on the provided [CoroutineContext].
 * The winner of the race cancels the other participants.
 * Cancelling the operation cancels all participants.
 *
 * Coroutine context is inherited from a [CoroutineScope], additional context elements can be specified with [ctx] argument.
 * If the combined context does not have any dispatcher nor any other [ContinuationInterceptor], then [Dispatchers.Default] is used.
 * **WARNING** If the combined context has a single threaded [ContinuationInterceptor], this function will not run [fa], [fb] & [fc] in parallel.
 *
 * @see raceN for a function that ensures operations run in parallel on the [Dispatchers.Default].
 */
@Suppress("LEAKED_IN_PLACE_LAMBDA")
public suspend inline fun  raceN(
  ctx: CoroutineContext = EmptyCoroutineContext,
  crossinline fa: suspend CoroutineScope.() -> A,
  crossinline fb: suspend CoroutineScope.() -> B,
  crossinline fc: suspend CoroutineScope.() -> C
): Race3 {
  contract {
    callsInPlace(fa, InvocationKind.AT_MOST_ONCE)
    callsInPlace(fb, InvocationKind.AT_MOST_ONCE)
    callsInPlace(fc, InvocationKind.AT_MOST_ONCE)
  }
  return coroutineScope {
    val a = async(ctx) { fa() }
    val b = async(ctx) { fb() }
    val c = async(ctx) { fc() }
    select> {
      a.onAwait.invoke { Race3.First(it) }
      b.onAwait.invoke { Race3.Second(it) }
      c.onAwait.invoke { Race3.Third(it) }
    }.also {
      when (it) {
        is Race3.First -> cancelAndCompose(b, c)
        is Race3.Second -> cancelAndCompose(a, c)
        is Race3.Third -> cancelAndCompose(a, b)
      }
    }
  }
}

@PublishedApi
internal suspend fun cancelAndCompose(first: Deferred<*>, second: Deferred<*>): Unit {
  val e1 = try {
    first.cancelAndJoin()
    null
  } catch (e: Throwable) {
    e.nonFatalOrThrow()
  }
  val e2 = try {
    second.cancelAndJoin()
    null
  } catch (e: Throwable) {
    e.nonFatalOrThrow()
  }
  (e1?.apply { e2?.let(::addSuppressed) } ?: e2)?.let { throw it }
}