Iterative consumption of flows

[lowasser]

This is the product of some library design work we've been doing at Google, and we wanted to take this upstream in case there was interest. The Flow API is almost exclusively "push"-based, not "pull" based. The main alternative approach is to adapt it to a Channel, but that often introduces inconvenient concurrency issues (arbitrary buffering, the risk of forgetting to cancel the channel).

What we want is more or less a scoped iteration method. Here is a sketch of some examples that we think would be better with the proposed API than what Flow can do today.

suspend fun <T> Flow<T>.all(predicate: suspend (T) -> Boolean): Boolean = iterate {
  while (hasNext()) {
    if (!predicate(next()) return@iterate false
  }
  true
}

fun <T> Flow<T>.pairElements(): Flow<Pair<T, T>> = flow {
  iterate {
    while (hasNext()) {
      val a = next()
      if (!hasNext()) break
      val b = next()
      emit(Pair(a, b))
    }
  }
}

// this API sort of generalizes e.g. https://github.com/cashapp/turbine
fun testStreamingRpcResponse() = runTest {
  val responses = stub.doRpc(request)
  responses.iterate {
    assertThat(next()).isEqualTo(response1)
    doStimulus(1)
    assertThat(next()).isEqualTo(response2)
    doStimulus(2)
    assertThat(hasNext()).isFalse()
  }
}

In short, what we're proposing is

suspend fun <T, R> Flow<T>.iterate(block: ??Iterator<T>.() -> R): R

It's not precisely clear what the ??Iterator should be: we might reuse ChannelIterator, but we're not particularly fond of the requirement to call hasNext() in the test case especially, and of course this isn't actually a Channel. For ourselves, we are inclined to introduce a separate FlowIterator type.

We have an implementation that does not introduce undue concurrency. To be a little more specific, it does not continue the backing Flow.collect until the the next call to hasNext, etc.; it's a linear transformation in the same way as map and filter but not buffer or conflate. We can PR it if there's interest.

[lowasser]

Continuing conversation from #3278 (comment).

We don't mind if it's slow

I've mentioned this in other contexts, but perhaps not when you were present, so let me say this as the coroutines lead on Google's internal Kotlin team: we consider performance to be our lowest priority for coroutines, behind correctness/ease of writing correct code, clear API design, and testability. We think this API has advantages for all three of those.

We're not using flows (or coroutines) for the performance, we're using them because they're a clear and clean API. The 30x slowdown you estimate is a little on the high side for us, but only a little; we would be excited to use this API even if it represented a 30x slowdown, even in production, even if it's only to make code easier to follow. Coroutines just aren't a bottleneck, or even within 30x of being a bottleneck; outages and bugs are much more expensive.

Test use cases

The canonical example of the test case looks something like this. Here's a very simple test:

// with Truth, other assertion libraries will presumably be similar
assertThat(stub.rpc(flowOf(1,2,3)).toList()).containsExactly(2, 4, 5, 6).inOrder()

This is great. But what if we want to assert that the RPC doesn't wait until the requests flow completes, but emits responses as we go?

val requests = SharedFlow()
stub.rpc(requests.take(3)).iterate {
  // this assumes that hasNext is not required as with ChannelIterator.
  requests.emit(1)
  assertThat(next()).isEqualTo(2)
  requests.emit(2)
  assertThat(next()).isEqualTo(4)
  assertThat(next()).isEqualTo(5)
  requests.emit(3)
  assertThat(next()).isEqualTo(6)
  assertThat(hasNext()).isFalse()
}

This makes different assertions at each request, which is a poor match for push-based APIs. I can imagine several push-based approaches to this problem, but they all separate stimulus and response, making the test harder to write, read, and debug. Most streaming RPCs need tests like this, and we have a lot of those.

https://github.com/cashapp/turbine is a preexisting library by @JakeWharton that addressed the specific use case of testing in this style; its README also discusses use cases in more detail. We've been hesitant to fully adopt Turbine because it dallies with UNDISPATCHED, which we've seen lead to impossible-to-debug issues, but its API is great for testing flows -- but a more general operator applicable to production would be even better. I'm sure Jake could comment more, but he did thumbs-up the proposal above.

Flow operators and production use cases

For the production use case, pairElements from the original comment is a simplified version of our real use case, but it gets across the idea that we want to express. The pull-based style looks like

fun Flow<T>.pairElements() = flow {
  var even = true
  lateinit var prev: T
  collect {
    if (even) {
      prev = it
    } else {
      emit(Pair(prev, it))
    }
    even = !even
  }
}

This is a very simple state machine, but it's much clearer to write

fun Flow<T>.pairElements() = flow {
  iterate {
    while (hasNext()) {
      emit(Pair(next(), next()))
    }
  }
}

Coroutines are all about capturing state machines from program control flow, and they do a great job here.

The use case being simplified is parsing streams of bytes from a Flow<ByteString> in the delimited proto format, in which the byte stream contains alternating varint-encoded message lengths, then messages of that length. These can be split arbitrarily across ByteStrings, which are simply chunks of unknown length read from files (in the cloud, via RPCs, thus the Flow). In this case, the algorithm is "consume byte strings until we have a length, parse it, consume byte strings until we have at least that length, emit that chunk, repeat until the stream ends." We started out with a push-based style to implement that, sure, but it was complex and hard to follow, because it had to manually manage a state machine that was much clearer in the pull style. Parsing the messages and juggling the RPC is much more expensive than the coroutine parts, so who cares if we slow down the flow consumption?

Tests for this use case also ran the hard way into the race condition in channel-based solutions, parsing byte streams interrupted by exceptions.

Debugging

The desirable state when debugging is that, at a breakpoint at B(X), A is suspended at the point where it just emitted X and has not done any more work.

It looks like it has to be solved not on the operator level, but at the infrastructure/testing one. E.g. (optional) injectable dispatchers to the production source code and sequential testing primitives: kotlinx-coroutines-test or turbine.

A clever dispatcher doesn't solve this problem, though! #3278 (comment) shows the control flow diagrams that are the root of the issue, which is that two separate coroutines are both eligible to run at the same point in the program, as far as any dispatcher is concerned, but we don't want one of them to run (yet) when debugging -- and the dispatcher API means that it can't possibly know which one we want. A fully sequential, single-threaded dispatcher prevents "races" in the sense of JVM races, but doesn't solve the "race" in the sense that the dispatcher chooses which of two eligible coroutines get resumed first, and debugging requires the dispatcher to make a specific choice there that a dispatcher can't guess.

The fix is to prevent the race, where the only coroutine eligible to run is the right one to step forward, and that means something more like iterate that doesn't use channels (in the obvious way).

[lowasser]

(To clarify, in the "parsing streams of bytes" use case, we are returning a Flow of the parsed messages.)

[qwwdfsad]

Hi, I'm still processing your input and trying to internalize all the things in order to further weigh the pros and cons of the solution. Let's discuss individual details in order to be on the same page.

We don't mind if it's slow
...
we would be excited to use this API even if it represented a 30x slowdown, even in production, even if it's only to make code easier to follow

It's good to hear that, meaning that we have a bit more space in our design. But it's important to remember that we have plenty of other consumers, including, but not limited to, Android developers, server-side developers, high-performance web-frameworks, high-throughput backend apps and IDEs. Having two similar APIs that, at the first glance, differs only in code-style preferences, but one being 30 times slower, will bring irreversible damage to the ecosystem in the long term, especially when both pull and push based APIs are equally appealing.

suspend fun <T> Flow<T>.all(predicate: suspend (T) -> Boolean): Boolean = iterate {
  while (hasNext()) {
    if (!predicate(next()) return@iterate false
  }
  true
}

This is an example of what can bring potential harm. One can implement it via collectWhile using a pull-based model, and one can implement it using iterate, but being an order of magnitude slower. While being acceptable in some cases, it eventually will separate the code style, will bring the debates over code-style and a lot of articles "Do and dont's for performance" (and the veil of mystery around the mythical abstract "performance" won't make it any easier for us).

#3278

Would it be okay for you to semi-finish the PR? It is a decent first peek at the API shape, but unfortunately, it doesn't compile and doesn't have any tests. It would be really nice to see API and tests that you already have written to give us a taste of how you are leveraging it. There are some examples, but they have their own pitfalls. Real-world use-cases are always better :)

This is a very simple state machine, but it's much clearer to write

fun Flow<T>.pairElements() = flow {
  iterate {
    while (hasNext()) {
      emit(Pair(next(), next()))
    }
  }
}

This particular example looks really concise, works and passes basic testing. As long as the original flow emits an even number of elements, otherwise, it throws NoSuchElementException.
Of course, channels and collection's iterators have the same pitfall, but their use pattern is completely different -- their iterators are almost never expected to be used directly, but rather with for-loop convention, consume extension or stdlib primitives. Here we explicitly delegate this responsibility to end-users, potentially creating the space for hard-to-spot bugs that type system cannot prevent.

A fully sequential, single-threaded dispatcher prevents "races" in the sense of JVM races, but doesn't solve the "race" in the sense that the dispatcher chooses which of two eligible coroutines get resumed first, and debugging requires the dispatcher to make a specific choice there that a dispatcher can't guess.

Thanks for the control flow diagram, it really made the picture clear. I've immediately started thinking about more intrusive UNDISPATCHED primitive for this purpose and noticed your comment

We've been hesitant to fully adopt Turbine because it dallies with UNDISPATCHED, which we've seen lead to impossible-to-debug issues

Could you please elaborate on that and show a few examples of such issues?

I'm looking forward to seeing more examples (I suppose the parsing is the most interesting and hard-to-emulate one) to see if there are ways to improve the API and get around potential drawbacks

[JakeWharton]

Do you mean Unconfined? Turbine's use of UNDISPATCHED merely starts the upstream flow collection synchronously.

[lowasser]

I don't have specific issues I can refer to about Unconfined or UNDISPATCHED, though both make us extremely nervous. Certainly in Java land we've had lots of issues caused by directExecutor. We specifically haven't seen any such issues with Turbine, but then, we have it on a tight allowlist just because we don't trust that there aren't issues with Unconfined and UNDISPATCHED too subtle for us to notice. Similarly, we would be unhappy about more intrusive primitives and their edge cases.

We care about performance, to be clear, but we think race conditions in common patterns that are subtle enough that they're not obvious to you or Roman at first glance are pretty damaging to the ecosystem themselves.

I've committed a working version with some of our tests to the PR, though it doesn't have as many tests as I'd like yet. I've also modified it to avoid the use of channels, instead just passing continuations back and forth. I'm less confident in its correctness, though.

(Possibly an advantage of the API in the current PR: it can't actually be for-eached over, which makes it more likely to be used in iterator-specific ways?)

[JakeWharton]

Turbines use of unconfined is an optimization to avoid double dispatch similar to one of your goals here. It isn't strictly required, and runTest seems to have broken its behavior so probably worth removing.

[jingibus]

It's not just Turbine-type interactions that have need to "pull" consumption. We also need it for interaction with fakes.

Internally at Cash, we've been attacking this problem by using a "pull" consumption tool that Kotlin already has: Channel.

We facade Channel's API with one better suited to our test requirements. Here's a simplified sketch of that API:

class AwaitChannel<T> {
  private val channel = Channel<T>(capacity = UNLIMITED)

  fun add(value: T) { channel.trySend(value) }

  val isEmpty: Boolean get() { return channel.isEmpty }

  /**
   * Wait for a value on this channel for 1s, throwing a [CancellationException] if none arrives.
   */
  suspend fun <T> awaitValue() = try {
    withTimeout(1000L) {
      channel.receive()
    }
  }

  /**
   * Return a value, or null if empty. 
   *
   * To be used from legacy, non-suspendful call sites. Will throw
   * when called from a non-suspending context.
   */
  fun <T> takeValue(): T? {
    assertCallingContextIsNotSuspended()
    return channel.tryReceive().getOrNull()
  }

This is essentially the same thing Turbine provides, but in a form that works in our Fakes.

I think the "it's just a channel" idea can work in Turbine, too. I did some homework on this last week, and haven't yet found a major roadblock.

[lowasser]

Does that approach address the race issue discussed in the PR here?

[jingibus]

We have only been concerned with races in test. Test has firm expectations of consistency: many immaterial changes in order may occur in multithreaded prod without harm, but immaterial changes in ordering under test will result in flaky test failures.

Since races are introduced by threading, not coroutines, we police dispatchers aggressively:

• All dispatchers are injected as EmptyCoroutineContext in test
• All Rx schedulers are injected as Schedulers.trampoline() in test
• We stay on guard against problematic Kotlin APIs that silently introduce Dispatchers.Default (e.g. rx(Single|Maybe|Completable|Observable), dispatcherless CoroutineScope())

[lowasser]

There are plenty of races that can happen as a result of coroutines and not threading! #3278 (comment) is an example of code that nondeterministically throws or succeeds, even run on a single-threaded dispatcher.