fix latency compensated query

Author: wu-hui

Firestore/Example/Tests/SpecTests/json/query_spec_test.json

@@ -1166,5 +1166,208 @@
         ]
       }
     ]
+  },
+  "Latency-compensated updates are included in query results": {
+    "describeName": "Queries:",
+    "itName": "Latency-compensated updates are included in query results",
+    "tags": [],
+    "config": {
+      "useGarbageCollection": true,
+      "numClients": 1
+    },
+    "steps": [
+      {
+        "userListen": [
+          2,
+          {
+            "path": "collection",
+            "filters": [],
+            "orderBys": []
+          }
+        ],
+        "stateExpect": {
+          "activeTargets": {
+            "2": {
+              "query": {
+                "path": "collection",
+                "filters": [],
+                "orderBys": []
+              },
+              "resumeToken": ""
+            }
+          }
+        }
+      },
+      {
+        "watchAck": [
+          2
+        ]
+      },
+      {
+        "watchEntity": {
+          "docs": [
+            {
+              "key": "collection/a",
+              "version": 1000,
+              "value": {
+                "match": false
+              },
+              "options": {
+                "hasLocalMutations": false,
+                "hasCommittedMutations": false
+              }
+            }
+          ],
+          "targets": [
+            2
+          ]
+        }
+      },
+      {
+        "watchCurrent": [
+          [
+            2
+          ],
+          "resume-token-1000"
+        ]
+      },
+      {
+        "watchSnapshot": {
+          "version": 1000,
+          "targetIds": []
+        },
+        "expect": [
+          {
+            "query": {
+              "path": "collection",
+              "filters": [],
+              "orderBys": []
+            },
+            "added": [
+              {
+                "key": "collection/a",
+                "version": 1000,
+                "value": {
+                  "match": false
+                },
+                "options": {
+                  "hasLocalMutations": false,
+                  "hasCommittedMutations": false
+                }
+              }
+            ],
+            "errorCode": 0,
+            "fromCache": false,
+            "hasPendingWrites": false
+          }
+        ]
+      },
+      {
+        "userPatch": [
+          "collection/a",
+          {
+            "match": true
+          }
+        ],
+        "expect": [
+          {
+            "query": {
+              "path": "collection",
+              "filters": [],
+              "orderBys": []
+            },
+            "modified": [
+              {
+                "key": "collection/a",
+                "version": 1000,
+                "value": {
+                  "match": true
+                },
+                "options": {
+                  "hasLocalMutations": true,
+                  "hasCommittedMutations": false
+                }
+              }
+            ],
+            "errorCode": 0,
+            "fromCache": false,
+            "hasPendingWrites": true
+          }
+        ]
+      },
+      {
+        "userListen": [
+          4,
+          {
+            "path": "collection",
+            "filters": [
+              [
+                "match",
+                "==",
+                true
+              ]
+            ],
+            "orderBys": []
+          }
+        ],
+        "stateExpect": {
+          "activeTargets": {
+            "2": {
+              "query": {
+                "path": "collection",
+                "filters": [],
+                "orderBys": []
+              },
+              "resumeToken": ""
+            },
+            "4": {
+              "query": {
+                "path": "collection",
+                "filters": [
+                  [
+                    "match",
+                    "==",
+                    true
+                  ]
+                ],
+                "orderBys": []
+              },
+              "resumeToken": ""
+            }
+          }
+        },
+        "expect": [
+          {
+            "query": {
+              "path": "collection",
+              "filters": [
+                [
+                  "match",
+                  "==",
+                  true
+                ]
+              ],
+              "orderBys": []
+            },
+            "added": [
+              {
+                "key": "collection/a",
+                "version": 1000,
+                "value": {
+                  "match": true
+                },
+                "options": {
+                  "hasLocalMutations": true,
+                  "hasCommittedMutations": false
+                }
+              }
+            ],
+            "errorCode": 0,
+            "fromCache": true,
+            "hasPendingWrites": true
+          }
+        ]
+      }
+    ]
   }
 }

Firestore/core/src/firebase/firestore/local/local_documents_view.h

@@ -103,6 +103,15 @@ class LocalDocumentsView {
   /** Queries the remote documents and overlays mutations. */
   model::DocumentMap GetDocumentsMatchingCollectionQuery(FSTQuery* query);
 
+  /**
+   * Find all `PatchMutation`s, and if their base documents are not in
+   * `existingDocs` add them to the returning `DocumentMap`. The returned map
+   * also has every document in `existingDocs`.
+   */
+  model::DocumentMap AddMissingBaseDocument(
+      const std::vector<FSTMutationBatch*>& matchingBatches,
+      const model::DocumentMap& existingDocs);
+
   RemoteDocumentCache* remote_document_cache_;
   MutationQueue* mutation_queue_;
   IndexManager* index_manager_;

Firestore/core/src/firebase/firestore/local/local_documents_view.mm

@@ -171,6 +171,8 @@
   std::vector<FSTMutationBatch*> matchingBatches =
       mutation_queue_->AllMutationBatchesAffectingQuery(query);
 
+  results = AddMissingBaseDocument(matchingBatches, results);
+
   for (FSTMutationBatch* batch : matchingBatches) {
     for (FSTMutation* mutation : [batch mutations]) {
       // Only process documents belonging to the collection.
@@ -215,6 +217,38 @@
   return results;
 }
 
+// It is possible that a `PatchMutation` can make a document match a query, even
+// if the version in the `RemoteDocumentCache` is not a match yet (waiting for
+// server to ack). To handle this, we find all document keys affected by the
+// `PatchMutation`s that are not in `existingDocs` yet, and back fill them via
+// `remoteDocumentCache.getAll`, otherwise those `PatchMutation`s will be
+// ignored because no base document can be found, and lead to missing results
+// for the query.
+DocumentMap LocalDocumentsView::AddMissingBaseDocument(
+    const std::vector<FSTMutationBatch*>& matchingBatches,
+    const DocumentMap& existingDocs) {
+  DocumentKeySet missingDocKeys;
+  for (FSTMutationBatch* batch : matchingBatches) {
+    for (FSTMutation* mutation : [batch mutations]) {
+      if ([mutation isKindOfClass:[FSTPatchMutation class]] &&
+          existingDocs.underlying_map().find([mutation key]) ==
+              existingDocs.underlying_map().end()) {
+        missingDocKeys = missingDocKeys.insert([mutation key]);
+      }
+    }
+  }
+
+  DocumentMap mergedDocs = existingDocs;
+  MaybeDocumentMap missingDocs = remote_document_cache_->GetAll(missingDocKeys);
+  for (const auto& kv : missingDocs) {
+    if (kv.second != nil && [kv.second isKindOfClass:[FSTDocument class]]) {
+      mergedDocs = mergedDocs.insert(kv.first, (FSTDocument*)kv.second);
+    }
+  }
+
+  return mergedDocs;
+}
+
 }  // namespace local
 }  // namespace firestore
 }  // namespace firebase