From 5a51fa907b54a275c2206fd534a376b62d189b2b Mon Sep 17 00:00:00 2001
From: agiardullo <agiardullo@fb.com>
Date: Tue, 29 Sep 2015 17:00:16 -0700
Subject: [PATCH] Fix accidental object copy in transactions

Summary: Should have used auto& instead of auto.  Also needed to change the code a bit due to const correctness.

Test Plan: unit tests

Reviewers: sdong, igor, yoshinorim, yhchiang

Reviewed By: yhchiang

Subscribers: dhruba, leveldb

Differential Revision: https://reviews.facebook.net/D47787
---
 utilities/transactions/transaction_base.cc |  2 +-
 utilities/transactions/transaction_impl.cc | 28 ++++++++++++++--------
 2 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/utilities/transactions/transaction_base.cc b/utilities/transactions/transaction_base.cc
index 3a5df96337..dc91677355 100644
--- a/utilities/transactions/transaction_base.cc
+++ b/utilities/transactions/transaction_base.cc
@@ -77,7 +77,7 @@ Status TransactionBaseImpl::RollbackToSavePoint() {
     assert(s.ok());
 
     // Rollback any keys that were tracked since the last savepoint
-    auto key_map = GetTrackedKeysSinceSavePoint();
+    const TransactionKeyMap* key_map = GetTrackedKeysSinceSavePoint();
     assert(key_map);
     for (auto& key_map_iter : *key_map) {
       uint32_t column_family_id = key_map_iter.first;
diff --git a/utilities/transactions/transaction_impl.cc b/utilities/transactions/transaction_impl.cc
index 6202970bd0..c2a93cf33d 100644
--- a/utilities/transactions/transaction_impl.cc
+++ b/utilities/transactions/transaction_impl.cc
@@ -126,7 +126,7 @@ void TransactionImpl::Rollback() { Clear(); }
 
 Status TransactionImpl::RollbackToSavePoint() {
   // Unlock any keys locked since last transaction
-  auto keys = GetTrackedKeysSinceSavePoint();
+  const TransactionKeyMap* keys = GetTrackedKeysSinceSavePoint();
   if (keys) {
     txn_db_impl_->UnLock(this, keys);
   }
@@ -227,18 +227,26 @@ Status TransactionImpl::TryLock(ColumnFamilyHandle* column_family,
   // TODO(agiardullo): could optimize by supporting shared txn locks in the
   // future
   bool check_snapshot = !untracked;
+  SequenceNumber tracked_seqno = kMaxSequenceNumber;
+
+  // Lookup whether this key has already been locked by this transaction
+  const auto& tracked_keys = GetTrackedKeys();
+  const auto tracked_keys_cf = tracked_keys.find(cfh_id);
+  if (tracked_keys_cf == tracked_keys.end()) {
+    previously_locked = false;
+  } else {
+    auto iter = tracked_keys_cf->second.find(key_str);
+    if (iter == tracked_keys_cf->second.end()) {
+      previously_locked = false;
+    } else {
+      previously_locked = true;
+      tracked_seqno = iter->second;
+    }
+  }
 
   // lock this key if this transactions hasn't already locked it
-  SequenceNumber tracked_seqno = kMaxSequenceNumber;
-  auto tracked_keys = GetTrackedKeys();
-  auto iter = tracked_keys[cfh_id].find(key_str);
-  if (iter == tracked_keys[cfh_id].end()) {
-    previously_locked = false;
-
+  if (!previously_locked) {
     s = txn_db_impl_->TryLock(this, cfh_id, key_str);
-  } else {
-    previously_locked = true;
-    tracked_seqno = iter->second;
   }
 
   if (s.ok()) {