tikv/resolvelock: resolve lock in a batch

src/server/service/kv.rs

@@ -12,6 +12,7 @@
 // limitations under the License.
 
 use std::boxed::FnBox;
+use std::collections::HashMap;
 use std::fmt::Debug;
 use std::io::Write;
 use std::sync::Arc;
@@ -475,14 +476,21 @@ impl<T: RaftStoreRouter + 'static> tikvpb_grpc::Tikv for Service<T> {
             .with_label_values(&[label])
             .start_coarse_timer();
 
-        let commit_ts = match req.get_commit_version() {
-            0 => None,
-            x => Some(x),
-        };
+        let mut txn_2_status: HashMap<u64, u64> = HashMap::new();
+
+        let start_ts = req.get_start_version();
+        if start_ts > 0 {
+            let commit_ts = req.get_commit_version();
+            txn_2_status.insert(start_ts, commit_ts);
+        } else {
+            for temp in req.take_txn_infos().into_iter() {
+                txn_2_status.insert(temp.txn, temp.status);
+            }
+        }
 
         let (cb, future) = make_callback();
         let res = self.storage
-            .async_resolve_lock(req.take_context(), req.get_start_version(), commit_ts, cb);
+            .async_resolve_lock(req.take_context(), txn_2_status, cb);
         if let Err(e) = res {
             self.send_fail_status(ctx, sink, Error::from(e), RpcStatusCode::ResourceExhausted);
             return;
@@ -507,6 +515,7 @@ impl<T: RaftStoreRouter + 'static> tikvpb_grpc::Tikv for Service<T> {
             });
 
         ctx.spawn(future);
