Repository: tikv/tikv
File: src/storage/txn/commands/acquire_pessimistic_lock.rs (Rust)

Code before:
            .iter()
            .map(|(key, _)| key.as_encoded().len())
            .sum()
    }

    gen_lock!(keys: multiple(|x| &x.0));
}

impl<S: Snapshot, L: LockManager> WriteCommand<S, L> for AcquirePessimisticLock {
    fn process_write(self, snapshot: S, context: WriteContext<'_, L>) -> Result<WriteResult> {
        if self.allow_lock_with_conflict && self.keys.len() > 1 {
            // Currently multiple keys with `allow_lock_with_conflict` set is not supported.
            return Err(Error::from(ErrorInner::Other(box_err!(
                "multiple keys in a single request with allowed_lock_with_conflict set is not allowed"
            ))));
        }

        let (start_ts, ctx, mut keys) = (self.start_ts, self.ctx, self.keys);
        let mut txn = MvccTxn::new(start_ts, context.concurrency_manager);

        let mut snapshot_reader = SnapshotReader::new_with_ctx(start_ts, snapshot, &ctx);
        if keys.len() > 1 {
            keys.sort();
            snapshot_reader.set_scan_mode(ScanMode::Forward);
            snapshot_reader.set_lower_bound(keys.first().unwrap().0.clone());
        }
        let mut reader = ReaderWithStats::new(snapshot_reader, context.statistics);

        let total_keys = keys.len();
        let mut res = PessimisticLockResults::with_capacity(total_keys);
        let mut encountered_locks = vec![];
        let need_old_value = context.extra_op == ExtraOp::ReadOldValue;
        let mut old_values = OldValues::default();
        for (k, should_not_exist) in keys {
            match acquire_pessimistic_lock(
                &mut txn,
                &mut reader,
                k.clone(),
                &self.primary,
                should_not_exist,
                self.lock_ttl,
                self.for_update_ts,
                self.return_values,
                self.check_existence,
                self.min_commit_ts,
                need_old_value,
                self.lock_only_if_exists,
                self.allow_lock_with_conflict,
            ) {
                Ok((key_res, old_value)) => {
                    res.push(key_res);

Review comment (refactor):
Can we wrap the logic into `SnapshotReader::new_with_ctx` like
```
SnapshotReader::new_with_ctx(start_ts, snapshot, &ctx, &keys);
```
, so the duplicated code could be avoided.

Code after:
            .iter()
            .map(|(key, _)| key.as_encoded().len())
            .sum()
    }

    gen_lock!(keys: multiple(|x| &x.0));
}

impl<S: Snapshot, L: LockManager> WriteCommand<S, L> for AcquirePessimisticLock {
    fn process_write(self, snapshot: S, context: WriteContext<'_, L>) -> Result<WriteResult> {
        if self.allow_lock_with_conflict && self.keys.len() > 1 {
            // Currently multiple keys with `allow_lock_with_conflict` set is not supported.
            return Err(Error::from(ErrorInner::Other(box_err!(
                "multiple keys in a single request with allowed_lock_with_conflict set is not allowed"
            ))));
        }

        let (start_ts, ctx, mut keys) = (self.start_ts, self.ctx, self.keys);
        let mut txn = MvccTxn::new(start_ts, context.concurrency_manager);

        let mut snapshot_reader = SnapshotReader::new_with_ctx(start_ts, snapshot, &ctx);
        snapshot_reader.setup_with_hint_items(&mut keys, |k| &k.0);
        let mut reader = ReaderWithStats::new(snapshot_reader, context.statistics);

        let total_keys = keys.len();
        let mut res = PessimisticLockResults::with_capacity(total_keys);
        let mut encountered_locks = vec![];
        let need_old_value = context.extra_op == ExtraOp::ReadOldValue;
        let mut old_values = OldValues::default();
        for (k, should_not_exist) in keys {
            match acquire_pessimistic_lock(
                &mut txn,
                &mut reader,
                k.clone(),
                &self.primary,
                should_not_exist,
                self.lock_ttl,
                self.for_update_ts,
                self.return_values,
                self.check_existence,
                self.min_commit_ts,
                need_old_value,
                self.lock_only_if_exists,
                self.allow_lock_with_conflict,
            ) {
                Ok((key_res, old_value)) => {
                    res.push(key_res);
                    // MutationType is unknown in AcquirePessimisticLock stage.
                    insert_old_value_if_resolved(&mut old_values, k, txn.start_ts, old_value, None);
                }
                Err(MvccError(box MvccErrorInner::KeyIsLocked(lock_info))) => {