From 5df62ac1721dad6caa1e2ca15789d4e2a044c788 Mon Sep 17 00:00:00 2001
From: Peter Johnson <johnson.peter@gmail.com>
Date: Mon, 20 Jul 2015 22:24:47 -0700
Subject: [PATCH] Storage: Disable use of update queue by default.

This ensures we don't "leak" memory for local use when the dispatch thread
is not running.
---
 src/Storage.cpp           | 24 +++++++++++++-----------
 src/Storage.h             | 10 ++++++++++
 test/unit/StorageTest.cpp |  3 +++
 3 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/src/Storage.cpp b/src/Storage.cpp
index 6610674fa5..8072583e04 100644
--- a/src/Storage.cpp
+++ b/src/Storage.cpp
@@ -17,7 +17,9 @@ using namespace nt;
 
 std::unique_ptr<Storage> Storage::m_instance;
 
-Storage::Storage() {}
+Storage::Storage() {
+  m_updates_enabled = false;
+}
 
 Storage::~Storage() {}
 
@@ -50,9 +52,9 @@ bool Storage::SetEntryValue(StringRef name, std::shared_ptr<Value> value) {
   entry->set_value(value);
   // put on update queue
   if (!old_value)
-    m_updates.push(Update{entry, Update::kAssign});
+    AddUpdate(entry, Update::kAssign);
   else if (*old_value != *value)
-    m_updates.push(Update{entry, Update::kValueUpdate});
+    AddUpdate(entry, Update::kValueUpdate);
   return true;
 }
 
@@ -65,9 +67,9 @@ void Storage::SetEntryTypeValue(StringRef name, std::shared_ptr<Value> value) {
   if (!old_value || *old_value != *value) {
     // put on update queue
     if (!old_value || old_value->type() != value->type())
-      m_updates.push(Update{entry, Update::kAssign});
+      AddUpdate(entry, Update::kAssign);
     else
-      m_updates.push(Update{entry, Update::kValueUpdate});
+      AddUpdate(entry, Update::kValueUpdate);
   }
 }
 
@@ -77,7 +79,7 @@ void Storage::SetEntryFlags(StringRef name, unsigned int flags) {
   if (!entry) return;
   if (entry->flags() != flags) {
     entry->set_flags(flags);
-    m_updates.push(Update{entry, Update::kFlagsUpdate});  // put on update queue
+    AddUpdate(entry, Update::kFlagsUpdate);  // put on update queue
   }
 }
 
@@ -94,14 +96,14 @@ void Storage::DeleteEntry(StringRef name) {
   auto entry = i->getValue();
   m_entries.erase(i);  // erase from map
   // if it had a value, put on update queue
-  if (entry->value()) m_updates.push(Update{entry, Update::kDelete});
+  if (entry->value()) AddUpdate(entry, Update::kDelete);
 }
 
 void Storage::DeleteAllEntries() {
   std::lock_guard<std::mutex> lock(m_mutex);
   if (m_entries.empty()) return;
   m_entries.clear();
-  m_updates.push(Update{nullptr, Update::kDeleteAll});  // put on update queue
+  AddUpdate(nullptr, Update::kDeleteAll);  // put on update queue
 }
 
 std::vector<EntryInfo> Storage::GetEntryInfo(StringRef prefix,
@@ -558,12 +560,12 @@ next_line:
 
       // put on update queue
       if (!old_value || old_value->type() != i.second->type())
-        m_updates.push(Update{entry, Update::kAssign});
+        AddUpdate(entry, Update::kAssign);
       else {
         if (*old_value != *i.second)
-          m_updates.push(Update{entry, Update::kValueUpdate});
+          AddUpdate(entry, Update::kValueUpdate);
         if (!was_persist)
-          m_updates.push(Update{entry, Update::kFlagsUpdate});
+          AddUpdate(entry, Update::kFlagsUpdate);
       }
     }
   }
diff --git a/src/Storage.h b/src/Storage.h
index 6447e4af2b..4f40d4de4b 100644
--- a/src/Storage.h
+++ b/src/Storage.h
@@ -91,8 +91,12 @@ class Storage {
   std::shared_ptr<StorageEntry> FindEntry(StringRef name) const;
   std::shared_ptr<StorageEntry> GetEntry(StringRef name);
 
+  // Accessors required by Dispatcher.
+  void EnableUpdates() { m_updates_enabled = true; }
+  void DisableUpdates() { m_updates_enabled = false; }
   UpdateQueue& updates() { return m_updates; }
 
+  // User functions
   std::shared_ptr<Value> GetEntryValue(StringRef name) const;
   bool SetEntryValue(StringRef name, std::shared_ptr<Value> value);
   void SetEntryTypeValue(StringRef name, std::shared_ptr<Value> value);
@@ -112,11 +116,17 @@ class Storage {
   Storage(const Storage&) = delete;
   Storage& operator=(const Storage&) = delete;
 
+  void AddUpdate(std::shared_ptr<StorageEntry> entry, Update::Kind kind) {
+    if (m_updates_enabled)
+      m_updates.push(Update{entry, kind});
+  }
+
   typedef llvm::StringMap<std::shared_ptr<StorageEntry>> EntriesMap;
 
   mutable std::mutex m_mutex;
   EntriesMap m_entries;
   UpdateQueue m_updates;
+  std::atomic_bool m_updates_enabled;
 
   static std::unique_ptr<Storage> m_instance;
 };
diff --git a/test/unit/StorageTest.cpp b/test/unit/StorageTest.cpp
index 0e1b8fdc31..29c9818c5d 100644
--- a/test/unit/StorageTest.cpp
+++ b/test/unit/StorageTest.cpp
@@ -16,6 +16,9 @@ namespace nt {
 
 class StorageTest : public ::testing::Test {
  public:
+  StorageTest() {
+    storage.EnableUpdates();
+  }
   Storage::EntriesMap& entries() { return storage.m_entries; }
   Storage::UpdateQueue& updates() { return storage.updates(); }
   Storage storage;