add "bufferCapacity" as a parameter

Author: martinbonnin

build.gradle

@@ -16,7 +16,7 @@ import com.google.firebase.gradle.plugins.license.LicenseResolverPlugin
 import com.google.firebase.gradle.MultiProjectReleasePlugin
 
 buildscript {
-    ext.kotlinVersion = '1.3.72'
+    ext.kotlinVersion = '1.4.32'
     repositories {
         google()
         mavenCentral()

firebase-firestore/ktx/ktx.gradle

@@ -51,7 +51,7 @@ dependencies {
     implementation project(':firebase-common:ktx')
     implementation project(':firebase-firestore')
     implementation 'androidx.annotation:annotation:1.1.0'
-    implementation 'org.jetbrains.kotlinx:kotlinx-coroutines-core:1.3.3'
+    implementation 'org.jetbrains.kotlinx:kotlinx-coroutines-core:1.6.1'
     implementation 'com.google.android.gms:play-services-base:17.0.0'
     testImplementation project(':firebase-database-collection')
     testImplementation 'org.mockito:mockito-core:2.25.0'

firebase-firestore/ktx/src/main/kotlin/com/google/firebase/firestore/ktx/Firestore.kt

@@ -14,27 +14,21 @@
 
 package com.google.firebase.firestore.ktx
 
+import android.support.multidex.BuildConfig
 import androidx.annotation.Keep
 import com.google.firebase.FirebaseApp
 import com.google.firebase.components.Component
 import com.google.firebase.components.ComponentRegistrar
-import com.google.firebase.firestore.DocumentSnapshot
-import com.google.firebase.firestore.FieldPath
-import com.google.firebase.firestore.FirebaseFirestore
-import com.google.firebase.firestore.FirebaseFirestoreSettings
-import com.google.firebase.firestore.QueryDocumentSnapshot
-import com.google.firebase.firestore.QuerySnapshot
-import com.google.firebase.firestore.FirebaseFirestoreSettings
-import com.google.firebase.firestore.MetadataChanges
-import com.google.firebase.firestore.ListenerRegistration
+import com.google.firebase.firestore.*
 import com.google.firebase.ktx.Firebase
 import com.google.firebase.platforminfo.LibraryVersionComponent
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.NonCancellable
+import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.buffer
+import kotlinx.coroutines.flow.callbackFlow
 import kotlinx.coroutines.flow.flow
-import kotlinx.coroutines.withContext
 
 /** Returns the [FirebaseFirestore] instance of the default [FirebaseApp]. */
 val Firebase.firestore: FirebaseFirestore
@@ -173,37 +167,69 @@ class FirebaseFirestoreKtxRegistrar : ComponentRegistrar {
 }
 
 /**
- * Attach a snapshotListener to a DocumentReference and use it as a coroutine flow
- * @param metadataChanges Indicates whether metadata-only changes
+ * Transforms a [DocumentReference] into a coroutine [Flow]
+ *
+ * **Backpressure handling**: by default this method conflates items. If the consumer isn't fast enough,
+ * it might miss some values but is always guaranteed to get the latest value. Use [bufferCapacity]
+ * to change that behaviour
+ *
+ * @param metadataChanges controls metadata-only changes. Default: [MetadataChanges.EXCLUDE]
+ * @param bufferCapacity the buffer capacity as in [Flow.buffer] or null to not buffer at all
  */
-fun DocumentReference.toFlow(metadataChanges: MetadataChanges = MetadataChanges.EXCLUDE) =
-    callbackFlow {
+fun DocumentReference.toFlow(
+    metadataChanges: MetadataChanges = MetadataChanges.EXCLUDE,
+    bufferCapacity: Int? = Channel.CONFLATED
+): Flow<DocumentSnapshot?> {
+    val flow = callbackFlow {
         val registration = addSnapshotListener(metadataChanges) { snapshot, exception ->
             if (exception != null) {
-                cancel(message = "Error getting DocumentReference snapshot", cause = exception)
+                cancel(CancellationException("Error getting DocumentReference snapshot", exception))
             }
 
             if (snapshot != null) {
                 trySend(snapshot)
             }
         }
         awaitClose { registration.remove() }
-    }.buffer(Channel.CONFLATED)
+    }
+
+    return if (bufferCapacity != null) {
+        flow.buffer(bufferCapacity)
+    } else {
+        flow
+    }
+}
 
 /**
- * Attach a snapshotListener to a Query and use it as a coroutine flow
- * @param metadataChanges Indicates whether metadata-only changes
+ * Transforms a [Query] into a coroutine [Flow]
+ *
+ * **Backpressure handling**: by default this method conflates items. If the consumer isn't fast enough,
+ * it might miss some values but is always guaranteed to get the latest value. Use [bufferCapacity]
+ * to change that behaviour
+ *
+ * @param metadataChanges controls metadata-only changes. Default: [MetadataChanges.EXCLUDE]
+ * @param bufferCapacity the buffer capacity as in [Flow.buffer] or null to not buffer at all
  */
-fun Query.toFlow(metadataChanges: MetadataChanges = MetadataChanges.EXCLUDE) =
-    callbackFlow {
+fun Query.toFlow(
+    metadataChanges: MetadataChanges = MetadataChanges.EXCLUDE,
+    bufferCapacity: Int? = Channel.CONFLATED
+): Flow<QuerySnapshot?> {
+    val flow = callbackFlow {
         val registration = addSnapshotListener(metadataChanges) { snapshot, exception ->
             if (exception != null) {
-                cancel(message = "Error getting Query snapshot", cause = exception)
+                cancel(CancellationException("Error getting Query snapshot", exception))
             }
 
             if (snapshot != null) {
                 trySend(snapshot)
             }
         }
         awaitClose { registration.remove() }
-    }.buffer(Channel.CONFLATED)
+    }
+
+    return if (bufferCapacity != null) {
+        flow.buffer(bufferCapacity)
+    } else {
+        flow
+    }
+}