Firestore: Fix memory leak in Query.whereField() (#14300)

Author: dconeybe

Diff for: Firestore/CHANGELOG.md

@@ -1,3 +1,6 @@
+# Unreleased
+- [fixed] Fixed memory leak in `Query.whereField()`. (#13978)
+
 # 11.8.0
 - [fixed] Fixed use-after-free bug when internally formatting strings. (#14306)
 - [changed] Update gRPC dependency to 1.69.

Diff for: Firestore/Example/Firestore.xcodeproj/project.pbxproj

@@ -17,6 +17,7 @@
 #include "Firestore/core/src/core/composite_filter.h"
 
 #include <algorithm>
+#include <memory>
 #include <utility>
 
 #include "Firestore/core/src/core/field_filter.h"
@@ -141,16 +142,14 @@ const FieldFilter* CompositeFilter::Rep::FindFirstMatchingFilter(
   return nullptr;
 }
 
-const std::vector<FieldFilter>& CompositeFilter::Rep::GetFlattenedFilters()
-    const {
-  return memoized_flattened_filters_->memoize([&]() {
-    std::vector<FieldFilter> flattened_filters;
-    for (const auto& filter : filters())
-      std::copy(filter.GetFlattenedFilters().begin(),
-                filter.GetFlattenedFilters().end(),
-                std::back_inserter(flattened_filters));
-    return flattened_filters;
-  });
+std::shared_ptr<std::vector<FieldFilter>>
+CompositeFilter::Rep::CalculateFlattenedFilters() const {
+  auto flattened_filters = std::make_shared<std::vector<FieldFilter>>();
+  for (const auto& filter : filters())
+    std::copy(filter.GetFlattenedFilters().begin(),
+              filter.GetFlattenedFilters().end(),
+              std::back_inserter(*flattened_filters));
+  return flattened_filters;
 }
 
 }  // namespace core

Diff for: Firestore/core/src/core/composite_filter.cc

@@ -138,7 +138,8 @@ class CompositeFilter : public Filter {
       return filters_.empty();
     }
 
-    const std::vector<FieldFilter>& GetFlattenedFilters() const override;
+    std::shared_ptr<std::vector<FieldFilter>> CalculateFlattenedFilters()
+        const override;
 
     std::vector<Filter> GetFilters() const override {
       return filters();

Diff for: Firestore/core/src/core/composite_filter.h

@@ -16,6 +16,7 @@
 
 #include "Firestore/core/src/core/field_filter.h"
 
+#include <memory>
 #include <utility>
 
 #include "Firestore/core/src/core/array_contains_any_filter.h"
@@ -122,12 +123,12 @@ FieldFilter::FieldFilter(std::shared_ptr<const Filter::Rep> rep)
     : Filter(std::move(rep)) {
 }
 
-const std::vector<FieldFilter>& FieldFilter::Rep::GetFlattenedFilters() const {
+std::shared_ptr<std::vector<FieldFilter>>
+FieldFilter::Rep::CalculateFlattenedFilters() const {
   // This is already a field filter, so we return a vector of size one.
-  return memoized_flattened_filters_->memoize([&]() {
-    return std::vector<FieldFilter>{
-        FieldFilter(std::make_shared<const Rep>(*this))};
-  });
+  auto filters = std::make_shared<std::vector<FieldFilter>>();
+  filters->push_back(FieldFilter(std::make_shared<const Rep>(*this)));
+  return filters;
 }
 
 std::vector<Filter> FieldFilter::Rep::GetFilters() const {

Diff for: Firestore/core/src/core/field_filter.cc

@@ -117,8 +117,6 @@ class FieldFilter : public Filter {
       return false;
     }
 
-    const std::vector<FieldFilter>& GetFlattenedFilters() const override;
-
     std::vector<Filter> GetFilters() const override;
 
    protected:
@@ -140,6 +138,9 @@ class FieldFilter : public Filter {
 
     bool MatchesComparison(util::ComparisonResult comparison) const;
 
+    std::shared_ptr<std::vector<FieldFilter>> CalculateFlattenedFilters()
+        const override;
+
    private:
     friend class FieldFilter;
 

Diff for: Firestore/core/src/core/field_filter.h

@@ -35,12 +35,6 @@ std::ostream& operator<<(std::ostream& os, const Filter& filter) {
   return os << filter.ToString();
 }
 
-Filter::Rep::Rep()
-    : memoized_flattened_filters_(
-          std::make_shared<
-              util::ThreadSafeMemoizer<std::vector<FieldFilter>>>()) {
-}
-
 }  // namespace core
 }  // namespace firestore
 }  // namespace firebase

Diff for: Firestore/core/src/core/filter.cc

@@ -17,6 +17,7 @@
 #ifndef FIRESTORE_CORE_SRC_CORE_FILTER_H_
 #define FIRESTORE_CORE_SRC_CORE_FILTER_H_
 
+#include <functional>
 #include <iosfwd>
 #include <memory>
 #include <string>
@@ -114,7 +115,7 @@ class Filter {
  protected:
   class Rep {
    public:
-    Rep();
+    Rep() = default;
 
     virtual ~Rep() = default;
 
@@ -147,20 +148,23 @@ class Filter {
 
     virtual bool IsEmpty() const = 0;
 
-    virtual const std::vector<FieldFilter>& GetFlattenedFilters() const = 0;
+    virtual const std::vector<FieldFilter>& GetFlattenedFilters() const {
+      const auto func = std::bind(&Rep::CalculateFlattenedFilters, this);
+      return memoized_flattened_filters_.value(func);
+    }
 
     virtual std::vector<Filter> GetFilters() const = 0;
 
+   protected:
+    virtual std::shared_ptr<std::vector<FieldFilter>>
+    CalculateFlattenedFilters() const = 0;
+
+   private:
     /**
      * Memoized list of all field filters that can be found by
      * traversing the tree of filters contained in this composite filter.
-     *
-     * Use a `std::shared_ptr<ThreadSafeMemoizer>` rather than using
-     * `ThreadSafeMemoizer` directly so that this class is copyable
-     * (`ThreadSafeMemoizer` is not copyable because of its `std::once_flag`
-     * member variable, which is not copyable).
      */
-    mutable std::shared_ptr<util::ThreadSafeMemoizer<std::vector<FieldFilter>>>
+    mutable util::ThreadSafeMemoizer<std::vector<FieldFilter>>
         memoized_flattened_filters_;
   };
 

Diff for: Firestore/core/src/core/filter.h

@@ -17,6 +17,7 @@
 #include "Firestore/core/src/core/query.h"
 
 #include <algorithm>
+#include <memory>
 #include <ostream>
 
 #include "Firestore/core/src/core/bound.h"
@@ -91,44 +92,42 @@ absl::optional<Operator> Query::FindOpInsideFilters(
   return absl::nullopt;
 }
 
-const std::vector<OrderBy>& Query::normalized_order_bys() const {
-  return memoized_normalized_order_bys_->memoize([&]() {
-    // Any explicit order by fields should be added as is.
-    std::vector<OrderBy> result = explicit_order_bys_;
-    std::set<FieldPath> fieldsNormalized;
-    for (const OrderBy& order_by : explicit_order_bys_) {
-      fieldsNormalized.insert(order_by.field());
-    }
+std::shared_ptr<std::vector<OrderBy>> Query::CalculateNormalizedOrderBys()
+    const {
+  // Any explicit order by fields should be added as is.
+  auto result = std::make_shared<std::vector<OrderBy>>(explicit_order_bys_);
+  std::set<FieldPath> fieldsNormalized;
+  for (const OrderBy& order_by : explicit_order_bys_) {
+    fieldsNormalized.insert(order_by.field());
+  }
 
-    // The order of the implicit ordering always matches the last explicit order
-    // by.
-    Direction last_direction = explicit_order_bys_.empty()
-                                   ? Direction::Ascending
-                                   : explicit_order_bys_.back().direction();
-
-    // Any inequality fields not explicitly ordered should be implicitly ordered
-    // in a lexicographical order. When there are multiple inequality filters on
-    // the same field, the field should be added only once. Note:
-    // `std::set<model::FieldPath>` sorts the key field before other fields.
-    // However, we want the key field to be sorted last.
-    const std::set<model::FieldPath> inequality_fields =
-        InequalityFilterFields();
-
-    for (const model::FieldPath& field : inequality_fields) {
-      if (fieldsNormalized.find(field) == fieldsNormalized.end() &&
-          !field.IsKeyFieldPath()) {
-        result.push_back(OrderBy(field, last_direction));
-      }
+  // The order of the implicit ordering always matches the last explicit order
+  // by.
+  Direction last_direction = explicit_order_bys_.empty()
+                                 ? Direction::Ascending
+                                 : explicit_order_bys_.back().direction();
+
+  // Any inequality fields not explicitly ordered should be implicitly ordered
+  // in a lexicographical order. When there are multiple inequality filters on
+  // the same field, the field should be added only once. Note:
+  // `std::set<model::FieldPath>` sorts the key field before other fields.
+  // However, we want the key field to be sorted last.
+  const std::set<model::FieldPath> inequality_fields = InequalityFilterFields();
+
+  for (const model::FieldPath& field : inequality_fields) {
+    if (fieldsNormalized.find(field) == fieldsNormalized.end() &&
+        !field.IsKeyFieldPath()) {
+      result->push_back(OrderBy(field, last_direction));
     }
+  }
 
-    // Add the document key field to the last if it is not explicitly ordered.
-    if (fieldsNormalized.find(FieldPath::KeyFieldPath()) ==
-        fieldsNormalized.end()) {
-      result.push_back(OrderBy(FieldPath::KeyFieldPath(), last_direction));
-    }
+  // Add the document key field to the last if it is not explicitly ordered.
+  if (fieldsNormalized.find(FieldPath::KeyFieldPath()) ==
+      fieldsNormalized.end()) {
+    result->push_back(OrderBy(FieldPath::KeyFieldPath(), last_direction));
+  }
 
-    return result;
-  });
+  return result;
 }
 
 LimitType Query::limit_type() const {
@@ -296,14 +295,12 @@ std::string Query::ToString() const {
   return absl::StrCat("Query(canonical_id=", CanonicalId(), ")");
 }
 
-const Target& Query::ToTarget() const& {
-  return memoized_target_->memoize(
-      [&]() { return ToTarget(normalized_order_bys()); });
+std::shared_ptr<Target> Query::CalculateTarget() const {
+  return std::make_shared<Target>(ToTarget(normalized_order_bys()));
 }
 
-const Target& Query::ToAggregateTarget() const& {
-  return memoized_aggregate_target_->memoize(
-      [&]() { return ToTarget(explicit_order_bys_); });
+std::shared_ptr<Target> Query::CalculateAggregateTarget() const {
+  return std::make_shared<Target>(ToTarget(explicit_order_bys_));
 }
 
 Target Query::ToTarget(const std::vector<OrderBy>& order_bys) const {

Diff for: Firestore/core/src/core/query.cc

@@ -17,6 +17,7 @@
 #ifndef FIRESTORE_CORE_SRC_CORE_QUERY_H_
 #define FIRESTORE_CORE_SRC_CORE_QUERY_H_
 
+#include <functional>
 #include <iosfwd>
 #include <limits>
 #include <memory>
@@ -148,7 +149,10 @@ class Query {
    * This might include additional sort orders added implicitly to match the
    * backend behavior.
    */
-  const std::vector<OrderBy>& normalized_order_bys() const;
+  const std::vector<OrderBy>& normalized_order_bys() const {
+    const auto func = std::bind(&Query::CalculateNormalizedOrderBys, this);
+    return memoized_normalized_order_bys_.value(func);
+  }
 
   bool has_limit() const {
     return limit_ != Target::kNoLimit;
@@ -246,15 +250,21 @@ class Query {
    * Returns a `Target` instance this query will be mapped to in backend
    * and local store.
    */
-  const Target& ToTarget() const&;
+  const Target& ToTarget() const& {
+    const auto func = std::bind(&Query::CalculateTarget, this);
+    return memoized_target_.value(func);
+  }
 
   /**
    * Returns a `Target` instance this query will be mapped to in backend
    * and local store, for use within an aggregate query. Unlike targets
    * for non-aggregate queries, aggregate query targets do not contain
    * normalized order-bys, they only contain explicit order-bys.
    */
-  const Target& ToAggregateTarget() const&;
+  const Target& ToAggregateTarget() const& {
+    const auto func = std::bind(&Query::CalculateAggregateTarget, this);
+    return memoized_aggregate_target_.value(func);
+  }
 
   friend std::ostream& operator<<(std::ostream& os, const Query& query);
 
@@ -295,20 +305,19 @@ class Query {
   // member variable, which is not copyable).
 
   // The memoized list of sort orders.
-  mutable std::shared_ptr<util::ThreadSafeMemoizer<std::vector<OrderBy>>>
-      memoized_normalized_order_bys_{
-          std::make_shared<util::ThreadSafeMemoizer<std::vector<OrderBy>>>()};
+  std::shared_ptr<std::vector<OrderBy>> CalculateNormalizedOrderBys() const;
+  mutable util::ThreadSafeMemoizer<std::vector<OrderBy>>
+      memoized_normalized_order_bys_;
 
   // The corresponding Target of this Query instance.
-  mutable std::shared_ptr<util::ThreadSafeMemoizer<Target>> memoized_target_{
-      std::make_shared<util::ThreadSafeMemoizer<Target>>()};
+  std::shared_ptr<Target> CalculateTarget() const;
+  mutable util::ThreadSafeMemoizer<Target> memoized_target_;
 
   // The corresponding aggregate Target of this Query instance. Unlike targets
   // for non-aggregate queries, aggregate query targets do not contain
   // normalized order-bys, they only contain explicit order-bys.
-  mutable std::shared_ptr<util::ThreadSafeMemoizer<Target>>
-      memoized_aggregate_target_{
-          std::make_shared<util::ThreadSafeMemoizer<Target>>()};
+  std::shared_ptr<Target> CalculateAggregateTarget() const;
+  mutable util::ThreadSafeMemoizer<Target> memoized_aggregate_target_;
 };
 
 bool operator==(const Query& lhs, const Query& rhs);