Fix MutationQueue issue resulting in re-sending acknowledged writes. (#909)

Port of: https://github.com/firebase/firebase-js-sdk/pull/559
Should address #772 once released.

getNextMutationBatchAfterBatchId() was not respecting
highestAcknowledgedBatchId and therefore we were resending writes after the
network was disabled / re-enabled.

5 files changed, 220 insertions, 4 deletions

diff --git a/Firestore/CHANGELOG.md b/Firestore/CHANGELOG.md
index 71b61ef..e1ae4dc 100644
--- a/Firestore/CHANGELOG.md
+++ b/Firestore/CHANGELOG.md
@@ -3,6 +3,9 @@
   neither succeeds nor fails within 10 seconds, the SDK will consider itself
   "offline", causing getDocument() calls to resolve with cached results, rather
   than continuing to wait.
+- [fixed] Fixed a potential race condition after calling `enableNetwork()` that
+  could result in a "Mutation batchIDs must be acknowledged in order" assertion
+  crash.
 
 # v0.10.3
 - [fixed] Fixed a regression in the 4.10.0 Firebase iOS SDK release that
diff --git a/Firestore/Example/Tests/Local/FSTMutationQueueTests.mm b/Firestore/Example/Tests/Local/FSTMutationQueueTests.mm
index 7d305d0..f5294d5 100644
--- a/Firestore/Example/Tests/Local/FSTMutationQueueTests.mm
+++ b/Firestore/Example/Tests/Local/FSTMutationQueueTests.mm
@@ -217,6 +217,22 @@ NS_ASSUME_NONNULL_BEGIN
   XCTAssertNil(notFound);
 }
 
+- (void)testNextMutationBatchAfterBatchIDSkipsAcknowledgedBatches {
+  if ([self isTestBaseClass]) return;
+
+  NSMutableArray<FSTMutationBatch *> *batches = [self createBatches:3];
+  XCTAssertEqualObjects([self.mutationQueue nextMutationBatchAfterBatchID:kFSTBatchIDUnknown],
+                        batches[0]);
+
+  [self acknowledgeBatch:batches[0]];
+  XCTAssertEqualObjects([self.mutationQueue nextMutationBatchAfterBatchID:kFSTBatchIDUnknown],
+                        batches[1]);
+  XCTAssertEqualObjects([self.mutationQueue nextMutationBatchAfterBatchID:batches[0].batchID],
+                        batches[1]);
+  XCTAssertEqualObjects([self.mutationQueue nextMutationBatchAfterBatchID:batches[1].batchID],
+                        batches[2]);
+}
+
 - (void)testAllMutationBatchesThroughBatchID {
   if ([self isTestBaseClass]) return;
 
diff --git a/Firestore/Example/Tests/SpecTests/json/write_spec_test.json b/Firestore/Example/Tests/SpecTests/json/write_spec_test.json
index 8e3f5d5..d4d1e7c 100644
--- a/Firestore/Example/Tests/SpecTests/json/write_spec_test.json
+++ b/Firestore/Example/Tests/SpecTests/json/write_spec_test.json
@@ -3407,6 +3407,198 @@
       }
     ]
   },
