How to propagate closing to a chain of flows in kotlin

Viewed 29

I am using kotlin and I wanted to stream over a possibly huge resultset using flows. I found some explanations around the web:

I implemented it and it works fine. I also needed to batch the results before sending them to an external services, so I implemented a chunked operation on flows. Something like that:

fun <T> Flow<T>.chunked(chunkSize: Int): Flow<List<T>> {
  return callbackFlow {
    val listOfResult = mutableListOf<T>()
    this@chunked.collect {
      listOfResult.add(it)
      if (listOfResult.size == chunkSize) {
        trySendBlocking(listOfResult.toList())
        listOfResult.clear()
      }
    }
    if (listOfResult.isNotEmpty()) {
      trySendBlocking(listOfResult)
    }
    close()
  }
}

To be sure that everything was working fine, I created some integration tests:

  • first flow + chuncked to consume all rows, passed
  • using the first flow (the one created from the jdbc repository) and applying take operator just to consider few x items. It passed correctly.
  • using first flow + chunked operator + take operator, it hangs forever

So the last test showed that there was something wrong in the implementation. I investigated a lot without finding nothing useful but, dumping the threads, I found a coroutine thread blocked in the trySendBlocking call on the first flow, the one created in the jdbc repository. I am wondering in which way the chunked operator is supposed to propagate the closing to the upstream flow since it seems this part is missing. In both cases I am propagating downstream the end of data with a close() call but I took a look the take operator and I saw it is triggering back the closing with an emitAbort(...) Should I do something similar in the callbackFlow{...}? After a bit of investigation, I was able to avoid the locking adding a timeout on the trySendBlocking inside the repository but I didn´t like that. At the end, I realized that I could cast the original flow (in the chunked operator) to a SendChannel and close it if the downstream flow is closed:

 trySendBlocking(listOfResult.toList()).onSuccess {
    LOGGER.debug("Sent")
  }.onFailure {
    LOGGER.warn("An error occurred sending data.", it)
  }.onClosed {
    LOGGER.info("Channel has been closed")
    (originalFlow as SendChannel<*>).close(it)
  }

Is this the correct way of closing flows backwards? Any hint to solve this issue? Thanks!

1 Answers

You shouldn't use trySendBlocking instead of send. You should never use a blocking function in a coroutine without wrapping it in withContext with a Dispatcher that can handle blocking code (e.g. Dispatchers.Default). But when there's a suspend function alternative, use that instead, in this case send().

Also, callbackFlow is more convoluted than necessary for transforming a flow. You should use the standard flow builder instead (and so you'll use emit() instead of send()).

fun <T> Flow<T>.chunked(chunkSize: Int): Flow<List<T>> = flow {
    val listOfResult = mutableListOf<T>()
    collect {
        listOfResult.add(it)
        if (listOfResult.size == chunkSize) {
            emit(listOfResult.toList())
            listOfResult.clear()
        }
    }
    if (listOfResult.isNotEmpty()) {
        emit(listOfResult)
    }
}
Related