Hooked up the PruneForest so that cached values are appropriately pruned.

PiperOrigin-RevId: 285262340

Author: amablue

database/CMakeLists.txt

@@ -72,6 +72,7 @@ set(desktop_SRCS
     src/desktop/connection/persistent_connection.cc
     src/desktop/connection/util_connection.cc
     src/desktop/connection/web_socket_client_impl.cc
+    src/desktop/core/cache_policy.cc
     src/desktop/core/child_event_registration.cc
     src/desktop/core/compound_write.cc
     src/desktop/core/constants.cc
@@ -98,6 +99,7 @@ set(desktop_SRCS
     src/desktop/persistence/in_memory_persistence_storage_engine.cc
     src/desktop/persistence/noop_persistence_storage_engine.cc
     src/desktop/persistence/persistence_manager.cc
+    src/desktop/persistence/prune_forest.cc
     src/desktop/push_child_name_generator.cc
     src/desktop/query_desktop.cc
     src/desktop/query_params_comparator.cc

database/src/desktop/core/repo.cc

@@ -374,8 +374,7 @@ void Repo::UpdateChildren(const Path& path, const Variant& data,
 
   // Start with our existing data and merge each child into it.
   Variant server_values = GenerateServerValues(server_time_offset_);
-  CompoundWrite resolved =
-      ResolveDeferredValueMerge(updates, server_values);
+  CompoundWrite resolved = ResolveDeferredValueMerge(updates, server_values);
 
   WriteId write_id = GetNextWriteId();
   std::vector<Event> events = server_sync_tree_->ApplyUserMerge(
@@ -425,14 +424,19 @@ void Repo::AckWriteAndRerunTransactions(WriteId write_id, const Path& path,
 
 static UniquePtr<SyncTree> InitializeSyncTree(
     UniquePtr<ListenProvider> listen_provider, Logger* logger) {
+  static const uint64_t kDefaultCacheSize = 10 * 1024 * 1024;
+
   UniquePtr<WriteTree> pending_write_tree = MakeUnique<WriteTree>();
   UniquePtr<PersistenceStorageEngine> persistence_storage_engine =
       MakeUnique<InMemoryPersistenceStorageEngine>(logger);
   UniquePtr<TrackedQueryManager> tracked_query_manager =
-      MakeUnique<TrackedQueryManager>(persistence_storage_engine.get());
+      MakeUnique<TrackedQueryManager>(persistence_storage_engine.get(), logger);
+  UniquePtr<CachePolicy> cache_policy =
+      MakeUnique<LRUCachePolicy>(kDefaultCacheSize);
   UniquePtr<PersistenceManager> persistence_manager =
       MakeUnique<PersistenceManager>(std::move(persistence_storage_engine),
-                                     std::move(tracked_query_manager));
+                                     std::move(tracked_query_manager),
+                                     std::move(cache_policy), logger);
   return MakeUnique<SyncTree>(std::move(pending_write_tree),
                               std::move(persistence_manager),
                               std::move(listen_provider));

database/src/desktop/core/tracked_query_manager.cc

@@ -16,12 +16,14 @@
 
 #include <algorithm>
 #include <cstdint>
+#include <ctime>
 #include <map>
 
 #include "app/src/assert.h"
 #include "app/src/path.h"
 #include "database/src/common/query_spec.h"
 #include "database/src/desktop/persistence/persistence_storage_engine.h"
+#include "database/src/desktop/persistence/prune_forest.h"
 #include "database/src/desktop/util_desktop.h"
 
 namespace firebase {
@@ -51,7 +53,7 @@ static bool IsQueryPrunablePredicate(const TrackedQuery& query) {
 
 // Returns true if the given TrackedQuery is not prunable. A query is considered
 // prunable if it is not active.
-static bool IsQueryUnPrunablePredicate(const TrackedQuery& query) {
+static bool IsQueryUnprunablePredicate(const TrackedQuery& query) {
   return query.active;
 }
 
@@ -69,10 +71,11 @@ static void AssertValidTrackedQuery(const QuerySpec& query_spec) {
 }
 
 TrackedQueryManager::TrackedQueryManager(
-    PersistenceStorageEngine* storage_engine)
+    PersistenceStorageEngine* storage_engine, LoggerBase* logger)
     : storage_engine_(storage_engine),
       tracked_query_tree_(),
-      next_query_id_(0) {
+      next_query_id_(0),
+      logger_(logger) {
   ResetPreviouslyActiveTrackedQueries();
 
   // Populate our cache from the storage layer.
@@ -84,6 +87,8 @@ TrackedQueryManager::TrackedQueryManager(
   }
 }
 
+TrackedQueryManager::~TrackedQueryManager() {}
+
 const TrackedQuery* TrackedQueryManager::FindTrackedQuery(
     const QuerySpec& query_spec) const {
   QuerySpec normalized_spec = GetNormalizedQuery(query_spec);
@@ -114,7 +119,9 @@ void TrackedQueryManager::SetQueryActiveFlag(
   QuerySpec normalized_spec = GetNormalizedQuery(query_spec);
   const TrackedQuery* tracked_query = FindTrackedQuery(normalized_spec);
 
-  uint64_t last_use = 0;
+  // TODO(amablue): Set up a more robust clock that won't get confused if, for
+  // example, the system time changes while the app is running.
+  uint64_t last_use = time(nullptr) * 1000;
   if (tracked_query != nullptr) {
     TrackedQuery updated_tracked_query = *tracked_query;
     updated_tracked_query.last_use = last_use;
@@ -173,6 +180,65 @@ bool TrackedQueryManager::IsQueryComplete(const QuerySpec& query_spec) {
   }
 }
 
+static uint64_t CalculateCountToPrune(const CachePolicy& cache_policy,
+                                      uint64_t prunable_count) {
+  uint64_t count_to_keep = prunable_count;
+
+  // Prune by percentage.
+  double percent_to_keep =
+      1.0 - cache_policy.GetPercentOfQueriesToPruneAtOnce();
+  count_to_keep = static_cast<uint64_t>(count_to_keep * percent_to_keep);
+
+  // Make sure we're not keeping more than the max.
+  count_to_keep =
+      std::min(count_to_keep, cache_policy.GetMaxNumberOfQueriesToKeep());
+
+  // Now we know how many to prune.
+  return prunable_count - count_to_keep;
+}
+
+PruneForest TrackedQueryManager::PruneOldQueries(
+    const CachePolicy& cache_policy) {
+  std::vector<TrackedQuery> prunable =
+      GetQueriesMatching(IsQueryPrunablePredicate);
+  uint64_t count_to_prune =
+      CalculateCountToPrune(cache_policy, prunable.size());
+
+  logger_->LogDebug("Pruning old queries. Prunable: %i Count to prune: %i",
+                    static_cast<int>(prunable.size()),
+                    static_cast<int>(count_to_prune));
+
+  std::sort(prunable.begin(), prunable.end(),
+            [](const TrackedQuery& q1, const TrackedQuery& q2) {
+              return q1.last_use < q2.last_use;
+            });
+
+  // Prune the queries that are no longer needed.
+  PruneForest forest;
+  PruneForestRef forest_ref(&forest);
+  for (uint64_t i = 0; i < count_to_prune; i++) {
+    const TrackedQuery& to_prune = prunable[i];
+    forest_ref.Prune(to_prune.query_spec.path);
+    RemoveTrackedQuery(to_prune.query_spec);
+  }
+  // Keep the rest of the prunable queries.
+  for (uint64_t i = count_to_prune; i < prunable.size(); i++) {
+    const TrackedQuery& to_keep = prunable[i];
+    forest_ref.Keep(to_keep.query_spec.path);
+  }
+  // Also keep the unprunable queries.
+  std::vector<TrackedQuery> unprunable =
+      GetQueriesMatching(IsQueryUnprunablePredicate);
+
+  logger_->LogDebug("Unprunable queries: %i",
+                    static_cast<int>(unprunable.size()));
+  for (const TrackedQuery& to_keep : unprunable) {
+    forest_ref.Keep(to_keep.query_spec.path);
+  }
+
+  return forest;
+}
+
 std::set<std::string> TrackedQueryManager::GetKnownCompleteChildren(
     const Path& path) {
   QuerySpec default_at_path(path);

database/src/desktop/core/tracked_query_manager.h

@@ -18,10 +18,14 @@
 #include <cstdint>
 #include <map>
 #include <set>
+
+#include "app/src/logger.h"
 #include "app/src/optional.h"
 #include "app/src/path.h"
 #include "database/src/common/query_spec.h"
+#include "database/src/desktop/core/cache_policy.h"
 #include "database/src/desktop/core/tree.h"
+#include "database/src/desktop/persistence/prune_forest.h"
 
 namespace firebase {
 namespace database {
@@ -35,7 +39,12 @@ struct TrackedQuery {
   enum CompletionStatus { kIncomplete, kComplete };
   enum ActivityStatus { kInactive, kActive };
 
-  TrackedQuery() : query_id(), query_spec(), last_use(), complete(), active() {}
+  TrackedQuery()
+      : query_id(0),
+        query_spec(),
+        last_use(0),
+        complete(false),
+        active(false) {}
 
   TrackedQuery(QueryId _query_id, const QuerySpec& _query_spec,
                uint64_t _last_use, CompletionStatus _complete,
@@ -91,6 +100,10 @@ class TrackedQueryManagerInterface {
   // complete.
   virtual bool IsQueryComplete(const QuerySpec& query) = 0;
 
+  // Remove queries that no longer need to be cached based on the given cache
+  // policy.
+  virtual PruneForest PruneOldQueries(const CachePolicy& cache_policy) = 0;
+
   // Return the keys of the completed TrackedQueries at the given location.
   virtual std::set<std::string> GetKnownCompleteChildren(const Path& path) = 0;
 
@@ -108,7 +121,10 @@ class TrackedQueryManagerInterface {
 
 class TrackedQueryManager : public TrackedQueryManagerInterface {
  public:
-  explicit TrackedQueryManager(PersistenceStorageEngine* storage_engine);
+  TrackedQueryManager(PersistenceStorageEngine* storage_engine,
+                      LoggerBase* logger);
+
+  ~TrackedQueryManager() override;
 
   // Find and return the TrackedQuery associated with the given QuerySpec, or
   // nullptr if there is no associated TrackedQuery.
@@ -136,6 +152,10 @@ class TrackedQueryManager : public TrackedQueryManagerInterface {
   // complete.
   bool IsQueryComplete(const QuerySpec& query) override;
 
+  // Remove queries that no longer need to be cached based on the given cache
+  // policy.
+  PruneForest PruneOldQueries(const CachePolicy& cache_policy) override;
+
   // Return the keys of the completed TrackedQueries at the given location.
   std::set<std::string> GetKnownCompleteChildren(const Path& path) override;
 
@@ -183,6 +203,8 @@ class TrackedQueryManager : public TrackedQueryManagerInterface {
 
   // ID we'll assign to the next tracked query.
   QueryId next_query_id_;
+
+  LoggerBase* logger_;
 };
 
 }  // namespace internal

database/src/desktop/persistence/in_memory_persistence_storage_engine.cc

@@ -22,6 +22,7 @@
 #include "app/src/assert.h"
 #include "app/src/include/firebase/variant.h"
 #include "app/src/log.h"
+#include "app/src/logger.h"
 #include "app/src/path.h"
 #include "app/src/variant_util.h"
 #include "database/src/common/query_spec.h"
@@ -36,7 +37,7 @@ namespace database {
 namespace internal {
 
 InMemoryPersistenceStorageEngine::InMemoryPersistenceStorageEngine(
-    Logger* logger)
+    LoggerBase* logger)
     : server_cache_(), inside_transaction_(false), logger_(logger) {}
 
 InMemoryPersistenceStorageEngine::~InMemoryPersistenceStorageEngine() {}
@@ -91,16 +92,68 @@ void InMemoryPersistenceStorageEngine::MergeIntoServerCache(
     const Path& path, const Variant& data) {
   VerifyInTransaction();
   Variant* target = MakeVariantAtPath(&server_cache_, path);
-  if (!target->is_map()) *target = Variant::EmptyMap();
-  PatchVariant(data, target);
+  if (data.is_map()) {
+    if (!target->is_map()) *target = Variant::EmptyMap();
+    PatchVariant(data, target);
+  } else {
+    *target = data;
+  }
   // Clean up in case anything was removed.
   PruneNulls(target);
 }
 
 void InMemoryPersistenceStorageEngine::MergeIntoServerCache(
     const Path& path, const CompoundWrite& children) {
-  // TODO(amablue)
   VerifyInTransaction();
+  children.write_tree().CallOnEach(
+      Path(), [this, &path](const Path& child_path, const Variant& value) {
+        this->MergeIntoServerCache(path.GetChild(child_path), value);
+      });
+}
+
+static uint64_t EstimateVariantMemoryUsage(const Variant& variant) {
+  switch (variant.type()) {
+    case Variant::kTypeNull:
+    case Variant::kTypeInt64:
+    case Variant::kTypeDouble:
+    case Variant::kTypeBool: {
+      return sizeof(Variant);
+    }
+    case Variant::kTypeStaticString: {
+      return sizeof(Variant) + strlen(variant.string_value());
+    }
+    case Variant::kTypeMutableString: {
+      return sizeof(Variant) + variant.mutable_string().size();
+    }
+    case Variant::kTypeVector: {
+      uint64_t sum_total = 0;
+      for (const auto& item : variant.vector()) {
+        sum_total += EstimateVariantMemoryUsage(item);
+      }
+      return sizeof(Variant) + sum_total;
+    }
+    case Variant::kTypeMap: {
+      uint64_t sum_total = 0;
+      for (const auto& key_value_pair : variant.map()) {
+        sum_total += EstimateVariantMemoryUsage(key_value_pair.first);
+        sum_total += EstimateVariantMemoryUsage(key_value_pair.second);
+      }
+      return sizeof(Variant) + sum_total;
+    }
+    case Variant::kTypeStaticBlob:
+    case Variant::kTypeMutableBlob: {
+      return sizeof(Variant) + variant.blob_size();
+    }
+    default: {
+      FIREBASE_DEV_ASSERT_MESSAGE(false, "Unhandled variant type.");
+      return 0;
+    }
+  }
+}
+
+uint64_t InMemoryPersistenceStorageEngine::ServerCacheEstimatedSizeInBytes()
+    const {
+  return EstimateVariantMemoryUsage(server_cache_);
 }
 
 void InMemoryPersistenceStorageEngine::SaveTrackedQuery(
@@ -139,7 +192,9 @@ void InMemoryPersistenceStorageEngine::UpdateTrackedQueryKeys(
 
   std::set<std::string>& tracked_keys = tracked_query_keys_[query_id];
   tracked_keys.insert(added.begin(), added.end());
-  tracked_keys.insert(removed.begin(), removed.end());
+  for (const std::string& to_remove : removed) {
+    tracked_keys.erase(to_remove);
+  }
 }
 
 std::set<std::string> InMemoryPersistenceStorageEngine::LoadTrackedQueryKeys(
@@ -157,16 +212,35 @@ std::set<std::string> InMemoryPersistenceStorageEngine::LoadTrackedQueryKeys(
   return result;
 }
 
+void PruneVariant(const Path& root, const PruneForestRef prune_forest,
+                  Variant* variant) {
+  Variant result = prune_forest.FoldKeptNodes(
+      Variant::Null(),
+      [&root, &variant](const Path& relative_path, Variant accum) {
+        Variant child = VariantGetChild(variant, root.GetChild(relative_path));
+        VariantUpdateChild(&accum, relative_path, child);
+        return accum;
+      });
+  VariantUpdateChild(variant, root, result);
+}
+
+void InMemoryPersistenceStorageEngine::PruneCache(
+    const Path& root, const PruneForestRef& prune_forest) {
+  PruneVariant(root, prune_forest, &server_cache_);
+}
+
 bool InMemoryPersistenceStorageEngine::BeginTransaction() {
   FIREBASE_DEV_ASSERT_MESSAGE(!inside_transaction_,
-                              "runInTransaction called when an existing "
+                              "RunInTransaction called when an existing "
                               "transaction is already in progress.");
   logger_->LogDebug("Starting transaction.");
   inside_transaction_ = true;
   return true;
 }
 
 void InMemoryPersistenceStorageEngine::EndTransaction() {
+  FIREBASE_DEV_ASSERT_MESSAGE(
+      inside_transaction_, "EndTransaction called when not in a transaction");
   inside_transaction_ = false;
   logger_->LogDebug("Transaction completed.");
 }