RocksDB上锁机制的实例详解

RocksDB作为一个开源的存储引擎支持事务的ACID特性，而要支持ACID中的I(Isolation),并发控制这块是少不了的，本文主要讨论RocksDB的锁机制实现，细节会涉及到源码分析，希望通过本文读者可以深入了解RocksDB并发控制原理。文章主要从以下4方面展开，首先会介绍RocksDB锁的基本结构，然后我会介绍RocksDB行锁数据结构设计下，锁空间开销，接着我会介绍几种典型场景的上锁流程，最后会介绍锁机制中必不可少的死锁检测机制。

1.行锁数据结构
RocksDB锁粒度最小是行，对于KV存储而言，锁对象就是key，每一个key对应一个LockInfo结构。所有key通过hash表管理，查找锁时，直接通过hash表定位即可确定这个key是否已经被上锁。但如果全局只有一个hash表，会导致这个访问这个hash表的冲突很多，影响并发性能。RocksDB首先按Columnfamily进行拆分，每个Columnfamily中的锁通过一个LockMap管理，而每个LockMap再拆分成若干个分片，每个分片通过LockMapStripe管理，而hash表(std::unordered_map<std::string, LockInfo>)则存在于Stripe结构中，Stripe结构中还包含一个mutex和condition_variable，这个主要作用是，互斥访问hash表，当出现锁冲突时，将线程挂起，解锁后，唤醒挂起的线程。这种设计很简单但也带来一个显而易见的问题，就是多个不相关的锁公用一个condition_variable，导致锁释放时，不必要的唤醒一批线程，而这些线程重试后，发现仍然需要等待，造成了无效的上下文切换。对比我们之前讨论的InnoDB锁机制，我们发现InnoDB是一个page里面的记录复用一把锁，而且复用是有条件的，同一个事务对一个page的若干条记录加锁才能复用；而且锁等待队列是精确等待，精确到记录级别，不会导致的无效的唤醒。虽然RocksDB锁设计比较粗糙，但也做了一定的优化，比如在管理LockMaps时，通过在每个线程本地缓存一份拷贝lock_maps_cache_，通过全局链表将每个线程的cache链起来，当LockMaps变更时(删除columnfamily)，则全局将每个线程的copy清空，由于columnfamily改动很少，所以大部分访问LockMaps操作都是不需要加锁的，提高了并发效率。
相关数据结构如下：

