MongoDB的upsert在并行程序中应该注意的问题

题图

Oct. 12th, 2018

最近使用python的多进程编程来处理一个MongoDB的数据库。这个程序中，每一个进程都要对数据库中的某些数据进行更新。如果目标数据不存在的话，则需要进行upsert。程序运行之后，速度确实比单进程快一些（在小数据集的情况下，大约快了三四倍），但是最终结果却与单进程产生的结果不符合。研究一番之后，把问题锁定在了多进程同时更新数据库时产生的upsert问题。这里讨论一个简化模型来阐明此问题并且论述解决方法。

简化模型

为了简化讨论，这里讨论一个投票程序。某城市人民选举该市市长。每个市民选择一个人为他/她投票。由于该选举为广泛选举，在选举之前并没有候选人集合，而是每个市民可以把票投给任何一个人。为简化问题，假设城市里没有人重名，因此每个人的名字可以作为其唯一代号使用。实际情况中，可以使用其他唯一表示，比如身份证号等等。

这里的数据集非常简单：

Collection: db.election

{'name': 'Alice', 'votes': 10}
{'name': 'Bob', 'votes': 21}
...

作为一个计票的进程，主要任务就是拿过一张选票，查看其name属性，在数据库中给名字为name的文档的票数加1。注意，这里name不一定已经存在于数据库中。如果此名字不存在，则应新建一条文档。因此，使用update_one时upsert属性应设为True。

 # name variable stores the candidate's name.
db.election.update_one({'name': name}, \
                {'$inc': {'votes': 1}}, upsert=True)

在只有一个进程运行的情况下，这段代码虽然速度并不快，但会给出正确的计票结果。如果我们使用多进程，创建几个worker，分别收集选票，给指定的被选人计票，会怎么样呢？

多进程下的写入矛盾

当简单地把上面的update_one交给几个进程来处理的时候，我们会发现运行结果出了这样的问题：每个被选人的票数似乎少了很多，而被选人的数量增加了。仔细检查会发现，其实是同一个候选人被创建了多个文档。为什么会导致这样的写入错误呢？不难想象到，这是多个进程同时试图更新一个文档的时候导致的。

需要注意的是，MongoDB本身是有文档级的写入锁的。也就是说，当一个进程开始修改一个文档时，该文档被锁定，其他文档不可以再对其进行写入甚至读取。这个写入锁的存在本身就是为了防止不同程序更新文档时产生的写入冲突。然而，update其实分为两步。首先是搜索文档位置，然后是文档更新。当两个程序同时试图更新一个不存在的文档的时候，假设程序A先发现文档不存在，然后程序B发现文档不存在。此时A还没来得及对文档进行写入，因此文档锁并没有挂起。或者说，由于文档不存在，讨论文档锁也就失去了意义。这个时候，两个进程就会分别创建文档并给其votes加1。于是就出现了不必要的重复。

如何解决？

解决方法其实很简单：unique index。 上文提到，name属性是唯一的。如果我们给它加一个唯一索引，不就可以从根本上避免一个人有多个不同的文档了吗？这个时候，即使两个进程经过搜索都得到了某个文档不存在的结果，假设A先一步创建了该文档，那么当B创建文档时，由于含有相同name的文档已经被A进程抢先创建，MongoDB就会拒绝B进程创建。pymongo对此类错误应该是有应对机制的，这是B进程会稍等片刻，重新尝试更新文档。这个时候，A进程已经完成计票并且释放了写入锁，文档被成功创建，而进程B再尝试时，也会检索到这个被新创建的文档，直接在上面把票数加1，而不是创建新文档。这样一个小的时间差，就解决了写入矛盾。

同时，我们还得到了额外的奖励：当name上创建了unique index之后，找到特定候选人的速度就会快很多。这个优势在计票初期，候选人数量不多时并没有显示，但当后期候选人数变多时，一方面再有新的候选人被加入的概率会变得很小（该被加的差不多都被加进来了），因此修改索引的几率越来越少；另一方面，在候选人基数变得很大的时候，相比于没有索引的情况，有唯一索引的情况下程序的速度优势会越发明显。这两个方面综合在一起，结果就是，添加唯一索引之后程序在后期速度优势会越来越明显。在我自己的程序中，运行初期多线程比单线程只快了三四倍，但在数据量较大时，多线程（加上唯一索引）会比单线程快10到20倍。这多处来的速度，就是唯一索引导致的。