Implement Observale.amb with Kotlin-coroutines?

user_2154011 0
Viewed 129

For instance I have 3 data source from different end-points. I'd like to call all of them in parallel and get the first answer (fastest), then other calls should be discarded. I know how do it using RxJava with Observable.amb(). How to implement it using Kotlin coroutines? The important thing is - no need to wait other calls after the first result.

suspend fun dataSourceOne(){
   delay(1_000L)
}

suspend fun dataSourceTwo(){
   delay(2_000L)
}

suspend fun dataSourceThree(){
   delay(3_000L)
}

// should call [dataSourceOne(), dataSourceTwo(), dataSourceThree()] in parallel 
// and discard [dataSourceTwo(), dataSourceThree()] after the getting a result from dataSourceOne()

PS: Android application.

rx-java2 kotlin-coroutines android
2 Answers

You can launch collectors, use an atomic integer to index the winner and keep track of all other jobs to cancel. For example:

import kotlinx.coroutines.*
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.FlowCollector
import kotlinx.coroutines.flow.collect
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.atomic.AtomicInteger

@FlowPreview
class FlowAmbIterable<T>(private val sources: Iterable<Flow<T>>) : Flow<T> {
    @InternalCoroutinesApi
    override suspend fun collect(collector: FlowCollector<T>) {
        val winner = AtomicInteger()
        val jobs = ConcurrentHashMap<Job, Int>()
        coroutineScope {
            var i = 1
            for (source in sources) {
                val idx = i
                val job = launch {
                    source.collect {
                        val w = winner.get()
                        if (w == idx) {
                            collector.emit(it)
                        } else if (w == 0 && winner.compareAndSet(0, idx)) {
                            for (j in jobs.entries) {
                                if (j.value != idx) {
                                    j.key.cancel()
                                }
                            }

                            collector.emit(it)
                        } else {
                            throw CancellationException()
                        }
                    }
                }

                jobs[job] = i
                val w = winner.get()
                if (w != 0 && w != i) {
                    job.cancel()
                    break
                }

                i++
            }
        }
    }
}
user_611580

References

import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.channels.ChannelResult
import kotlinx.coroutines.channels.onFailure
import kotlinx.coroutines.channels.onSuccess
import kotlinx.coroutines.channels.produce
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.emitAll
import kotlinx.coroutines.flow.flow
import kotlinx.coroutines.selects.select
import kotlinx.coroutines.yield

@ExperimentalCoroutinesApi
public fun <T> race(flows: Iterable<Flow<T>>): Flow<T> = flow {
  coroutineScope {
    // 1. Collect to all source Flows
    val channels = flows.map { flow ->
      // Produce the values using the default (rendezvous) channel
      produce {
        flow.collect {
          send(it)
          yield() // Emulate fairness, giving each flow chance to emit
        }
      }
    }

    // If channels List is empty, just return and complete result Flow.
    if (channels.isEmpty()) {
      return@coroutineScope
    }
    
    // If channels List has single element, just forward all events from it.
    channels
      .singleOrNull()
      ?.let { return@coroutineScope emitAll(it) }

    // 2. When a new event arrives from a source Flow, pass it down to a collector.
    // Select expression makes it possible to await multiple suspending functions simultaneously
    // and select the first one that becomes available.
    val (winnerIndex, winnerResult) = select<Pair<Int, ChannelResult<T>>> {
      channels.forEachIndexed { index, channel ->
        channel.onReceiveCatching {
          index to it
        }
      }
    }

    // 3. Cancel all other Flows.
    channels.forEachIndexed { index, channel ->
      if (index != winnerIndex) {
        channel.cancel()
      }
    }

    // 4. Forward all events from the winner Flow .
    winnerResult
      .onSuccess {
        emit(it)
        emitAll(channels[winnerIndex])
      }
      .onFailure {
        it?.let { throw it }
      }
  }
}
user_111914240
Related