+
     }
 
     fn kv_gc(&self, ctx: RpcContext, mut req: GCRequest, sink: UnarySink<GCResponse>) {

src/storage/mod.rs

@@ -12,6 +12,7 @@
 // limitations under the License.
 
 use std::thread;
+use std::collections::HashMap;
 use std::boxed::FnBox;
 use std::fmt::{self, Debug, Display, Formatter};
 use std::sync::mpsc::{self, Receiver};
@@ -35,6 +36,7 @@ pub use self::engine::{new_local_engine, CFStatistics, Cursor, Engine, Error as
                        FlowStatistics, Modify, ScanMode, Snapshot, Statistics, StatisticsSummary,
                        TEMP_DIR};
 pub use self::engine::raftkv::RaftKv;
+use self::mvcc::Lock;
 pub use self::txn::{Msg, Scheduler, SnapshotStore, StoreScanner};
 pub use self::types::{make_key, Key, KvPair, MvccInfo, Value};
 pub type Callback<T> = Box<FnBox(Result<T>) + Send>;
@@ -131,10 +133,9 @@ pub enum Command {
     ScanLock { ctx: Context, max_ts: u64 },
     ResolveLock {
         ctx: Context,
-        start_ts: u64,
-        commit_ts: Option<u64>,
+        txn2status: HashMap<u64, u64>,
         scan_key: Option<Key>,
-        keys: Vec<Key>,
+        key_locks: Vec<(Key, Lock)>,
     },
     Gc {
         ctx: Context,
@@ -241,18 +242,7 @@ impl Display for Command {
             Command::ScanLock {
                 ref ctx, max_ts, ..
             } => write!(f, "kv::scan_lock {} | {:?}", max_ts, ctx),
-            Command::ResolveLock {
-                ref ctx,
-                start_ts,
-                commit_ts,
-                ..
-            } => write!(
-                f,
-                "kv::resolve_txn {} -> {:?} | {:?}",
-                start_ts,
-                commit_ts,
-                ctx
-            ),
+            Command::ResolveLock { .. } => write!(f, "kv::resolve_lock"),
             Command::Gc {
                 ref ctx,
                 safe_point,
@@ -328,7 +318,7 @@ impl Command {
             Command::Pause { .. } |
             Command::MvccByKey { .. } |
             Command::MvccByStartTs { .. } => true,
-            Command::ResolveLock { ref keys, .. } |
+            Command::ResolveLock { ref key_locks, .. } => key_locks.is_empty(),
             Command::Gc { ref keys, .. } => keys.is_empty(),
             _ => false,
         }
@@ -379,11 +369,11 @@ impl Command {
             Command::Prewrite { start_ts, .. } |
             Command::Cleanup { start_ts, .. } |
             Command::Rollback { start_ts, .. } |
-            Command::ResolveLock { start_ts, .. } |
             Command::MvccByStartTs { start_ts, .. } => start_ts,
             Command::Commit { lock_ts, .. } => lock_ts,
             Command::ScanLock { max_ts, .. } => max_ts,
             Command::Gc { safe_point, .. } => safe_point,
+            Command::ResolveLock { .. } |
             Command::RawGet { .. } |
             Command::RawScan { .. } |
             Command::DeleteRange { .. } |
@@ -449,8 +439,8 @@ impl Command {
             Command::Gc { ref keys, .. } |
             Command::BatchGet { ref keys, .. } |
             Command::Commit { ref keys, .. } |
-            Command::Rollback { ref keys, .. } |
-            Command::ResolveLock { ref keys, .. } => keys.len(),
+            Command::Rollback { ref keys, .. } => keys.len(),
+            Command::ResolveLock { ref key_locks, .. } => key_locks.len(),
             Command::Prewrite { ref mutations, .. } => mutations.len(),
         }
     }
@@ -768,16 +758,14 @@ impl Storage {
     pub fn async_resolve_lock(
         &self,
         ctx: Context,
-        start_ts: u64,
-        commit_ts: Option<u64>,
+        txn2status: HashMap<u64, u64>,
         callback: Callback<()>,
     ) -> Result<()> {
         let cmd = Command::ResolveLock {
             ctx: ctx,
-            start_ts: start_ts,
-            commit_ts: commit_ts,
+            txn2status: txn2status,
             scan_key: None,
-            keys: vec![],
+            key_locks: vec![],
         };
         let tag = cmd.tag();
         self.send(cmd, StorageCb::Boolean(callback))?;

src/storage/txn/scheduler.rs

@@ -31,6 +31,7 @@
 //! is ensured by the transaction protocol implemented in the client library, which is transparent
 //! to the scheduler.
 
+use std::collections::HashMap;
 use std::fmt::{self, Debug, Formatter};
 use std::sync::mpsc::Receiver;
 use std::time::Duration;
@@ -52,7 +53,6 @@ use raftstore::store::engine::IterOption;
 use util::transport::{Error as TransportError, SyncSendCh};
 use util::threadpool::{Context as ThreadContext, ThreadPool, ThreadPoolBuilder};
 use util::time::SlowTimer;
-use util::collections::HashMap;
 
 use super::Result;
 use super::Error;
@@ -585,12 +585,9 @@ fn process_read(
                 Err(e) => ProcessResult::Failed { err: e.into() },
             }
         }
-        // Scan the locks with timestamp `start_ts`, then either commit them if the command has
-        // commit timestamp populated or rollback otherwise.
         Command::ResolveLock {
             ref ctx,
-            start_ts,
-            commit_ts,
+            ref txn2status,
             ref mut scan_key,
             ..
         } => {
@@ -605,24 +602,23 @@ fn process_read(
             let res = reader
                 .scan_lock(
                     scan_key.take(),
-                    |lock| lock.ts == start_ts,
+                    |lock| txn2status.contains_key(&lock.ts),
                     Some(RESOLVE_LOCK_BATCH_SIZE),
                 )
                 .map_err(Error::from)
                 .and_then(|(v, next_scan_key)| {
-                    let keys: Vec<Key> = v.into_iter().map(|x| x.0).collect();
+                    let key_locks: Vec<_> = v.into_iter().map(|x| x).collect();
                     KV_COMMAND_KEYREAD_HISTOGRAM_VEC
                         .with_label_values(&[tag])
-                        .observe(keys.len() as f64);
-                    if keys.is_empty() {
+                        .observe(key_locks.len() as f64);
+                    if key_locks.is_empty() {
                         Ok(None)
                     } else {
                         Ok(Some(Command::ResolveLock {
                             ctx: ctx.clone(),
-                            start_ts: start_ts,
-                            commit_ts: commit_ts,
+                            txn2status: txn2status.clone(),
                             scan_key: next_scan_key,
-                            keys: keys,
+                            key_locks: key_locks,
                         }))
                     }
                 });
@@ -884,52 +880,56 @@ fn process_write_impl(
         }
         Command::ResolveLock {
             ref ctx,
-            start_ts,
-            commit_ts,
+            ref txn2status,
             ref mut scan_key,
-            ref keys,
+            ref key_locks,
         } => {
-            if let Some(cts) = commit_ts {
-                if cts <= start_ts {
-                    return Err(Error::InvalidTxnTso {
-                        start_ts: start_ts,
-                        commit_ts: cts,
-                    });
-                }
-            }
             let mut scan_key = scan_key.take();
-            let mut txn = MvccTxn::new(
-                snapshot,
-                statistics,
-                start_ts,
-                None,
-                ctx.get_isolation_level(),
-                !ctx.get_not_fill_cache(),
-            );
-            let rows = keys.len();
-            for k in keys {
-                match commit_ts {
-                    Some(ts) => txn.commit(k, ts)?,
-                    None => txn.rollback(k)?,
+            let mut modifies: Vec<Modify> = vec![];
+            let rows = key_locks.len();
+            for &(ref current_key, ref current_lock) in key_locks {
+                let mut txn = MvccTxn::new(
+                    snapshot,
+                    statistics,
+                    current_lock.ts,
+                    None,
+                    ctx.get_isolation_level(),
+                    !ctx.get_not_fill_cache(),
+                );
+                let status = txn2status.get(&current_lock.ts);
+                let ts = match status {
+                    Some(ts) => *ts,
+                    None => panic!("txn status not found!"),
+                };
+                if ts > 0 {
+                    if current_lock.ts >= ts {
+                        return Err(Error::InvalidTxnTso {
+                            start_ts: current_lock.ts,
+                            commit_ts: ts,
+                        });
+                    }
+                    txn.commit(current_key, ts)?;
+                } else {
+                    txn.rollback(current_key)?;
                 }
-                if txn.write_size() >= MAX_TXN_WRITE_SIZE {
-                    scan_key = Some(k.to_owned());
+                modifies.append(&mut txn.modifies());
+                if modifies.len() >= MAX_TXN_WRITE_SIZE {
+                    scan_key = Some(current_key.to_owned());
                     break;
                 }
             }
             if scan_key.is_none() {
-                (ProcessResult::Res, txn.modifies(), rows)
+                (ProcessResult::Res, modifies, rows)
             } else {
                 let pr = ProcessResult::NextCommand {
                     cmd: Command::ResolveLock {
                         ctx: ctx.clone(),
-                        start_ts: start_ts,
-                        commit_ts: commit_ts,
+                        txn2status: txn2status.clone(),
                         scan_key: scan_key.take(),
-                        keys: vec![],
+                        key_locks: vec![],
                     },
                 };
-                (pr, txn.modifies(), rows)
+                (pr, modifies, rows)
             }
         }
         Command::Gc {
@@ -1375,6 +1375,7 @@ impl Scheduler {
             .with_label_values(&[self.get_ctx_tag(cid), "write"])
             .inc();
         if to_be_write.is_empty() {
+            info!("modifies is empty");
             return self.on_write_finished(cid, pr, Ok(()));
         }
         let engine_cb = make_engine_cb(cmd.tag(), cid, pr, self.schedch.clone(), rows);
@@ -1521,9 +1522,13 @@ pub fn gen_command_lock(latches: &Latches, cmd: &Command) -> Lock {
             let keys: Vec<&Key> = mutations.iter().map(|x| x.key()).collect();
             latches.gen_lock(&keys)
         }
-        Command::Commit { ref keys, .. } |
-        Command::Rollback { ref keys, .. } |
-        Command::ResolveLock { ref keys, .. } => latches.gen_lock(keys),
+        Command::ResolveLock { ref key_locks, .. } => {
+            let keys: Vec<&Key> = key_locks.iter().map(|x| &x.0).collect();
+            latches.gen_lock(&keys)
+        }
+        Command::Commit { ref keys, .. } | Command::Rollback { ref keys, .. } => {
+            latches.gen_lock(keys)
+        }
         Command::Cleanup { ref key, .. } => latches.gen_lock(&[key]),
         _ => Lock::new(vec![]),
     }
@@ -1532,12 +1537,16 @@ pub fn gen_command_lock(latches: &Latches, cmd: &Command) -> Lock {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use std::collections::HashMap;
     use kvproto::kvrpcpb::Context;
     use storage::txn::latch::*;
     use storage::{make_key, Command, Mutation, Options};
+    use storage::mvcc;
 
     #[test]
     fn test_command_latches() {
+        let mut temp_map = HashMap::new();
+        temp_map.insert(10, 20);
         let readonly_cmds = vec![
             Command::Get {
                 ctx: Context::new(),
@@ -1562,10 +1571,9 @@ mod tests {
             },
             Command::ResolveLock {
                 ctx: Context::new(),
-                start_ts: 10,
-                commit_ts: Some(20),
+                txn2status: temp_map.clone(),
                 scan_key: None,
-                keys: vec![],
+                key_locks: vec![],
             },
             Command::Gc {
                 ctx: Context::new(),
@@ -1609,10 +1617,14 @@ mod tests {
             },
             Command::ResolveLock {
                 ctx: Context::new(),
-                start_ts: 10,
-                commit_ts: Some(20),
+                txn2status: temp_map.clone(),
                 scan_key: None,
-                keys: vec![make_key(b"k")],
+                key_locks: vec![
+                    (
+                        make_key(b"k"),
+                        mvcc::Lock::new(mvcc::LockType::Put, b"k".to_vec(), 10, 20, None),
+                    ),
+                ],
             },
         ];
 

tests/raftstore/test_service.rs

@@ -332,6 +332,8 @@ fn test_mvcc_rollback_and_cleanup() {
 
 #[test]
 fn test_mvcc_resolve_lock_gc_and_delete() {
+    use kvproto::kvrpcpb::*;
+    use protobuf;
     let (_cluster, client, ctx) = must_new_cluster_and_kv_client();
     let (k, v) = (b"key".to_vec(), b"value".to_vec());
 
@@ -388,9 +390,12 @@ fn test_mvcc_resolve_lock_gc_and_delete() {
     ts += 1;
     let resolve_lock_commit_version = ts;
     let mut resolve_lock_req = ResolveLockRequest::new();
+    let mut temp_txninfo = TxnInfo::new();
+    temp_txninfo.txn = prewrite_start_version2;
+    temp_txninfo.status = resolve_lock_commit_version;
+    let vec_txninfo = vec![temp_txninfo];
     resolve_lock_req.set_context(ctx.clone());
-    resolve_lock_req.start_version = prewrite_start_version2;
-    resolve_lock_req.commit_version = resolve_lock_commit_version;
+    resolve_lock_req.set_txn_infos(protobuf::RepeatedField::from_vec(vec_txninfo));
     let resolve_lock_resp = client.kv_resolve_lock(resolve_lock_req).unwrap();
     assert!(!resolve_lock_resp.has_region_error());
     assert!(!resolve_lock_resp.has_error());