{
    "componentChunkName": "component---src-templates-blog-blog-detail-tsx",
    "path": "/blog/tikv-source-code-reading-10",
    "result": {"pageContext":{"blog":{"id":"Blogs_293","title":"TiKV 源码解析系列文章（十）Snapshot 的发送和接收","tags":["TiKV 源码解析","社区"],"category":{"name":"产品技术解读"},"summary":"TiKV 针对 Snapshot 收发场景做了特殊处理，解决了消息包过大会导致的一系列问题。","body":"## 背景知识\n\nTiKV 使用 Raft 算法来提供高可用且具有强一致性的存储服务。在 Raft 中，Snapshot 指的是整个 State Machine 数据的一份快照，大体上有以下这几种情况需要用到 Snapshot：\n\n1. 正常情况下 leader 与 follower/learner 之间是通过 append log 的方式进行同步的，出于空间和效率的考虑，leader 会定期清理过老的 log。假如 follower/learner 出现宕机或者网络隔离，恢复以后可能所缺的 log 已经在 leader 节点被清理掉了，此时只能通过 Snapshot 的方式进行同步。\n2. Raft 加入新的节点的，由于新节点没同步过任何日志，只能通过接收 Snapshot 的方式来同步。实际上这也可以认为是 1 的一种特殊情形。\n3. 出于备份/恢复等需求，应用层需要 dump 一份 State Machine 的完整数据。\n\nTiKV 涉及到的是 1 和 2 这两种情况。在我们的实现中，Snapshot 总是由 Region leader 所在的 TiKV 生成，通过网络发送给 Region follower/learner 所在的 TiKV。\n\n理论上讲，我们完全可以把 Snapshot 当作普通的 `RaftMessage` 来发送，但这样做实践上会产生一些问题，主要是因为 Snapshot 消息的尺寸远大于其他 `RaftMessage`：\n\n1. Snapshot 消息需要花费更长的时间来发送，如果共用网络连接容易导致网络拥塞，进而引起其他 Region 出现 Raft 选举超时等问题。\n2. 构建待发送 Snapshot 消息需要消耗更多的内存。\n3. 过大的消息可能导致超出 gRPC 的 Message Size 限制等问题。\n\n基于上面的原因，TiKV 对 Snapshot 的发送和接收进行了特殊处理，为每个 Snapshot 创建单独的网络连接，并将 Snapshot 拆分成 1M 大小的多个 Chunk 进行传输。\n\n## 源码解读\n\n下面我们分别从 RPC 协议、发送 Snapshot、收取 Snapshot 三个方面来解读相关源代码。本文的所有内容都基于 v3.0.0-rc.2 版本。\n\n### Snapshot RPC call 的定义\n\n与普通的 raft message 类似，Snapshot 消息也是使用 gRPC 远程调用的方式来传输的。在 [pingcap/kvproto](https://github.com/pingcap/kvproto) 项目中可以找到相关 RPC Call 的定义，具体在 [tikvpb.proto](https://github.com/pingcap/kvproto/blob/5cb23649b361013f929e0d46a166ae24848fbcbb/proto/tikvpb.proto#L57) 和 [raft_serverpb.proto](https://github.com/pingcap/kvproto/blob/5cb23649b361013f929e0d46a166ae24848fbcbb/proto/raft_serverpb.proto#L42-L47) 文件中。\n\n```proto\nrpc Snapshot(stream raft_serverpb.SnapshotChunk) returns (raft_serverpb.Done) {}\n...\nmessage SnapshotChunk {\n  RaftMessage message = 1;\n  bytes data = 2;\n}\n\nmessage Done {}\n```\n\n可以看出，Snapshot 被定义成 client streaming 调用，即对于每个 Call，客户端依次向服务器发送多个相同类型的请求，服务器接收并处理完所有请求后，向客户端返回处理结果。具体在这里，每个请求的类型是 `SnapshotChunk`，其中包含了 Snapshot 对应的 `RaftMessage`，或者携带一段 Snapshot 数据；回复消息是一个简单的空消息 `Done`，因为我们在这里实际不需要返回任何信息给客户端，只需要关闭对应的 stream。\n\n### Snapshot 的发送流程\n\nSnapshot 的发送过程的处理比较简单粗暴，直接在将要发送 `RaftMessage` 的地方截获 Snapshot 类型的消息，转而通过特殊的方式进行发送。相关代码可以在 [server/transport.rs](https://github.com/tikv/tikv/blob/892c12039e0213989940d29c232bddee9cbe4686/src/server/transport.rs#L313-L344) 中找到：\n\n```rust\nfn write_data(&self, store_id: u64, addr: &str, msg: RaftMessage) {\n  if msg.get_message().has_snapshot() {\n      return self.send_snapshot_sock(addr, msg);\n  }\n  if let Err(e) = self.raft_client.wl().send(store_id, addr, msg) {\n      error!(\"send raft msg err\"; \"err\" => ?e);\n  }\n}\n\nfn send_snapshot_sock(&self, addr: &str, msg: RaftMessage) {\n  ...\n  if let Err(e) = self.snap_scheduler.schedule(SnapTask::Send {\n      addr: addr.to_owned(),\n      msg,\n      cb,\n  }) {\n      ...\n  }\n}\n```\n\n从代码中可以看出，这里简单地把对应的 `RaftMessage` 包装成一个 `SnapTask::Send` 任务，并将其交给独立的 `snap-worker` 去处理。值得注意的是，这里的 `RaftMessage` 只包含 Snapshot 的元信息，而不包括真正的快照数据。TiKV 中有一个单独的模块叫做 `SnapManager` ，用来专门处理数据快照的生成与转存，稍后我们将会看到从 `SnapManager` 模块读取 Snapshot 数据块并进行发送的相关代码。\n\n我们不妨顺藤摸瓜来看看 `snap-worker` 是如何处理这个任务的，相关代码在 [server/snap.rs](https://github.com/tikv/tikv/blob/892c12039e0213989940d29c232bddee9cbe4686/src/server/snap.rs#L329-L398)，精简掉非核心逻辑后的代码引用如下：\n\n```rust\nfn run(&mut self, task: Task) {\n  match task {\n      Task::Recv { stream, sink } => {\n           ...\n           let f = recv_snap(stream, sink, ...).then(move |result| {\n               ...\n           });\n           self.pool.spawn(f).forget();\n      }\n      Task::Send { addr, msg, cb } => {\n          ...\n          let f = future::result(send_snap(..., &addr, msg))\n              .flatten()\n              .then(move |res| {\n                  ...\n              });\n          self.pool.spawn(f).forget();\n      }\n  }\n}\n```\n\n`snap-worker` 使用了 `future` 来完成收发 Snapshot 任务：通过调用 `send_snap()` 或 `recv_snap()` 生成一个 future 对象，并将其交给 `FuturePool` 异步执行。\n\n现在我们暂且只关注 `send_snap()` 的 [实现](https://github.com/tikv/tikv/blob/892c12039e0213989940d29c232bddee9cbe4686/src/server/snap.rs#L103-L175)：\n\n```rust\nfn send_snap(\n  ...\n  addr: &str,\n  msg: RaftMessage,\n) -> Result<impl Future<Item = SendStat, Error = Error>> {\n  ...\n  let key = {\n      let snap = msg.get_message().get_snapshot();\n      SnapKey::from_snap(snap)?\n  };\n  ...\n  let s = box_try!(mgr.get_snapshot_for_sending(&key));\n  if !s.exists() {\n      return Err(box_err!(\"missing snap file: {:?}\", s.path()));\n  }\n  let total_size = s.total_size()?;\n  let chunks = {\n      let mut first_chunk = SnapshotChunk::new();\n      first_chunk.set_message(msg);\n\n      SnapChunk {\n          first: Some(first_chunk),\n          snap: s,\n          remain_bytes: total_size as usize,\n      }\n  };\n\n  let cb = ChannelBuilder::new(env);\n  let channel = security_mgr.connect(cb, addr);\n  let client = TikvClient::new(channel);\n  let (sink, receiver) = client.snapshot()?;\n\n  let send = chunks.forward(sink).map_err(Error::from);\n  let send = send\n      .and_then(|(s, _)| receiver.map_err(Error::from).map(|_| s))\n      .then(move |result| {\n          ...\n      });\n  Ok(send)\n}\n```\n\n这一段流程还是比较清晰的：先是用 Snapshot 元信息从 `SnapManager` 取到待发送的快照数据，然后将 `RaftMessage` 和 `Snap` 一起封装进 `SnapChunk` 结构，最后创建全新的 gRPC 连接及一个 Snapshot stream 并将 `SnapChunk` 写入。这里引入 `SnapChunk` 是为了避免将整块 Snapshot 快照一次性加载进内存，它 impl 了 `futures::Stream` 这个 trait 来达成按需加载流式发送的效果。如果感兴趣可以参考它的 [具体实现](https://github.com/tikv/tikv/blob/892c12039e0213989940d29c232bddee9cbe4686/src/server/snap.rs#L55-L92)，本文就暂不展开了。\n\n### Snapshot 的收取流程\n\n最后我们来简单看一下 Snapshot 的收取流程，其实也就是 gRPC Call 的 server 端对应的处理，整个流程的入口我们可以在 [server/service/kv.rs](https://github.com/tikv/tikv/blob/892c12039e0213989940d29c232bddee9cbe4686/src/server/service/kv.rs#L714-L729) 中找到：\n\n```rust\nfn snapshot(\n  &mut self,\n  ctx: RpcContext<'_>,\n  stream: RequestStream<SnapshotChunk>,\n  sink: ClientStreamingSink<Done>,\n) {\n  let task = SnapTask::Recv { stream, sink };\n  if let Err(e) = self.snap_scheduler.schedule(task) {\n      ...\n  }\n}\n```\n\n与发送过程类似，也是直接构建 `SnapTask::Recv` 任务并转发给 `snap-worker` 了，这里会调用上面出现过的 `recv_snap()` 函数，[具体实现](https://github.com/tikv/tikv/blob/892c12039e0213989940d29c232bddee9cbe4686/src/server/snap.rs#L237-L291) 如下：\n\n```rust\nfn recv_snap<R: RaftStoreRouter + 'static>(\n  stream: RequestStream<SnapshotChunk>,\n  sink: ClientStreamingSink<Done>,\n  ...\n) -> impl Future<Item = (), Error = Error> {\n  ...\n  let f = stream.into_future().map_err(|(e, _)| e).and_then(\n      move |(head, chunks)| -> Box<dyn Future<Item = (), Error = Error> + Send> {\n          let context = match RecvSnapContext::new(head, &snap_mgr) {\n              Ok(context) => context,\n              Err(e) => return Box::new(future::err(e)),\n          };\n\n          ...\n          let recv_chunks = chunks.fold(context, |mut context, mut chunk| -> Result<_> {\n              let data = chunk.take_data();\n              ...\n              if let Err(e) = context.file.as_mut().unwrap().write_all(&data) {\n                  ...\n              }\n              Ok(context)\n          });\n\n          Box::new(\n              recv_chunks\n                  .and_then(move |context| context.finish(raft_router))\n                  .then(move |r| {\n                      snap_mgr.deregister(&context_key, &SnapEntry::Receiving);\n                      r\n                  }),\n          )\n      },\n  );\n  f.then(move |res| match res {\n      ...\n  })\n  .map_err(Error::from)\n}\n```\n\n值得留意的是 stream 中的第一个消息（其中包含有 `RaftMessage`）被用来创建 `RecvSnapContext` 对象，其后的每个 chunk 收取后都依次写入文件，最后调用 `context.finish()` 把之前保存的 `RaftMessage` 发送给 `raftstore` 完成整个接收过程。\n\n## 总结\n\n以上就是 TiKV 发送和接收 Snapshot 相关的代码解析了。这是 TiKV 代码库中较小的一个模块，它很好地解决了由于 Snapshot 消息特殊性所带来的一系列问题，充分应用了 `grpc-rs` 组件及 `futures`/`FuturePool` 模型，大家可以结合本系列文章的 [《TiKV 源码解析系列文章（七）gRPC Server 的初始化和启动流程》](https://pingcap.com/blog-cn/tikv-source-code-reading-7/) 和 [《TiKV 源码解析系列文章（八）grpc-rs 的封装与实现》](https://pingcap.com/blog-cn/tikv-source-code-reading-8/) 进一步拓展学习。\n\n> 点击查看更多 [TiKV 源码解析系列文章](https://pingcap.com/zh/blog/?tag=TiKV%20%E6%BA%90%E7%A0%81%E8%A7%A3%E6%9E%90)","date":"2019-07-09","author":"黄梦龙","fillInMethod":"writeDirectly","customUrl":"tikv-source-code-reading-10","file":null,"relatedBlogs":[]}}},
    "staticQueryHashes": ["1327623483","1820662718","3081853212","3430003955","3649515864","4265596160","63159454"]}