Expand the functionality of the BudgetDatabase.

chrome/browser/budget_service/budget_database.h

Index: chrome/browser/budget_service/budget_database.h
diff --git a/chrome/browser/budget_service/budget_database.h b/chrome/browser/budget_service/budget_database.h
index 155f23d22c8f6e708c1f7282669e55b7ce0a9ae0..389304992c50a74900d8dd8ef82d186e3f28c844 100644
--- a/chrome/browser/budget_service/budget_database.h
+++ b/chrome/browser/budget_service/budget_database.h
@@ -5,9 +5,11 @@
 #ifndef CHROME_BROWSER_BUDGET_SERVICE_BUDGET_DATABASE_H_
 #define CHROME_BROWSER_BUDGET_SERVICE_BUDGET_DATABASE_H_
 
+#include <list>
 #include <memory>
 
 #include "base/callback_forward.h"
+#include "base/containers/hash_tables.h"

[Peter Beverloo, 2016/07/26 17:26:21]

delete

[harkness, 2016/07/27 10:59:03]

Done.

 #include "base/macros.h"
 #include "base/memory/weak_ptr.h"
 #include "components/leveldb_proto/proto_database.h"
@@ -26,6 +28,9 @@ class GURL;
 // assigned to an origin. The class uses an underlying LevelDB.
 class BudgetDatabase {
  public:
+  // Data structure for returing the budget decay expectations to the caller.
+  using BudgetExpectation = std::list<std::pair<double, double>>;
+
   // Callback for the basic GetBudget call.
   using GetValueCallback =
       base::Callback<void(bool success,
@@ -34,6 +39,18 @@ class BudgetDatabase {
   // Callback for setting a budget value.
   using SetValueCallback = base::Callback<void(bool success)>;
 
+  // Callback for getting a list of all budget chunks.
+  using GetBudgetDetailsCallback = base::Callback<
+      void(bool success, double debt, const BudgetExpectation& expectation)>;
+
+ private:
+  // Data structure for caching budget information.
+  using BudgetChunks = std::vector<std::pair<double, double>>;
+  using BudgetInfo = std::pair<double, BudgetChunks>;
+
+  using AddToCacheCallback = base::Callback<void(bool success)>;
+
+ public:
   // The database_dir specifies the location of the budget information on
   // disk. The task_runner is used by the ProtoDatabase to handle all blocking
   // calls and disk access.
@@ -46,12 +63,30 @@ class BudgetDatabase {
                 const budget_service::Budget& budget,
                 const SetValueCallback& callback);
 
+  // Get the full budget expectation for the origin. This will return any
+  // debt as well as a sequence of time points and the expected budget at
+  // those times.
+  void GetBudgetDetails(const GURL& origin,
+                        const GetBudgetDetailsCallback& callback);
+
  private:
   void OnDatabaseInit(bool success);
 
+  void AddToCache(const GURL& origin,
+                  const AddToCacheCallback& callback,
+                  bool success,
+                  std::unique_ptr<budget_service::Budget> budget);
+
+  void DidGetBudget(const GURL& origin,
+                    const GetBudgetDetailsCallback& callback,
+                    bool success);
+
   // The database for storing budget information.
   std::unique_ptr<leveldb_proto::ProtoDatabase<budget_service::Budget>> db_;
 
+  // Cached data for the origins which have been loaded.
+  std::unordered_map<std::string, BudgetInfo> budget_map_;
+
   base::WeakPtrFactory<BudgetDatabase> weak_ptr_factory_;
 
   DISALLOW_COPY_AND_ASSIGN(BudgetDatabase);