+  "Held writes are not re-sent after disable/enable network.": {
+    "describeName": "Writes:",
+    "itName": "Held writes are not re-sent after disable/enable network.",
+    "tags": [],
+    "config": {
+      "useGarbageCollection": true
+    },
+    "steps": [
+      {
+        "userListen": [
+          2,
+          {
+            "path": "collection",
+            "filters": [],
+            "orderBys": []
+          }
+        ],
+        "stateExpect": {
+          "activeTargets": {
+            "2": {
+              "query": {
+                "path": "collection",
+                "filters": [],
+                "orderBys": []
+              },
+              "resumeToken": ""
+            }
+          }
+        }
+      },
+      {
+        "watchAck": [
+          2
+        ]
+      },
+      {
+        "watchEntity": {
+          "docs": [],
+          "targets": [
+            2
+          ]
+        }
+      },
+      {
+        "watchCurrent": [
+          [
+            2
+          ],
+          "resume-token-500"
+        ],
+        "watchSnapshot": 500,
+        "expect": [
+          {
+            "query": {
+              "path": "collection",
+              "filters": [],
+              "orderBys": []
+            },
+            "errorCode": 0,
+            "fromCache": false,
+            "hasPendingWrites": false
+          }
+        ]
+      },
+      {
+        "userSet": [
+          "collection/a",
+          {
+            "v": 1
+          }
+        ],
+        "expect": [
+          {
+            "query": {
+              "path": "collection",
+              "filters": [],
+              "orderBys": []
+            },
+            "added": [
+              [
+                "collection/a",
+                0,
+                {
+                  "v": 1
+                },
+                "local"
+              ]
+            ],
+            "errorCode": 0,
+            "fromCache": false,
+            "hasPendingWrites": true
+          }
+        ]
+      },
+      {
+        "writeAck": {
+          "version": 1000,
+          "expectUserCallback": true
+        },
+        "stateExpect": {
+          "writeStreamRequestCount": 2
+        }
+      },
+      {
+        "enableNetwork": false,
+        "stateExpect": {
+          "activeTargets": {},
+          "limboDocs": [],
+          "writeStreamRequestCount": 3
+        },
+        "expect": [
+          {
+            "query": {
+              "path": "collection",
+              "filters": [],
+              "orderBys": []
+            },
+            "errorCode": 0,
+            "fromCache": true,
+            "hasPendingWrites": true
+          }
+        ]
+      },
+      {
+        "enableNetwork": true,
+        "stateExpect": {
+          "activeTargets": {
+            "2": {
+              "query": {
+                "path": "collection",
+                "filters": [],
+                "orderBys": []
+              },
+              "resumeToken": "resume-token-500"
+            }
+          },
+          "writeStreamRequestCount": 3
+        }
+      },
+      {
+        "watchAck": [
+          2
+        ]
+      },
+      {
+        "watchEntity": {
+          "docs": [
+            [
+              "collection/a",
+              1000,
+              {
+                "v": 1
+              }
+            ]
+          ],
+          "targets": [
+            2
+          ]
+        }
+      },
+      {
+        "watchCurrent": [
+          [
+            2
+          ],
+          "resume-token-2000"
+        ],
+        "watchSnapshot": 2000,
+        "expect": [
+          {
+            "query": {
+              "path": "collection",
+              "filters": [],
+              "orderBys": []
+            },
+            "metadata": [
+              [
+                "collection/a",
+                1000,
+                {
+                  "v": 1
+                }
+              ]
+            ],
+            "errorCode": 0,
+            "fromCache": false,
+            "hasPendingWrites": false
+          }
+        ]
+      }
+    ]
+  },
   "Held writes are released when there are no queries left.": {
     "describeName": "Writes:",
     "itName": "Held writes are released when there are no queries left.",
diff --git a/Firestore/Source/Local/FSTLevelDBMutationQueue.mm b/Firestore/Source/Local/FSTLevelDBMutationQueue.mm
index 575e98d..d7b5eca 100644
--- a/Firestore/Source/Local/FSTLevelDBMutationQueue.mm
+++ b/Firestore/Source/Local/FSTLevelDBMutationQueue.mm
@@ -315,7 +315,12 @@ static ReadOptions StandardReadOptions() {
 }
 
 - (nullable FSTMutationBatch *)nextMutationBatchAfterBatchID:(FSTBatchID)batchID {
-  std::string key = [self mutationKeyForBatchID:batchID + 1];
+  // All batches with batchID <= self.metadata.lastAcknowledgedBatchId have been acknowledged so
+  // the first unacknowledged batch after batchID will have a batchID larger than both of these
+  // values.
+  FSTBatchID nextBatchID = MAX(batchID, self.metadata.lastAcknowledgedBatchId) + 1;
+
+  std::string key = [self mutationKeyForBatchID:nextBatchID];
   std::unique_ptr<Iterator> it(_db->NewIterator(StandardReadOptions()));
   it->Seek(key);
 
@@ -336,7 +341,7 @@ static ReadOptions StandardReadOptions() {
     return nil;
   }
 
-  FSTAssert(rowKey.batchID > batchID, @"Should have found mutation after %d", batchID);
+  FSTAssert(rowKey.batchID >= nextBatchID, @"Should have found mutation after %d", nextBatchID);
   return [self decodedMutationBatch:it->value()];
 }
 
diff --git a/Firestore/Source/Local/FSTMemoryMutationQueue.mm b/Firestore/Source/Local/FSTMemoryMutationQueue.mm
index 5b2fca6..7e5cc02 100644
--- a/Firestore/Source/Local/FSTMemoryMutationQueue.mm
+++ b/Firestore/Source/Local/FSTMemoryMutationQueue.mm
@@ -192,10 +192,10 @@ static const NSComparator NumberComparator = ^NSComparisonResult(NSNumber *left,
 
   // All batches with batchID <= self.highestAcknowledgedBatchID have been acknowledged so the
   // first unacknowledged batch after batchID will have a batchID larger than both of these values.
-  batchID = MAX(batchID + 1, self.highestAcknowledgedBatchID);
+  FSTBatchID nextBatchID = MAX(batchID, self.highestAcknowledgedBatchID) + 1;
 
   // The requested batchID may still be out of range so normalize it to the start of the queue.
-  NSInteger rawIndex = [self indexOfBatchID:batchID];
+  NSInteger rawIndex = [self indexOfBatchID:nextBatchID];
   NSUInteger index = rawIndex < 0 ? 0 : (NSUInteger)rawIndex;
 
   // Finally return the first non-tombstone batch.