From a786470623c145041b20192c8d7a321ced7b6d92 Mon Sep 17 00:00:00 2001
From: Thad House <ThadHouse@users.noreply.github.com>
Date: Mon, 15 Aug 2016 22:46:49 -0700
Subject: [PATCH] Only allows 1 blocking call per Rpc Call Id (#93)

---
 src/Storage.cpp | 10 +++++++++-
 src/Storage.h   |  3 +++
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/src/Storage.cpp b/src/Storage.cpp
index c68430b39f..b43e7deff1 100644
--- a/src/Storage.cpp
+++ b/src/Storage.cpp
@@ -1420,6 +1420,8 @@ bool Storage::GetRpcResult(bool blocking, unsigned int call_uid, double time_out
         m_rpc_results.find(std::make_pair(call_uid >> 16, call_uid & 0xffff));
     if (i == m_rpc_results.end()) {
       if (!blocking || m_terminating) return false;
+      // only allow one blocking call per rpc call uid
+      if (!m_rpc_blocking_calls.insert(call_uid).second) return false;
       if (time_out < 0) {
         m_rpc_results_cond.wait(lock);
       } else {
@@ -1428,14 +1430,20 @@ bool Storage::GetRpcResult(bool blocking, unsigned int call_uid, double time_out
         while (!m_terminating) {
           auto timed_out = m_rpc_results_cond.wait_until(lock, timeout_time);
           if (timed_out == std::cv_status::timeout) {
+            m_rpc_blocking_calls.erase(call_uid);
             return false;
           }
         }
       }
-      if (m_terminating) return false;
+      if (m_terminating) {
+        m_rpc_blocking_calls.erase(call_uid);
+        return false;
+      }
       continue;
     }
     result->swap(i->getSecond());
+    // safe to erase even if id does not exist
+    m_rpc_blocking_calls.erase(call_uid);
     m_rpc_results.erase(i);
     return true;
   }
diff --git a/src/Storage.h b/src/Storage.h
index 46362f5824..6e9964b6d2 100644
--- a/src/Storage.h
+++ b/src/Storage.h
@@ -18,6 +18,7 @@
 #include <vector>
 
 #include "llvm/DenseMap.h"
+#include "llvm/SmallSet.h"
 #include "llvm/StringMap.h"
 #include "support/atomic_static.h"
 #include "Message.h"
@@ -143,11 +144,13 @@ class Storage {
   typedef std::vector<Entry*> IdMap;
   typedef llvm::DenseMap<std::pair<unsigned int, unsigned int>, std::string>
       RpcResultMap;
+  typedef llvm::SmallSet<unsigned int, 12> RpcBlockingCallSet;
 
   mutable std::mutex m_mutex;
   EntriesMap m_entries;
   IdMap m_idmap;
   RpcResultMap m_rpc_results;
+  RpcBlockingCallSet m_rpc_blocking_calls;
   // If any persistent values have changed
   mutable bool m_persistent_dirty = false;