struct LockInfo { bool exclusive; //排它锁或是共享锁 autovector<TransactionID> txn_ids; //事务列表，对于共享锁而言，同一个key可以对应多个事务 // Transaction locks are not valid after this time in us uint64_t expiration_time; } struct LockMapStripe { // Mutex must be held before modifying keys map std::shared_ptr<TransactionDBMutex> stripe_mutex; // Condition Variable per stripe for waiting on a lock std::shared_ptr<TransactionDBCondVar> stripe_cv; // Locked keys mapped to the info about the transactions that locked them. std::unordered_map<std::string, LockInfo> keys; } struct LockMap { const size_t num_stripes_; //分片个数 std::atomic<int64_t> lock_cnt{0}; //锁数目 std::vector<LockMapStripe*> lock_map_stripes_; //锁分片 } class TransactionLockMgr { using LockMaps = std::unordered_map<uint32_t, std::shared_ptr<LockMap>>; LockMaps lock_maps_; // Thread-local cache of entries in lock_maps_. This is an optimization // to avoid acquiring a mutex in order to look up a LockMap std::unique_ptr<ThreadLocalPtr> lock_maps_cache_; }

2.行锁空间代价
由于锁信息是常驻内存，我们简单分析下RocksDB锁占用的内存。每个锁实际上是unordered_map中的一个元素，则锁占用的内存为key_length+8+8+1，假设key为bigint，占8个字节，则100w行记录，需要消耗大约22M内存。但是由于内存与key_length正相关，导致RocksDB的内存消耗不可控。我们可以简单算算RocksDB作为MySQL存储引擎时，key_length的范围。对于单列索引，最大值为2048个字节，具体可以参考max_supported_key_part_length实现；对于复合索引，索引最大长度为3072个字节，具体可以参考max_supported_key_length实现。假设最坏的情况，key_length=3072，则100w行记录，需要消耗3G内存，如果是锁1亿行记录，则需要消耗300G内存，这种情况下内存会有撑爆的风险。因此RocksDB提供参数配置max_row_locks，确保内存可控，默认RDB_MAX_ROW_LOCKS设置为1G，对于大部分key为bigint场景，极端情况下，也需要消耗22G内存。而在这方面，InnoDB则比较友好，hash表的key是(space_id, page_no)，所以无论key有多大，key部分的内存消耗都是恒定的。前面我也提到了InnoDB在一个事务需要锁大量记录场景下是有优化的，多个记录可以公用一把锁，这样也间接可以减少内存。

3.上锁流程分析
前面简单了解了RocksDB锁数据结构的设计以及锁对内存资源的消耗。这节主要介绍几种典型场景下，RocksDB是如何加锁的。与InnoDB一样，RocksDB也支持MVCC，读不上锁，为了方便，下面的讨论基于RocksDB作为MySQL的一个引擎来展开，主要包括三类，基于主键的更新，基于二级索引的更新，基于主键的范围更新等。在展开讨论之前，有一点需要说明的是，RocksDB与InnoDB不同，RocksDB的更新也是基于快照的，而InnoDB的更新基于当前读，这种差异也使得在实际应用中，相同隔离级别下，表现有所不一样。对于RocksDB而言，在RC隔离级别下，每个语句开始都会重新获取一次快照；在RR隔离级别下，整个事务中只在第一个语句开始时获取一次快照，所有语句共用这个快照，直到事务结束。

3.1.基于主键的更新
这里主要接口是TransactionBaseImpl::GetForUpdate
1).尝试对key加锁，如果锁被其它事务持有，则需要等待
2).创建snapshot
3).调用ValidateSnapshot，Get key,通过比较Sequence判断key是否被更新过
4).由于是加锁后，再获取snapshot，所以检查一定成功。
5).执行更新操作
这里有一个延迟获取快照的机制，实际上在语句开始时，需要调用acquire_snapshot获取快照，但为了避免冲突导致的重试，在对key加锁后，再获取snapshot，这就保证了在基于主键更新的场景下，不会存在ValidateSnapshot失败的场景。

堆栈如下：

1-myrocks::ha_rocksdb::get_row_by_rowid 2-myrocks::ha_rocksdb::get_for_update 3-myrocks::Rdb_transaction_impl::get_for_update 4-rocksdb::TransactionBaseImpl::GetForUpdate { //加锁 5-rocksdb::TransactionImpl::TryLock 6-rocksdb::TransactionDBImpl::TryLock 7-rocksdb::TransactionLockMgr::TryLock //延迟获取快照，与acquire_snapshot配合使用 6-SetSnapshotIfNeeded() //检查key对应快照是否过期 6-ValidateSnapshot 7-rocksdb::TransactionUtil::CheckKeyForConflict 8-rocksdb::TransactionUtil::CheckKey 9-rocksdb::DBImpl::GetLatestSequenceForKey //第一次读取 //读取key 5-rocksdb::TransactionBaseImpl::Get 6-rocksdb::WriteBatchWithIndex::GetFromBatchAndDB 7-rocksdb::DB::Get 8-rocksdb::DBImpl::Get 9-rocksdb::DBImpl::GetImpl //第二次读取 }

3.2.基于主键的范围更新
1).创建Snapshot，基于迭代器扫描主键
2).通过get_row_by_rowid，尝试对key加锁
3).调用ValidateSnapshot，Get key,通过比较Sequence判断key是否被更新过
4).如果key被其它事务更新过(key对应的SequenceNumber比Snapshot要新)，触发重试
5).重试情况下，会释放老的快照并释放锁，通过tx->acquire_snapshot(false)，延迟获取快照(加锁后，再拿snapshot)
5).再次调用get_for_update，由于此时key已经被加锁，重试一定可以成功。
6).执行更新操作
7).跳转到1，继续执行，直到主键不符合条件时，则结束。

3.3.基于二级索引的更新
这种场景与3.2类似，只不过多一步从二级索引定位主键过程。
1).创建Snapshot，基于迭代器扫描二级索引
2).根据二级索引反向找到主键，实际上也是调用get_row_by_rowid，这个过程就会尝试对key加锁
3).继续根据二级索引遍历下一个主键，尝试加锁
4).当返回的