From e5c804416c135b966545b3e847d53d1ed0514119 Mon Sep 17 00:00:00 2001
From: Tao He <linzhu.ht@alibaba-inc.com>
Date: Tue, 12 Jan 2021 01:11:27 +0800
Subject: [PATCH] Protect implicit keepalive maps using a lexical scoped lock.

Signed-off-by: Tao He <linzhu.ht@alibaba-inc.com>
---
 etcd/Client.hpp |  2 ++
 src/Client.cpp  | 36 ++++++++++++++++++++++--------------
 2 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/etcd/Client.hpp b/etcd/Client.hpp
index abed256..11e7ceb 100644
--- a/etcd/Client.hpp
+++ b/etcd/Client.hpp
@@ -4,6 +4,7 @@
 #include "etcd/Response.hpp"
 
 #include <map>
+#include <mutex>
 #include <string>
 
 #include <grpc++/grpc++.h>
@@ -262,6 +263,7 @@ namespace etcd
     std::unique_ptr<Lease::Stub> leaseServiceStub;
     std::unique_ptr<Lock::Stub> lockServiceStub;
 
+    std::mutex mutex_for_keepalives;
     std::map<std::string, int64_t> leases_for_locks;
     std::map<int64_t, std::shared_ptr<KeepAlive>> keep_alive_for_locks;
 
diff --git a/src/Client.cpp b/src/Client.cpp
index 8245c0d..a98841d 100644
--- a/src/Client.cpp
+++ b/src/Client.cpp
@@ -523,7 +523,11 @@ pplx::task<etcd::Response> etcd::Client::lock(std::string const &key) {
   // (base on our experiences in vineyard and GraphScope).
   auto resp = this->leasegrant(DEFAULT_LEASE_TTL_FOR_LOCK).get();
   int64_t lease_id = resp.value().lease();
-  this->keep_alive_for_locks[lease_id].reset(new KeepAlive(*this, DEFAULT_LEASE_TTL_FOR_LOCK, lease_id));
+  {
+    std::lock_guard<std::mutex> lexical_scope_lock(mutex_for_keepalives);
+    this->keep_alive_for_locks[lease_id].reset(
+      new KeepAlive(*this, DEFAULT_LEASE_TTL_FOR_LOCK, lease_id));
+  }
   params.key = key;
   params.lease_id = lease_id;
   params.lock_stub = lockServiceStub.get();
@@ -531,10 +535,13 @@ pplx::task<etcd::Response> etcd::Client::lock(std::string const &key) {
   return Response::create(call).then(
     [this, lease_id](pplx::task<etcd::Response> const &resp_task) -> etcd::Response {
       auto const& resp = resp_task.get();
-      if (resp.is_ok()) {
-        this->leases_for_locks[resp.lock_key()] = lease_id;
-      } else {
-        this->keep_alive_for_locks.erase(lease_id);
+      {
+        std::lock_guard<std::mutex> lexical_scope_lock(mutex_for_keepalives);
+        if (resp.is_ok()) {
+          this->leases_for_locks[resp.lock_key()] = lease_id;
+        } else {
+          this->keep_alive_for_locks.erase(lease_id);
+        }
       }
       return resp;
     }
@@ -555,16 +562,17 @@ pplx::task<etcd::Response> etcd::Client::lock(std::string const &key,
 pplx::task<etcd::Response> etcd::Client::unlock(std::string const &lock_key) {
   std::cout << "begin unlock" << std::endl;
   // cancel the KeepAlive first, it exists
-  auto p_leases = this->leases_for_locks.find(lock_key);
-  if (p_leases != this->leases_for_locks.end()) {
-    std::cout << "Unlock for " << lock_key << " and revoke lease " << std::hex << p_leases->second << std::dec << std::endl;
-    auto p_keeps_alive = this->keep_alive_for_locks.find(p_leases->second);
-    if (p_keeps_alive != this->keep_alive_for_locks.end()) {
-      this->keep_alive_for_locks.erase(p_keeps_alive);
+  {
+    std::lock_guard<std::mutex> lexical_scope_lock(mutex_for_keepalives);
+    auto p_leases = this->leases_for_locks.find(lock_key);
+    if (p_leases != this->leases_for_locks.end()) {
+      std::cout << "Unlock for " << lock_key << " and revoke lease " << std::hex << p_leases->second << std::dec << std::endl;
+      auto p_keeps_alive = this->keep_alive_for_locks.find(p_leases->second);
+      if (p_keeps_alive != this->keep_alive_for_locks.end()) {
+        this->keep_alive_for_locks.erase(p_keeps_alive);
+      }
+      this->leases_for_locks.erase(p_leases);
     }
-    this->leases_for_locks.erase(p_leases);
-  } else {
-    std::cout << "Unable to find lease_id for " << lock_key;
   }
 
   // issue a "unlock" request