建议对照着 laixintao 老师的 这篇文章 来补齐一些概念

tcptrace 图中的基本元素

上面的 tcptrace 示例中（关注其中的元素，其它的时序不真实）：

• 横轴是时间，纵轴是 Sequence Number，可以理解为包的字节数
• 其中每 发送 一个数据包，就会画一条蓝色的 I 型竖线段，长度表示数据量
• 上方的绿色阶梯线代表接收方的窗口大小（Receive Window），即最多能发多少数据
• 下方的棕黄色阶梯线代表已经被 ACK 的数据量，即接收方实际收到的数据
• 红色的线段代表选择确认（SACK），即接收方收到的数据中间有缺失，红线代表收到了哪些数据
• 棕黄线下方的小竖线代表重复 ACK（dup ACK），通常代表接收方收到了乱序的包

所以这个图也只是把发的包和收到包的信息以图形化的方式展现出来，方便我们分析。

tcptrace 图中的“距离”信息

除了直接标记出来的元素，还可以通过图中的一些“距离”，得知其它一些概念

• Bytes in Flight: 发送方已经发出去但还没有被 ACK 的数据量，某时刻蓝色线和棕黄色线的垂直距离
• Window Room（窗口余量）： 接收方还能接收的数据量，某时刻绿色线和棕黄色线的垂直距离
• RTT（往返时间）: 数据发送到确认接收所需的时间，图中蓝色线和棕黄色线的水平距离

还有就是整体的斜率，代表吞吐量（Throughput），斜率越大，吞吐量越高。

发送方与接收方抓的包有差异

发送方作为最初发包和接收 ACK 的一方，可以看到整体的 RTT 及链路信息，如下图中可以看到蓝线和 ACK 线的距离(代表 RTT)：

而接收方中中的蓝线和 ACK 线的距离，则代表的是本机接收到包后本机发 ACK 的距离，并不能代表 RTT。但是接收方额外地可以看到发送方的一些乱序包，而发送方是看不到的。

结论是两方的抓包有各自无法替代的信息，如果可能尽量两方同时抓包，如果只有一方，尽量是发送方。

常见网络问题下 tcptrace 图的表现

丢包

1. 看到红色的 SACK 线，说明接收方有包没收到
2. 看到不断有递增的红色 SACK 线，说明后续的包都收到了，但是丢失的包依旧没收到
3. 丢包期间，接收方的 buffer 不能释放，导致窗口没有增长
4. 终于收到丢失的包后，综色的 ACK 线跳跃式上升，且绿色窗口线也上升，说明接收方释放了 buffer
5. 由于接收方窗口增大，发送方马上利用了新增的窗口，蓝色线也跟着上升

额外注意这个图里可以看到发送方重传了没有收到数据。

吞吐受到接收方窗口限制

1. 蓝线和距离绿线非常接近，说明发送方一直发送，接近了接收方窗口的上限
2. 绿线一上升，蓝线也跟着上升，说明发送方马上利用了接收方新增的窗口

结合起来看，说明接收方窗口太小，限制了发送方的吞吐。

吞吐受到发送方窗口限制

发送方能发多少数据除了受接收方窗口限制外，还受两个因素影响，一个是发送窗口（代码里使用 SO_SNDBUF 指定，内核里通过 net.ipv4.tcp_wmem 指定），另一个是拥塞窗口(cwnd)。例如下面的示例（构造时把 net.ipv4.tcp_wmem 设置成了 4096 4096 4096）^[1]:

上图中蓝线不断上升，但远不到绿线的位置就停止了，直到有 ACK 一定数据后才继续发送。说明是发送方做了限制。但是究竟是 send_buffer 的限制还是 cwnd 的限制？这就需要使用 “Window Scaling”：

简单理解 Windown Scaling 里计算的是每个包对应的 “Bytes In Flight” 的数据，也就是“发出但未被 ACK”的数据，通常这个信息可以认为约等于发送方的 cwnd。但是上图中这个窗口信息很“平”，而一般 cwnd 是会波动的，因此可以认为不是 cwnd 引起的，那么就只能是发送窗口了。作为对比，我们看一个不显式设置 send buffer 的示例：

可以看到上图中随着一些丢包的发生，窗口也在不断变化，这与 cwnd 随着 RTT 和丢包不断变化相匹配。

特例：零窗口

上面提到速率可能跟发送方窗口有关，还有一种特例是因为各种原因（如消费的速度太慢）导致发送方窗口为 0, 此时 wireshark 会特地标记为 x：

网络状况很差

模拟 5% 的丢包与 10% 的乱序得到如下示例，一眼看就是很“不均匀”，有红色的丢包，棕色线有坑坑挖挖的乱序，整体的上升趋势（斜率）一会高一会低：

我们再放大看一些特征：

完美的连接

完美的情况下，会是一个非常干净的图，笔直的线，没有 SACK、乱序等复杂元素，也能达到网络的带宽：

附：模拟网络条件用到的一些命令

上面的实验是在 Mac 上用 podman 的两个 container 完成的，使用以下方式模拟网络条件:

1. 先 podman machine ssh 进入 podman 虚拟机
2. 这个虚拟机缺少 sch_netem 内核模块，通过 rpm-ostree install kernel-modules-extra 安装
3. 重启 podman 的虚拟机
4. 通过 podman machine ssh 进入虚拟机执行 modprobe sch_netem
5. 视情况执行如下 tc 命令来设置网络条件

tc qdisc del dev veth0 root
tc qdisc del dev veth1 root
tc qdisc add dev veth0 root handle 1: netem delay 30ms loss 5% reorder 10%
tc qdisc add dev veth0 parent 1:1 handle 10: tbf rate 10mbit burst 500kbit limit 1m
tc qdisc add dev veth1 root handle 1: netem delay 30ms loss 5% reorder 10%
tc qdisc add dev veth1 parent 1:1 handle 10: tbf rate 10mbit burst 500kbit limit 1m

其中 veth0, veth1 这两个网卡，是在虚拟机中通过如下命令定位

1. podman ps 获取 container id
2. podman inspect <container id> | grep Pid 获取 Pid
3. nsenter -t <pid> -n ip link 看到 container 网卡对应的 id 为 eth0@if5，注意这里的 5
4. ip link 找到编号为 5 的网卡 5: veth1@if2，因此为 veth1

pcap 文件

• iperf-bad-net.pcap
• iperf-client-1mb.pcap
• iperf-client-buf-4k.pcap
• iperf-client-buf-auto-with-loss.pcap
• iperf-client-buf-auto.pcap
• tcptrace-client-defaultwindow.pcapng
• tcptrace-client-loss.pcapng
• Lab1-GreerBombal_ItsNotTheNetwork.pcapng

参考

• 用 Wireshark 分析 TCP 吞吐瓶颈laixintao 老师的博客，看了之后才第一次真正读懂了 tcptrace 图的含义
• Wireshark TCP Trace Graph Tutorial 细致地介绍了一些常见网络问题下 tcptrace 图的表现，可以作为 laixintao 老师博客的补充
• 在Wireshark的tcptrace图中看清TCP拥塞控制算法的细节(CUBIC/BBR算法为例) 逻辑性地分析不同拥塞算法下 tcptrace 图的表现，但我本身对算法不熟悉，只能浅尝则止

RAG 的作用

RAG 能解决纯靠大模型生成的一些问题：

1. 幻觉问题：大模型可能会“凭空捏造”一些看似合理的信息，但实际上是错误的。
2. 知识更新问题：大模型无法及时更新知识，如新闻、最新研究等。
3. 私有知识问题：企业内部的资料非公开，大模型训练时不包含这些信息。

当然还有一些其它工程问题，如生成速度、资源消耗、上下文限制等。

Naive RAG

一个简单的 RAG 流程主要包含三个部分：

1. 索引器 Indexer: 解析文档，对文档做切片并向量化，存储在向量数据库中。
2. 检索器 Retriever: ，从向量库中检索与问题最相关的 top-N 文本切片。
3. 生成器 Generator: 根据检索到的文本切片生成回答。

为什么要这么做？本质上还是因为“知识”太多，大模型的上下文放不下，因此需要提前筛选出相关的内容。但由于筛选的过程是很耗时的（不管是因为数量太大，还是筛选的逻辑很复杂），因此需要提前建立索引。

RAG 缺少大局观

我们看下面这个文档，假设它被切成了右边的 4 个切片并分别建立索引:

现在我们提出问题：RAG 有什么优点，假设 Retriever 只返回 top-1 的结果，上帝视角来看就会遇到问题：

1. 单看片段 ③，我们并不知道它说的是什么东西的优缺点，它并没有提到 RAG，很可能不会被召回
2. 因为只返回 top-1，但实际上“优点”可能即在片段 ③，也在片段 ④ 中，于是回答不完整

这就是 RAG 缺少大局观的问题，处理的目标对象是“文档的切片”，而切片的局部性注定了它无法看到整篇文档的信息，也无法看到整个知识库的信息，因此：

1. RAG 检索切片时，可能会因为切片本身包含的上下文信息不足，而漏掉正确的切片
2. RAG 无法判断需要多少个切片才能回答问题，于是诸如“列举 XXX”，“总结 XXX”的问题常常回答不完整
3. RAG 无法判断文档间的联系，如法律条文，新的解释覆盖旧的解释，但 RAG 无法判断哪个是最新的

无法望文生义

RAG 是基于文档的，但文档都有自己的语境，有时仅凭文档内容根本无法理解，必须要结合语境。有时需要参考其它文档，有时需要了解文档所处的背景（如是在哪个知识库下的），有时需要有一些先验知识（如行业黑话）。

例如银行发了一张卡叫“宝贝成长卡”，你问“宝贝成长卡怎么销卡”，可能找遍所有的文档都找不到答案。这时需要分析知道“宝贝成长卡”是一张“借计卡”，再去找“借记卡销卡”相关的流程才能找到答案。

再例如有一篇文章说的是“A 卡的种类”，你可以理解是“AMD 显卡的种类”，但如果金融背景下，它可能就是指“信用卡评分模型的种类”了。单看这个文章是很难决定是否和问题匹配的。

再有像是一些行业黑话，口语，网络用语等，这些都是文档本身没有的但是用户会问的，RAG 仅靠文档是无法建立这两者之间的联系的。

一些提升手段

文档结构化

做法是提取文档的结构化信息，如标题、目录等，将这些信息也加入到切片中一起做索引。

这种方法背后的想法是文档的各级标题通常包含大量信息，在切片中加入可极大增加切片的信息量。另外还有很多细节需要优化，如目录相可能会干扰正文，可视情况移除；遇到列表最好放在同一个切片。

当然文档结构化本身是个脏活累活，一方面有许多不同类型的文档需要支持，如 PDF、Word、Excel、PPT 等。另一方面即便只是 doc，如何提取标题、目录也是个大问题。毕竟绝大多数文档都是不规范的。

多路召回与重排序

上面的 Naive RAG 模型中，只使用了向量化这一种手段，通常也称为语义检索(SemanticSearch)或 Dense Retrieval。但向量化在一些场景下效果并不好，例如：

• 专有名词：一些领域特定的词汇在向量模型训练时可能没有出现过
• 缺乏精确性：返回的可能是语义相近但并不是精确匹配的内容
• 长度敏感：文档切片、Query 的长度不同，很可能有不同的效果

所以通常会使用多路召回，再加上重排序，以提高检索的效果。

实践中最重要的是加上关键词检索，例如 ElasticSearch、Solr 等。但由于评分体系不同，最后需要重新对所有的文档进行排序。常见的做法有：

• 可以专门使用一个重排序模型，如 bge-reranker
• 也可以使用 Reciprocal Rank Fusion (RRF) 这类简单的融合方法
• 还可以对各种分数归一化后再加权求和（如 Milvus 会对关键词搜索的分数使用arctan 归一化^[1]）

查询改写/扩展

对于像行业黑话，或者其它很难直接通过文档建立联系的概念，可以考虑由人工来提供一些信息，在查询前对 Query 进行改写或扩展。亦或者使用知识图谱或 LLM 来补充信息。

文档 meta 利用与查询理解

用户的问题常常也会直接或间接包含“结构化”的查询条件，如“二孩的专项附加扣除是多少”隐含时间信息“今年”；再如“XX 公司 2023 年的营收是多少”指定了公司名和时间。这些信息是过滤条件的强表达，如果在检索时不体现，则通常意味着结果不准确。

这通常需要结合文档的 meta 信息，例如维护文档的时间信息、作者信息、来源信息等。再利用 LLM 从用户 Query 中抽取相关的过滤条件，在检索时加入这些条件。

解决大局观问题其它尝试

上面提到的文档结构化、多路召回、查询改写等都可以认为是在打补丁，怎么为文档片段增加上下文信息。但它们依旧解决不了”如何决定需要多少文档片段”这类问题。

目前看到最多的尝试就是使用知识图谱来解决这个问题。其中的典型是GraphRAG^[2]，它的流程如下:

它的问答流程其实已经不是一个传统的 RAG 了，可以理解成它是在索引阶段，通过对文档片段抽取实体关系及对应的描述，聚类成了一个（非典型的）知识图谱。而在检索阶段，就直接使用生成的聚类描述来回答问题。说实话我觉得这套方法大概率是无法落地的，有如下质疑：

1. 如何正确的且无监督地抽取实体关系
2. 如何决定何时使用聚类来回答，何时使用普通的文档片段
3. 如何保证响应时间（GraphRAG 回答时需要先使用聚类描述生成答案，由 LLM 筛选正确后再汇总）

但这个方法其实让我想到传统机器学习中的混合高斯模型(GMM^[3])，它代表的混合模型思想，我粗浅的理解是：如果用一个分布无法解释数据，那就用多个分布来解释。在RAG 中，可以理解成：如果用一个长度的文档片段无法解释问题，那就用多个文档片段来解决。

例如在建索引时，可以分别以 200、400、600、… 个字为单位建立索引。并在检索时定位到对应长度的片段。这样可以保证大概率有一个足够长的片段是包含了足够多上下文的。当然，向量模型的能力可能并足以理解这么长的文本，导致效果不佳。

落地的最大问题: 人力

遇到效果问题，很多客户都会提“微调行不行”，可以，但是得加钱。

我只调过 Ranking 模型和 LLM，都需要高质量的标注数据。但一般产出这些数据需要对应的业务专家，尤其是 LLM 的 SFT 数据，还需要提供结果的回答。我们见到的很多客户，连知识库的数据预期从哪来，哪些合适放在知识库都回答不了，更别说调动资源来准备这些数据了。

也因此，ToB 的 RAG 项目很难落地，无监督的各种方法很快就会达到上限，而后续的数据治理又没法推进，于是死局。

未来在哪？

其实这个问题才是我的这碟醋，上面的内容都是为这碟醋包的饺子。当然也是在脑洞与胡思乱想。

我觉得不管是简单的向量化，还是像 GraphRAG，或者其它使用 LLM 提前建图谱的各种方法，本质上都是在尝试提前更好地“理解”文档，甚至建立文档间的联系。这个问题其实是“知识的表示”。

如果说大模型是世界知识的“无损”压缩，那么 RAG 就是在尝试压缩知识库中的这些知识，并在回答问题时提取出最相关的知识。但显然不管是向量化还是 GraphRAG，都不是一个足够好的压缩方式。

因此我认为未来会有一种高效的训练方式，能够在“理解”知识库的内容后，高效地压缩成一个模型，并且具有很高的压缩比。同时知识可以叠加，在问答时，可以结合世界知识模型和知识库浓缩模型来回答问题。

另一个问题是：搜索引擎是一个足够好的压缩方式吗？我认为对绝大多数场景来说是的。我也认为当前的 RAG 如果做的足够好，检索方面可能也就是一个复杂的搜索引擎了。

参考

• Retrieval-Augmented Generation for Large Language Models: A Survey 23年12月的一篇综述论文
• Retrieval-Augmented Generation for AI-Generated Content: A Survey 另一篇 24 年 2 月的综述论文
• Graph Retrieval-Augmented Generation: A Survey 使用图来增强 RAG 的综述，24 年 8 月
• 大模型RAG 场景、数据、应用难点与解决（四） 知乎的一篇 RAG 文章，里面很多痛我们也实际遇到了

字符串匹配难在哪里

字符串匹配其实就是要实现字符串的 contains 方法，判断一个字符串 s 中是否包含另一个字符串 p。例如 "hello".contains("ll") 应该返回 true。应该如何实现呢？

一个很自然的想法是一个个匹配（如下图左侧）。从 s 的第一个字符开始，依次和 p 进行匹配 ①，如果匹配失败，从 s 的下一个字符开始，再次尝试匹配 ②，依此类推 ③。

这种方法的问题是需要的匹配次数太多。它的时间复杂度是 O(mn)，m 和 n分别是 s 和 p 的长度。

哪里可以优化

通常在匹配过程中，待查找/匹配的字符串 p 是固定的，而被查找的字符串 s 是未知的，这意味着我们可以充分对 p 进行分析，从而优化匹配过程。

观察匹配过程，虽然我们并不知道 s 有哪些字符，但在 ① 的匹配过程中，前几个字符和 p 是匹配的，因此我们能确认 s 的前 5 个字符一定是 s[0:5] = "ababa"。

那么，当我们尝试把 s 向后移位，去匹配 s[1:5] 与 p 时，我们其实是在浪费时间，因此我们已经知道 s[1] = 'b' 而 p[0] = 'a'，肯定是不匹配的。这就是 ② 的情况。

再考虑 ③ 中的匹配，同样由于已经知道了 s[0:5] = "ababa"，移 2 位后 s[2:5] = "aba"，肯定是能和 p 的前三个字符 p[0:3] = "aba" 匹配的。没必要再匹配一遍，可以直接从s[5] 和 p[3] 开始匹配。

换句话说，通过 ① 中的匹配得到的关于 s 的信息，以及对字符串 p 的分析，在 ①匹配失败后，我们其实可以直接跳到 s[5] 和 p[3] 开始匹配（如上图右侧）。以此节省许多无效的匹配。

跳到哪里

上面的分析中，为什么我们能判断 ① 匹配失败后，直接尝试匹配 s[5] 和 p[3] 就可以呢？这是因为我们不断地向后移位，直到移了 x 位时，发现 s[x:5] 与p[0: 5-x] 匹配。在上例中，x = 2。于是下一次就可以从 p[5-x] 开始匹配。

• 这个分析其实完全不需要 s，因为所有 s 中要用到的信息（匹配的字符）都包含在 p 中了
• 移位和跳过匹配的过程，可以看作是在 p 中找到一个最长的前缀，使得这个前缀同时也是 p 的后缀

因此我们实际上要求的是：对于一个字符串 p[0:n]，找到一个最长的前缀，使得这个前缀（长度为 k）同时也是 p 的后缀，即 p[0:k] = p[n-k:n]。（注意k与上面的 x 是不同的，x = n-k）。

注：下文开始，我们提到“前缀”时指的都是同时既是前缀，也是后缀的字符串。

如何找到最长前缀

这个问题需要递归地考虑，我们记 T[i] 为 p[0:i] 的最长前缀长度（满足前缀等于后缀）。假设我们已经知道了所有的 T[0], ..., T[i-1]，现在我们要求 T[i]。

已知 p[0:i] 字符串最长前缀长度为 T[i-1]，前缀和后缀字符串分别记为 X 和 Y，有 X=Y:

① 现在 T[i] 最好的情况是 T[i-1] + 1。此时需要满足 p[i] = p[T[i-1]]:

② 如果 p[i] != p[T[i-1]]，则我们可以假设已经找到了 T[i]，看看它满足什么条件。首先我们知道 T[i] 一定小于 T[i-1]，所以假设找到了 T[i] 并记前缀为 A，后缀为 B，则有下图：

由于 X = Y，所以一定可以在 X 中找到后缀字符串 C，满足 C = B = A。于是我们发现，前缀 A 即是 p[0:i] 的前缀，也是 p[0:T[i-1]] 的前缀。因此我们可以得出结论：T[i] <= T[T[i-1]]。

接着可以从最大值 T[T[i-1]] 开始，判断 p[i] 和 p[T[T[i-1]]] 是否相等，此时就递归加了情况 ①。

最终，如果匹配到 p[0] 还是不匹配，则认为不存在前缀，此时 T[0] = 0。

建表代码

建表的逻辑就是上面描述的递归过程，只是用循环来实现：

def kmp_build_table(pattern):
    table = [0] * len(pattern)
    i = 0
    for j in range(1, len(pattern)):  # ①
        while i > 0 and pattern[i] != pattern[j]:
            i = table[i - 1]  # ②
        if pattern[i] == pattern[j]:
            i += 1 # ③
        table[j] = i # ④
    return table

这个代码的时间复杂度是 O(n)，其中 n 是 pattern 的长度。

外层循环 ① 从 1 到 n = len(pattern) 运行 O(n) 次比较容易理解。内层循环首先注意到 ③ 中，i 每次循环最多只增加 1，而 ④ 中 table 赋值为 i，因此可推出 table 中任意 m > n 两个元素，满足 table[m] - table[n] <= (m-n)。换句话说，在 ② 中的操作使得 i 是不断减小的，且全局减小的次数 < n，于是减少的次数是 O(n) 的。另一方面 i 每次最多增加 1，增加的次数也是 O(n) 的。因此整体的时间复杂度是 O(n)。

字符串匹配代码

KMP 的算法就很简单了，只需要在匹配失败时，根据表中的值跳到下一个需要匹配的位置即可：

def kmp_search(text, pattern):
    table = kmp_build_table(pattern)
    i = 0
    for j in range(len(text)):
        # skip the non matching part
        while i > 0 and text[j] != pattern[i]:
            i = table[i - 1] # next char in pattern to match
        if text[j] == pattern[i]:
            i += 1 # text shift to next
        if i == len(pattern):
            return j - i + 1
    return -1

这里的分析和建表代码一样，重点是内层循环是回退操作，且它的值不会因为 j 而被重置，因此内层的总时间复杂度是 O(m)，m 是 pattern 的长度。因此总的搜索时间为 O(m+n)。

总结

个人觉得理解过程中有两个关键点：

1. 为什么最后问题等价于找到一个最长的前缀，使得这个前缀同时也是后缀
2. 在建表的过程中，为什么可以递归？（即分析中 X = Y 的性质，使得可以找到 C = B = A 的字符串）

如果能理解这两点，KMP 其它的部分相信就不难理解了。希望这篇文章对你有帮助。

参考

• 前缀函数与 KMP 算法 这里的分析更详细，里面学习到了分析中 X = Y 的性质

题外话

在写代码的时候 Github Copilot 给我补全了下面的错误代码。怎么看都觉得不对，但是像这种很经典的代码它又不太应该出错，于是怀疑自己怀疑了半天。问 ChatGPT 倒是直接指出了错误，但让它给 test case 也老是给不对。

感慨下目前 AI 写代码，还是需要很强的鉴别能力的。这个代码在很多常见 case 下和正确代码是一样的，但它的逻辑就是错误的，如果埋在代码里得被坑死。

def kmp_build_table(pattern):
    table = [0] * len(pattern)
    i = 0
    for j in range(1, len(pattern)):
        if pattern[i] == pattern[j]:
            i += 1
            table[j] = i
        else:
            i = 0
    return table

引子

一个业务系统是用 Webflux 写的，发现后台在做批量任务时，会卡住页面的访问。排查发现是把 r2dbc 的 IO 线程给卡住了，导致页面请求时从数据库捞数据的请求卡死。但这个批量操作本身并没有特别多的 DB 操作，为什么会卡住呢？

Reactor Stream 简介

Java 9 引入了 java.util.concurrent.Flow 接口，支持[Reactive Streams 规范](https://www.reactive-streams.org/)。规范的核心是定义了发布者(Publisher)和订阅者(Subscriber)的交互逻辑，规定Subscriber 必须以 PULL (拉取)的方式获取数据，以此解决异步流式处理中的背压问题^[1]。

Reactive Streams 规定的交互流程如下（由于标准中有些部分留白，实际有两种常见模式）：

主体流程分为这么几步：

1. Subscriber 向 Publisher 订阅。onSubscribe 的入参是订阅者 Subscriber。
2. Publisher 通知 Subscriber 订阅成功，并发送一个 Subscription 对象用于后续交互。
3. 当 Subscriber 有处理能力时，调用 Subscription 的 request 方法通知Publisher 发送 N 个数据
4. 每有一个新数据，调用 Subscriber 的 onNext 方法一次，直到发送了 N 个数据。

为什么需要有 Subscription 这个接口呢？为什么不直接把 request 方法定义在Publisher 中呢？有个大前提是 Reactive Streams 规范中，一个 Publisher 可以有多个 Subscriber，于是如果没有 Subscription，则 Publisher 需要在内部维护这个Subscriber 与数据的关系，增加了复杂度。因此不管是从概念上的解耦还是减小实现复杂性及提高性能性能方面考虑，把 Subscriber 与 Publisher 之间交互的生命周期抽象成Subscription，都是一个不错的选择。

另外注意到图里有两种模式。Reactive Sterams 只规定调用了 Subscription.request之后，如果有新的数据需要调用 Subscriber 的 onNext 方法。但是并没有规定onNext 谁来调用。于是根据 Publisher 中数据是否需要共享，可以分为 Cold 和 Hot两种模式。

Cold 模式下数据是分离的，每个 Subscriber 都有自己的数据流，例如Flux.range，每个 subscriber 都会从头开始计数。于是 Publisher 可以把当前消费的位置保存在 Subscription 中，由 Subscription 来调用 onNext 方法。

Hot 模式下数据是共享的，例如 Flux.interval(..).share()，记录了开始到现有的秒数，每个 Subscriber 在订阅时都希望得到当前秒数，而不是从第 1s 开始。于是秒数信息必须由 Publisher 保存，并且对 Subscriber 共享，此时 subscription.request就只是个传话筒了。

Reactor 与 Reactive Streams 规范

在流式代码中，通常只有一个数据源（例如调用某个 API），之后会对这个数据做一系列的 map, filter 等操作，每个这样的操作符，从逻辑上都可以等价于既是一个publisher 又是一个 subscriber。例如下面这样的代码:

var myPub = Flux.range(1, 10)
                .map(x -> x * 2)
                .filter(x -> x > 10);

myPub.subscribe(System.out::println);

首先是构建 publisher，的过程，每个操作符^[2]都会保留它的父 publisher：

于是当我们执行 myPub.subscribe^[3] 时，每个操作符本身作为一个 Subscriber，会不断调用父 Publisher 的 subscribe 方法；而父 Publisher 在调用 onSubscribe 时，每个操作符作为一个 Subscriber，会不断调用下一个操作符的 onSubscribe 方法:

而当 Subscriber 调用 request 方法时，也是相同的路径^[4]:

线程如何调度

上面我们讲解了如何组装流式代码以及它的内部执行流程，但这些代码是在哪个线程上执行的呢？我们知道对于同步代码，代码会在同一个线程上执行，于是上面的示例中，所有的调用都在同一个线程上：

图中的棕线代表线程。但是这个例子比较特殊，因为 Flux.range 的数据是就绪的，而如果需要使用诸如 WebClient 调用 API 后做处理，则涉及到异步调用 IO，此时则会是这样：

上图会假设 WebClient 调用了外部服务，当外部服务返回时会在另一个线程上执行回调函数 callback，而这个 callback 会调用 B.onNext 方法，以此类推后续的 onNext都会在这个线程上执行。

这就有大问题了！例如底层调用使用的是 Netty，则执行 callback 的线程一般就是Netty的 worker 线程，但现在我们必须在这个线程上执行所有的 onNext 方法，如果某个操作符（如某个 map）是 CPU 密集型的，就会导致该 worker 线程被长时间阻塞，此时 Netty 的 Worker 线程池成为瓶颈，造成其它子模块的请求没有 worker 线程能处理而卡死，子功能之间互相耦合、干扰。

无奈下的 subscribeOn 与 publishOn

为了解决上面的问题，Reactor 提供了 subscribeOn 和 publishOn 两个方法，可以分别影响 request 和 onNext 方法的执行线程。例如：

var myPub = Flux.range(1, 10)
                .map(x -> x * 2)
                .subscribeOn(Schedulers.elastic())
                .filter(x -> x > 10);
myPub.subscribe(System.out::println);

则执行的流程如下:

可以看到 subscribeOn 方法会影响 request 方法的执行线程，另外由于整个流程没有另外的线程切换（如上节提到的 WebClient），因此 onNext 方法也会在同一个线程执行。我们又知道诸如 map(x -> x * 2) 这样的操作是在 onNext 方法中执行的，于是也会在新的线程上执行。

由于 request 方法调用顺序从代码的视角是由下到上的，因此一般说 subscribeOn影响的是向上的调用链，直到 publishOn 或其它的 subscribeOn 方法为止。

同样的，publishOn 方法会影响 onNext 方法的执行线程，例如：

var myPub = Flux.range(1, 10)
                .map(x -> x * 2)
                .publishOn(Schedulers.elastic())
                .filter(x -> x > 10);
myPub.subscribe(System.out::println);

由于 onNext 方法调用顺序从代码的视角是由上到下的，因此一般说 publishOn影响的是向下的调用链，直到其它的 publishOn 为止。

但要注意，如果 subscribeOn 和 publishOn 同时存在，则 subscribeOn 的作用会“穿过” publishOn：

1
2
3
4
5
6

var myPub = Flux.range(1, 10)
                .map(x -> x * 2)
                .publishOn(Schedulers.elastic())
                .subscribeOn(Schedulers.elastic())
                .filter(x -> x > 10);
myPub.subscribe(System.out::println);

在这种情况下，第 2 行的 map 还是会被第 4 行的 subscribeOn 影响；而第 5行的 fitler 最终会被 publishOn 影响。

后记

学习这个线程模型距离我开始学习 Webflux 几乎有 4 年以上的时间了，在我自认为对Webflux 了解还算充分的时候被教育了。时至今日，我依然有两个暴论：

1. Webflux 只适用于诸如网关这样的业务简单但高并发的场景。
2. 对于绝大多数人来说，green thread 类型的异步模型才是最好的。

本文会从零构建一个 Java Agent，让Jar 包在运行时打印每一个调用的方法名，其中涉及到 Java Agent 的整体结构，ASM 库的基础操作，文章较长，建议跟着走一遍。

Java Agent 项目结构

先创建如下目录结构：

.
├── pom.xml
└── src
    └── main
        ├── java
        │   └── me
        │       └── lotabout
        │           └── Launcher.java
        └── resources
            └── META-INF
                └── MANIFEST.MF

premain 与 agentmain

我们知道常规 Java 程序的入口是 main 函数，而 Java Agent 在不同的架构模式下有不同的入口^[1]：

• 静态加载入口为 premain：如 java -javaagent:my-agent.jar -jar app.jar，在启动 Jar 包时指定要加载的 agent，权限较高。
• 动态加载入口为 agentmain：已经通过 java -jar app.jar 等方式运行的 JVM，可以动态 Attach 后加载 Agent，权限较低，如无法新增属性、方法等。

两个方法定义如下（定义放在哪个类中都可以，下面会在 MANIFEST.MF 文件中声明）：

package me.lotabout;
public class Launcher {
    public static void premain(String agentArgs, Instrumentation inst) {}
    public static void agentmain(String agentArgs, Instrumentation inst) {}
}

• 参数中的 agentArgs 是传递给 Agent 的参数。例如这样调用 java -javaagent:my-agent.jar=my-agent-args app.jar，则 my-agent.jar 中的premain 函数中的 agentArgs 参数的值，就是字符串 "my-agent-args"。
• 参数中的 Instrumentation 是 Java 提供的修改字节码的 API. 通常 Java Agent作者的任务，就是利用 Instrumentation 定位到希望修改的类并做出修改。

另外容易踩坑的一点是，调用 Instrumentation.addTransformer 添加的 transformer默认只对“ 未来加载的类 ”才会生效。而动态加载(agentmain)通常是在应用程序启动后才加载，就会出现添加的 transformer 不生效的情况。对静态加载(premain)则一般不会有这个问题，因为它是在 main 函数之前加载的，

动态加载(agentmain) 如果想修改 main 中就已经加载的类，则需要在添加transformer 再调用Instrumentation#retransformClasses 对已加载的类执行转换才能生效。

MANIFEST

上面提到 premain 和 agentmain 可以定义在任何类中，那 JVM 怎么知道去哪找呢？我们需要在 jar包的 MANIFEST.MF 文件^[2] 中指定 agent 的入口类是什么，以及 agent会有哪些能力：

Premain-Class: me.lotabout.Launcher # 静态加载(premain) Agent 时的入口类
Agent-Class: me.lotabout.Launcher   # 动态加载(agentmain) Agent 时的入口类
Can-Redefine-Classes: true          # 该 Agent 能否重新定义类
Can-Retransform-Classes: true       # 该 Agent 能否修改已有类
Can-Set-Native-Method-Prefix: true  # 是否允许修改 Native 方法的前缀

• Premain-Class: 静态加载(premain) Agent 时的入口类
• Agent-Class: 动态加载(agentmain) Agent 时的入口类
• Can-Redefine-Classes: 该 Agent 能否重新定义类
• Can-Retransform-Classes: 该 Agent 能否修改已有类
• Can-Set-Native-Method-Prefix: 是否允许修改 Native 方法的前缀。Native 方法不是字节码实现的，Agent 修改不了它的逻辑。通常修改 Native 是Proxy 的做法，把原有的 Native 方法重命名，新建同名的 Java 方法来调用老方法。此时需要修改Native 方法前缀的能力。

pom.xml

打包本身也比较烦，比如 maven 打包时需要指定 MANIFEST.MF 路径，示例如下：

<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
  xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
  xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
  <modelVersion>4.0.0</modelVersion>

  <groupId>me.lotabout</groupId>
  <artifactId>my-agent</artifactId>
  <version>1.0-SNAPSHOT</version>

  <properties>
    <maven.compiler.source>11</maven.compiler.source>
    <maven.compiler.target>11</maven.compiler.target>
    <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
  </properties>

  <dependencies>
    <dependency>
      <groupId>org.ow2.asm</groupId>
      <artifactId>asm</artifactId>
      <version>9.4</version>
    </dependency>
    <dependency>
      <groupId>org.ow2.asm</groupId>
      <artifactId>asm-tree</artifactId>
      <version>9.4</version>
    </dependency>
  </dependencies>

  <build>
    <plugins>
      <plugin>
        <groupId>org.apache.maven.plugins</groupId>
        <artifactId>maven-assembly-plugin</artifactId>
        <version>3.6.0</version>
        <configuration>
          <descriptorRefs>
            <descriptorRef>jar-with-dependencies</descriptorRef>
          </descriptorRefs>
          <archive>
            <manifestFile>src/main/resources/META-INF/MANIFEST.MF</manifestFile>
          </archive>
        </configuration>
        <executions>
          <execution>
            <phase>package</phase>
            <goals>
              <goal>single</goal>
            </goals>
          </execution>
        </executions>
      </plugin>
    </plugins>
  </build>
</project>

我们额外引入了 asm 与 asm-tree 库，我们后面要用它们来操作字节码。

mvn clean package 后得到 target/my-agent-1.0-SNAPSHOT-jar-with-dependencies.jar，之后就可以用java -javaagent:target/my-agent-1.0-SNAPSHOT-jar-with-dependencies.jar -jar app.jar 来调用了。这个名字太长了，后面我们写命令时会简写成 my-agent.jar。

动态加载 Attach

假设我们已经执行了 java -jar app.jar，希望加载 my-agent.jar，要怎么做？需要利用 Attach API。

1. 先得到 app.jar 进程的 PID，并 attach 得到 app.jar 的 VirtualMachine 实例：VirtualMachine vm = VirtualMachine.attach(PID);
2. 调用 VirtualMachine#loadAgent("my-agent.jar") 让 app.jar 进程加载 agent

为了方便上述操作，我们可以把这段逻辑写到 Launcher 的 main 函数中：

public static void main(String[] args)
    throws IOException, AttachNotSupportedException, AgentLoadException, AgentInitializationException {
    String pid = args[0];
    String path = Launcher.class.getProtectionDomain().getCodeSource().getLocation().getPath();
    VirtualMachine vm = VirtualMachine.attach(pid);
    try {
        vm.loadAgent(path);
    } finally {
        vm.detach();
    }
}

再在 MANIFEST.MF 中增加一行：

Main-Class: me.lotabout.Launcher

现在就可以使用 java -jar my-agent.jar <目标 PID> 来动态加载 agent 了。注意此时调用的是 agent 的 agentmain 方法。

Instrumentation

上节中的内容是建立 java agent 项目结构，目标是产出一个能被 JVM 识别的Agent。接下来的任务是找到 app.jar 中感兴趣的类并修改这些类的字节码。这些工作都要基于JDK 提供的 Instrumentation API。

Instrumentation API

Instrumentation的核心抽象是 ClassFileTransformer，对字节码的修改逻辑都在这个接口中实现，而 Instrumentation 接口则是用来添加、删除 transformer 的。Instrumentation 常见的使用流程（伪代码）为:

// 对于未加载的类，addTransformer 后就能生效
instrument.addTransformer(myTransformer, true);
// 对于已经加载的类，需要调用 retransformClasses 来触发修改
for (Class clazz: instrument.getAllLoadedClasses()) {
    if (needToTransform(clazz)) {
        instrument.retransformClasses(clazz);
    }
}

Instrumentation 的一些常用接口定义如下：

• getAllLoadedClasses() 获取所有加载的类，得到数组后我们可以自己筛选出关心的类
• redefineClasses(ClassDefinition... definitions) 使用参数中的类定义重新定义类
• retransformClasses(Class<?>... classes) 使用添加的 transformers 修改指定的类
• addTransformer(ClassFileTransformer transformer) 注册 transformer
• removeTransformer(ClassFileTransformer transformer) 注销 transformer

ClassFileTransformer

对字节码的修改逻辑需要定义在ClassFileTransformer的 transform 方法中，方法的签名如下：

byte[] transform(ClassLoader loader,
                 String className,
                 Class<?> classBeingRedefined,
                 ProtectionDomain protectionDomain,
                 byte[] classfileBuffer)
    throws IllegalClassFormatException {
}

• 通常我们会使用各种信息来过滤掉不感兴趣的类（不想修改就直接直接返回原字节码）。
• 核心输入输出是 class 二进制流(byte[])，即 transformer 假定字节码的修改是在二进制层面进行的。

直接修改类的二进制不是人能干的事，于是通常会使用一些库把 byte[] 转成一些库定义的结构，操作后再转回 byte[] 返回。下面是常用的一些库：

• asm JDK 内部也用了它，性能好，但 API 的抽象层度很低
• javaassist API 的抽象比 ASM 更高，更适合普通用户，支持直接写 Java 源码
• bytebuddy API 抽象度更高，例如有专门的 builder 来创建 Agent

ASM API 简介

ASM 核心 API

ASM^[3] 有两套 API： Event-Based 和 Tree-Based。简单来说 Event-Based 就是 visitor模式，用户需要定义各种元素的 visitor，扫描字节码中过程中遇到什么元素就调用对应元素的 visitor；Tree-Based 可以理解成先扫一遍字节码组装成一棵树，再对这棵树做后续编辑、修改等操作。Event-Based API 性能更好但 Tree-Based API 更容易理解和使用。

ASM 的整体流程是 byte[] -> ClassNode -> (修改) -> byte[]，其中ClassNode 是Tree-Based API 对“类”的抽象。基于 Tree-Based API 来修改字节码的 pattern 如下：

ClassNode cn = new ClassNode(ASM4);                       // 定义解析后的类
ClassReader cr = new ClassReader(origin_classfile_bytes); // 创建 reader 读取原始字节码
cr.accept(cn, 0);                                         // 解析原始字节码，填充到 cn 中
...                                                       // 这里可对 cn 做修改
ClassWriter cw = new ClassWriter(0);                      // 创建 writer
cn.accept(cw);                                            // 把修改后的 cn 写回到 writer
byte[] b = cw.toByteArray();                              // 把 writer 中的字节码转成 byte[]

对 ClassNode 的操作，最常见的是遍历其中的 cn.methods 属性来遍历该类的所有方法，之后通过修改 method.instructions 来修改字节码。

类型描述符(Type Descriptor)

ASM 中对于类型的描述有自己的一套规则，严格来说也不是 ASM 自创的，而是 JVM Spec中定义的^[4]，定义如下：

基本类型的描述符就是对应的大写字母（除了 boolean 用 Z 代替，因为字母冲突）；其中类的描述符是 L<classname>; 的格式，数组的描述符是 [<array_type>，如果多维就以此类推。

方法描述符(Method Descriptor)

方法描述符是一个字符串，格式为 (<参数类型1><参数类型2>...)<返回类型>，其中参数类型就是上节的类型描述符，如果返回 void 则写 V，例如：

示例-打印每个调用的方法

由于 ASM 的 API 基本是直接添加字节码，但如果对字节码不熟悉其实很难直接写出，于是一种方法是先用 javap 等工具把一个类的字节码反编译出来，再根据反编译的结果来写。

println 字节码

例如我们想在方法被调用时执行如下代码：

System.out.println(">> calling Method: <my_method>");

于是我们先写一个类，然后用 javap -c 来反编译：

class Hello {
    public static void main(String[] args) {
        System.out.println(">> calling Method: <my_method>");
    }
}

执行如下命令：

$ javac Hello.java
$ javap -c Hello
Compiled from "Hello.java"
class Hello {
  Hello();
    Code:
       0: aload_0
       1: invokespecial #1                  // Method java/lang/Object."<init>":()V
       4: return

  public static void main(java.lang.String[]);
    Code:
       0: getstatic     #2                  // Field java/lang/System.out:Ljava/io/PrintStream;
       3: ldc           #3                  // String >> calling Method: <my_method>
       5: invokevirtual #4                  // Method java/io/PrintStream.println:(Ljava/lang/String;)V
       8: return
}

比较关键的是 getstatic, ldc, invokevirtual 这三个指令，分别代表先获取System.out，再加载常量 ">> calling Method: <my_method>"，最后调用println 三个操作。

自定义 Transformer

接下来我们定义一个 ClassFileTransformer 来实现上述逻辑：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

public class MyTransformer implements ClassFileTransformer {

  private String prefixOfclassToPrint = "";

  public MyTransformer(String prefixOfclassToPrint) {
      this.prefixOfclassToPrint = prefixOfclassToPrint.replace(".", "/");
  }

  @Override
  public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
      ProtectionDomain protectionDomain, byte[] classfileBuffer) {
    if (!className.startsWith(this.prefixOfclassToPrint)) {
      return classfileBuffer;
    }

    System.out.println("transforming class: " + className);
    ClassNode cn = new ClassNode(Opcodes.ASM4);
    ClassReader cr = new ClassReader(classfileBuffer);
    cr.accept(cn, 0);

    for (var method : cn.methods) {
      System.out.println("patching Method: " + method.name);
      var list = new InsnList();
      list.add(new FieldInsnNode(Opcodes.GETSTATIC, "java/lang/System", "out",
          "Ljava/io/PrintStream;"));
      list.add(new LdcInsnNode(">> calling Method: " + method.name));
      list.add(new MethodInsnNode(Opcodes.INVOKEVIRTUAL, "java/io/PrintStream", "println",
          "(Ljava/lang/String;)V", false));
      method.instructions.insert(list);
    }

    ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_MAXS);
    cn.accept(cw);
    return cw.toByteArray();
  }
}

• 第 6 行要注意在 transformer 中拿到的类，包名是以 / 分隔的，而Instrument.getAllLoadedClasses() 中拿到的类名是以 . 分隔的我们允许传入的参数是以 . 分隔的，所以需要转换一下。
• 第 12~14 行是过滤掉不感兴趣的类，不感兴趣的类直接返回原字节码。
• 第 16~19, 32~34 行是上文所说的 ASM 框架代码，反序列化二进制和序列化二进制的过程。
• 第 21 行开始遍历该类的所有方法，每个方法都插入我们的逻辑
• 第 24~28 行是插入字节码的逻辑，对应上小节说的 getstatic, ldc,invokevirtual 三个指令。其中也看到了类型描述符、方法描述符的使用。
• 第 29 行是把生成的字节码 “insert” 到方法的字节码中，“insert” 是在最前面插入

组装与测试

最后，我们在 premain 和 agentmain 中注册我们的 MyTransformer，最终 Launcher 类如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68

public class Launcher {

  public static void main(String[] args)
      throws IOException, AttachNotSupportedException,
             AgentLoadException, AgentInitializationException {
    String pid = args[0];
    String prefix = args[1];
    String path = Launcher.class.getProtectionDomain().getCodeSource().getLocation().getPath();
    VirtualMachine vm = VirtualMachine.attach(pid);
    try {
      vm.loadAgent(path, prefix);
    } finally {
      vm.detach();
    }
  }

  public static void premain(String agentArgs, Instrumentation inst)
      throws UnmodifiableClassException {
    inst.addTransformer(new MyTransformer(agentArgs), true);
    for (var clazz : inst.getAllLoadedClasses()) {
      if (inst.isModifiableClass(clazz) && clazz.getName().startsWith(agentArgs)) {
        inst.retransformClasses(clazz);
      }
    }
  }

  public static void agentmain(String agentArgs, Instrumentation inst)
      throws UnmodifiableClassException {
    premain(agentArgs, inst);
  }

  private static class MyTransformer implements ClassFileTransformer {

    private String prefixOfclassToPrint = "";

    public MyTransformer(String prefixOfclassToPrint) {
        this.prefixOfclassToPrint = prefixOfclassToPrint.replace(".", "/");
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
        ProtectionDomain protectionDomain, byte[] classfileBuffer) {
      if (!className.startsWith(this.prefixOfclassToPrint)) {
        return classfileBuffer;
      }

      System.out.println("transforming class: " + className);
      ClassNode cn = new ClassNode(Opcodes.ASM4);
      ClassReader cr = new ClassReader(classfileBuffer);
      cr.accept(cn, 0);

      for (var method : cn.methods) {
        System.out.println("patching Method: " + method.name);
        var list = new InsnList();
        list.add(new FieldInsnNode(Opcodes.GETSTATIC, "java/lang/System", "out",
            "Ljava/io/PrintStream;"));
        list.add(new LdcInsnNode(">> calling Method: " + method.name));
        list.add(new MethodInsnNode(Opcodes.INVOKEVIRTUAL, "java/io/PrintStream", "println",
            "(Ljava/lang/String;)V", false));
        method.instructions.insert(list);
      }

      ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_MAXS);
      cn.accept(cw);
      return cw.toByteArray();
    }
  }
}

之后有两种调用方式（类名前缀的包名以 . 分隔）：

$ java -javaagent:my-agent-1.0-SNAPSHOT-jar-with-dependencies.jar=<类名前缀> -jar app.jar
$ java -jar my-agent-1.0-SNAPSHOT-jar-with-dependencies.jar <PID> <类名前缀>

对于如下的示例 Hello 类：

public class Hello {

  public static void main(String[] args) throws Exception {
    for (int i = 0; i < 2; i++) {
      outer();
    }
  }

  public static void outer() {
    test();
  }

  public static void test() {
    System.out.println("Hello world!");
  }
}

可以在运行时挂上 agent 来看到输出:

$ javac Hello.java
$ java -javaagent:my-agent-1.0-SNAPSHOT-jar-with-dependencies.jar=Hello Hello
transforming class: Hello
patching Method: <init>
patching Method: main
patching Method: outer
patching Method: test
>> calling Method: main
>> calling Method: outer
>> calling Method: test
Hello world!
>> calling Method: outer
>> calling Method: test
Hello world!

可以看到我们成功的在每个方法调用时打印了一行信息。

小结

本文介绍了 Java Agent 的基本代码结构，简单介绍了 ASM 库来修改字节码的方法，最后给出了示例，让 Agent 能动态修改类的方法，在方法开始处打印一行信息。

另外一些常见的字节码修改场景可以参考 ASM 的文档或使用其它字节码修改库。例如希望打印每个方法的返回值，理论上需要遍历每个方法的字节码，找到 return 指令，然后在该指令前插入打印指令，这种常见 pattern 通常都有库封装好，可以直接使用。

看到 plantegg 大佬的文章 MySQL线程池导致的延时卡顿排查中提到 MySQL 线程池中 oversubscribe 的行为。也看到有小伙伴wych42在尝试复现文章里提到的现象。

自己也尝试复现（基于 Percona 8.0.29-21），但现象和 wych42 的结果有细节上的差异，引申发现自己对 Percona 线程池模型的理解有问题，记录一下。

假想的 oversubscribe 模型

我们知道 oversubscribe 是用来限制 thread group 内同时运行的线程数量的，于是猜想工作原理（类比 Java 中的 ThreadPoolExecutor，oversubscribe 类比corePoolSize）：

1. 由 thread group 中的 listener 线程将任务从网络连接挪到 thread group 中的队列
2. 在需要时创建若干个 worker 来消费队列，此时如果数量超过 oversubscribe 则停止创建
3. worker 在空闲若干时间后退出

差异一：耗时不稳定

仿照在 wych42 的 执行方案一中，设置如下：

| thread_handling                         | pool-of-threads |
| thread_pool_high_prio_mode              | transactions    |
| thread_pool_high_prio_tickets           | 4294967295      |
| thread_pool_idle_timeout                | 60              |
| thread_pool_max_threads                 | 100000          |
| thread_pool_oversubscribe               | 1               |
| thread_pool_size                        | 1               |
| thread_pool_stall_limit                 | 500             |

• 执行SQL select sleep(2)
• 执行并发：8

按前一节的假想模型，预期 thread group 每次有 2 个线程执行(1+oversubscribe)，结果是每批次 2 个 SQL 输出，耗时分别是 2s, 4s, 6s, 8s。

实测发现并不总是符合预期，各种情况都有，比如会有两批次就跑完的(4 个SQL 2s 加上2 个 SQL 4s）：

16:05:08.765+08:00: thread 0 iteration 1 start
16:05:08.765+08:00: thread 7 iteration 1 start
16:05:08.765+08:00: thread 3 iteration 1 start
16:05:08.765+08:00: thread 4 iteration 1 start
16:05:08.765+08:00: thread 1 iteration 1 start
16:05:08.765+08:00: thread 5 iteration 1 start
16:05:08.766+08:00: thread 2 iteration 1 start
16:05:08.766+08:00: thread 6 iteration 1 start
16:05:10.985+08:00: thread 4 iteration 1 took 2220 ms
16:05:10.985+08:00: thread 0 iteration 1 took 2220 ms
16:05:10.985+08:00: thread 3 iteration 1 took 2220 ms
16:05:11.029+08:00: thread 5 iteration 1 took 2264 ms
16:05:11.029+08:00: thread 7 iteration 1 took 2264 ms
16:05:13.080+08:00: thread 6 iteration 1 took 4313 ms
16:05:13.080+08:00: thread 1 iteration 1 took 4314 ms
16:05:13.085+08:00: thread 2 iteration 1 took 4320 ms
16:05:13.086+08:00: thread 2 iteration 2 start
16:05:13.086+08:00: thread 6 iteration 2 start
16:05:13.086+08:00: thread 5 iteration 2 start
16:05:13.086+08:00: thread 3 iteration 2 start
16:05:13.086+08:00: thread 1 iteration 2 start
16:05:13.086+08:00: thread 4 iteration 2 start
16:05:13.086+08:00: thread 0 iteration 2 start
16:05:13.086+08:00: thread 7 iteration 2 start
16:05:15.167+08:00: thread 6 iteration 2 took 2081 ms
16:05:15.167+08:00: thread 1 iteration 2 took 2081 ms
16:05:17.165+08:00: thread 2 iteration 2 took 4080 ms
16:05:17.165+08:00: thread 3 iteration 2 took 4079 ms
16:05:19.170+08:00: thread 5 iteration 2 took 6084 ms
16:05:19.170+08:00: thread 4 iteration 2 took 6084 ms
16:05:21.164+08:00: thread 0 iteration 2 took 8078 ms
16:05:21.164+08:00: thread 7 iteration 2 took 8078 ms
16:05:21.165+08:00: thread 1 iteration 3 start
16:05:21.165+08:00: thread 4 iteration 3 start
16:05:21.165+08:00: thread 7 iteration 3 start
16:05:21.165+08:00: thread 2 iteration 3 start
16:05:21.165+08:00: thread 3 iteration 3 start
16:05:21.165+08:00: thread 6 iteration 3 start
16:05:21.165+08:00: thread 5 iteration 3 start
16:05:21.165+08:00: thread 0 iteration 3 start
16:05:23.252+08:00: thread 6 iteration 3 took 2087 ms
16:05:23.252+08:00: thread 7 iteration 3 took 2087 ms
16:05:25.271+08:00: thread 2 iteration 3 took 4106 ms
16:05:25.271+08:00: thread 4 iteration 3 took 4106 ms
16:05:27.252+08:00: thread 3 iteration 3 took 6087 ms
16:05:27.251+08:00: thread 1 iteration 3 took 6086 ms
16:05:29.243+08:00: thread 0 iteration 3 took 8078 ms
16:05:29.243+08:00: thread 5 iteration 3 took 8078 ms
16:05:29.244+08:00: thread 5 iteration 4 start
16:05:29.244+08:00: thread 2 iteration 4 start
16:05:29.244+08:00: thread 6 iteration 4 start
16:05:29.244+08:00: thread 7 iteration 4 start
16:05:29.244+08:00: thread 4 iteration 4 start
16:05:29.244+08:00: thread 1 iteration 4 start
16:05:29.244+08:00: thread 3 iteration 4 start
16:05:29.244+08:00: thread 0 iteration 4 start
16:05:31.342+08:00: thread 7 iteration 4 took 2098 ms
16:05:31.343+08:00: thread 5 iteration 4 took 2099 ms
16:05:33.324+08:00: thread 2 iteration 4 took 4080 ms
16:05:33.324+08:00: thread 1 iteration 4 took 4080 ms
16:05:33.411+08:00: thread 3 iteration 4 took 4167 ms
16:05:33.417+08:00: thread 6 iteration 4 took 4173 ms
16:05:35.424+08:00: thread 0 iteration 4 took 6180 ms
16:05:35.424+08:00: thread 4 iteration 4 took 6180 ms

这说明要么是我们的假想模型有问题，要么是 Percona 实现有 BUG。那实际情况是怎么样的呢？拉代码看半天也看不出所以然，只能自行编译并加了很多 debug 日志，大概是明白了。

Percona thread group 模型

这里只说明 thread group 内的机制（不考虑全局的限制）。

1. 线程有两种状态：active & waiting，执行 SQL 过程中阻塞则记为 waiting，如等锁或 SLEEP
2. 线程的角色分成 listener 和 worker。listener 将网络上的请求挪到队列中，worker 从队列中获取任务来执行。一些情况下 listener 会自己变成worker 执行任务，worker 发现没有 listener 时也会变成 listener
3. 每个 group 内部有两个队列，高优队列和低优队列。如默认模式下，处于 XA 事务、持有表锁等情形下会被认为高优^[1]
4. thread_pool_oversubscribe 用来限制同时运行的线程数量，它是通过限制从队列获取任务来达到目的：
   1. active 线程数 >= 1 + oversubscribe 时 worker 不取任务，直接休眠，取任务时额外考虑下面规则
   2. active + waiting 线程数 > 1 + oversubscribe 时 worker 不取任务，但只限制低优队列中的任务
5. 为了防止各种假死的情况，会有专门的定时线程，检测两次执行间是否有进展，如果没有进展则会创建新的 worker（且此时规则 4.1 失效）。如 listener 变成 worker后创建新的线程承担工作
6. worker 线程在空闲一段时间(thread_pool_idle_timeout)后会退出

有几个推论：

1. thread group 的线程数可能大于 1+oversubscribe。没有机制限制线程的生成
2. 同时运行的任务可能会大于 1+oversubscribe。一方面 listener 变成 worker 时接任务不通过队列，因此不受限制；另一方面 waiting 的任务不参与计算规则 4.1；再者高优任务不参与计算规则 4.2.
3. 即使允许执行，任务的开始时间也可能会有延迟。如当前 listener 在干活，新任务只能等定时任务生成新的 worker 来执行，运气差的，延时可能会接近2*thread_pool_stall_limit。

差异一解释

结合更新后的模型以及实际的 debug 日志，差异一解释如下：

1. SELECT SLEEP 语句在执行时，线程会变成 waiting 状态
2. 由于 active 线程为 0, 很多代码位置上会尝试创建新的线程
3. 新的 worker 线程由于规则 4.2 的限制会取不到任务，进入休眠
4. 但由于某些时刻所有线程都在做任务，没有 listener，此时 worker 不休眠，而进入listener 模式
5. 进入 listener 模式的的线程在一些情况（发现有新任务且已有的队列为空）下会决定自己执行任务
6. 当 listener 决定自己执行任务，它会直接从网络连接中获取任务而不经过任务队列，因此不受限制
7. listener 开始执行任务时变成 worker 角色，有可能重新触发情况 #4

差异一补充 case：高优队列不受规则 4.2 限制

已知锁表能让事务成为高优，我们把负载改成两句 SQL：

LOCK TABLES t? READ ; 这里每个线程锁不同的表
SELECT SLEEP(2)

测试的结果如下，可以看到在第一个迭代中创建了 N 个 worker，第二个迭代中每个worker 都实际执行了任务，因此结果只有一个批次，都是 2s 左右。

11:47:36.251+08:00: thread 0 iteration 0 start
11:47:36.251+08:00: thread 7 iteration 0 start
11:47:36.251+08:00: thread 5 iteration 0 start
11:47:36.251+08:00: thread 2 iteration 0 start
11:47:36.251+08:00: thread 1 iteration 0 start
11:47:36.251+08:00: thread 6 iteration 0 start
11:47:36.251+08:00: thread 4 iteration 0 start
11:47:36.251+08:00: thread 3 iteration 0 start
11:47:39.059+08:00: thread 0 iteration 0 took 2808 ms
11:47:39.162+08:00: thread 5 iteration 0 took 2911 ms
11:47:39.425+08:00: thread 7 iteration 0 took 3174 ms
11:47:41.282+08:00: thread 1 iteration 0 took 5031 ms
11:47:41.352+08:00: thread 6 iteration 0 took 5101 ms
11:47:41.463+08:00: thread 2 iteration 0 took 5212 ms
11:47:41.613+08:00: thread 4 iteration 0 took 5362 ms
11:47:43.532+08:00: thread 3 iteration 0 took 7281 ms
11:47:43.533+08:00: thread 3 iteration 1 start
11:47:43.533+08:00: thread 5 iteration 1 start
11:47:43.533+08:00: thread 7 iteration 1 start
11:47:43.533+08:00: thread 0 iteration 1 start
11:47:43.533+08:00: thread 2 iteration 1 start
11:47:43.533+08:00: thread 4 iteration 1 start
11:47:43.533+08:00: thread 1 iteration 1 start
11:47:43.533+08:00: thread 6 iteration 1 start
11:47:45.727+08:00: thread 3 iteration 1 took 2194 ms
11:47:45.751+08:00: thread 7 iteration 1 took 2218 ms
11:47:45.751+08:00: thread 2 iteration 1 took 2218 ms
11:47:45.751+08:00: thread 4 iteration 1 took 2218 ms
11:47:45.755+08:00: thread 5 iteration 1 took 2222 ms
11:47:45.756+08:00: thread 1 iteration 1 took 2223 ms
11:47:45.760+08:00: thread 0 iteration 1 took 2227 ms
11:47:45.765+08:00: thread 6 iteration 1 took 2232 ms

差异二：SELECT SLEEP 可能不是好负载

通过源码我们知道执行 SLEEP 的线程是 waiting 状态，会绕过某些 oversubscribe 的限制。我们尝试使用下面的负载：

MySQL> select benchmark(9999999, md5('when will it end?'));
1 row in set
Time: 2.079s

测试的结果就更看不出“批次”的模式了。当然由于多个任务并行执行，实际的耗时也增加了（3.8s）。

16:58:01.905+08:00: thread 2 iteration 0 start
16:58:01.905+08:00: thread 6 iteration 0 start
16:58:01.905+08:00: thread 1 iteration 0 start
16:58:01.905+08:00: thread 4 iteration 0 start
16:58:01.905+08:00: thread 7 iteration 0 start
16:58:01.905+08:00: thread 5 iteration 0 start
16:58:01.905+08:00: thread 3 iteration 0 start
16:58:01.905+08:00: thread 0 iteration 0 start
16:58:05.757+08:00: thread 0 iteration 0 took 3852 ms
16:58:06.679+08:00: thread 5 iteration 0 took 4774 ms
16:58:08.376+08:00: thread 1 iteration 0 took 6471 ms
16:58:10.425+08:00: thread 2 iteration 0 took 8520 ms
16:58:11.620+08:00: thread 6 iteration 0 took 9715 ms
16:58:13.604+08:00: thread 4 iteration 0 took 11699 ms
16:58:14.413+08:00: thread 3 iteration 0 took 12508 ms
16:58:16.843+08:00: thread 7 iteration 0 took 14938 ms

但是，我们预期仍是一个批次执行两个 SQL，为什么第二个请求 4774ms 才返回？下面我们看看在 Percona 中增加的 debug 信息，来了解内部工作的机制

   time         thread id       message
   16:58:02.197 123145417097216 command: 3: select benchmark(9999999, md5('when will it end?'));
   16:58:02.206 4863544832 add connection(active: 1, waiting: 0, stalled: 0)
① 16:58:02.899 123145390891008 check_stall #1> wake_or_create(active: 1, waiting: 0, stalled: 0)
   16:58:02.899 123145390891008 wake or create thread
   16:58:02.899 123145390891008 thread waked (active: 1, waiting: 0, stalled: 0)
   16:58:02.899 123145418162176 get_event > after wakeup(active: 2, waiting: 0, stalled: 0)
② 16:58:02.899 123145418162176 get_event poll (active: 2, waiting: 0, stalled: 0, oversubscribed: 1)
   16:58:02.899 123145418162176 get_event current listener(0)
③ 16:58:02.899 123145418162176 get_event become listener(active: 1, waiting: 0, stalled: 0)
   16:58:02.899 123145418162176 get_event become listener get lock(active: 1, waiting: 0, stalled: 0)
   16:58:02.899 123145418162176 listener #0(active: 1, waiting: 0, stalled: 0)
④ 16:58:03.402 123145390891008 check_stall #2> wake_or_create(active: 1, waiting: 0, stalled: 1)
   16:58:03.402 123145390891008 wake or create thread
   16:58:03.402 123145390891008 waked failed (active: 1, waiting: 0, stalled: 1)
   16:58:03.402 123145390891008 throttle create worker #2(active: 1, waiting: 0, stalled: 1)
   16:58:03.402 123145390891008 create worker called(active: 1, waiting: 0, stalled: 1)
   16:58:03.402 123145419227136 worker main start (active: 2, waiting: 0, stalled: 1)
   16:58:03.402 123145419227136 get_event start (active: 2, waiting: 0, stalled: 1)
   16:58:03.402 123145419227136 get_event poll (active: 2, waiting: 0, stalled: 1, oversubscribed: 0)
⑤ 16:58:03.402 123145419227136 queue_get #0 (active: 2, waiting: 0, stalled: 1, toomany: 0)
   16:58:03.402 123145419227136 queue_get #2 (active: 2, waiting: 0, stalled: 1)
   16:58:03.402 123145419227136 get_event connection = 8c148620(active: 2, waiting: 0, stalled: 1, oversubscribed: 0)
   16:58:03.402 123145419227136 get_event end (active: 2, waiting: 0, stalled: 1)
⑥ 16:58:03.402 123145419227136 command: 3: select benchmark(9999999, md5('when will it end?'));

• 由于第一个 listener 线程执行了第一个任务，① 处 check stall 线程触发，尝试创建一个新的 worker。
• 在 ② 处，该 worker 尝试获取任务，但因为此时 active 为 2 (>= 1+oversubscribed=2），触发了限制，因此从队列中获取不到任务。
• 接着 ③ 中，worker 发现当前没有 listener，于是自己成为 listener，但此时网络上没有数据，进入休眠。
• ④ 中 check stall 第二次唤醒，发现有 listener，但队列不为空，于是尝试唤醒或新建线程。此时没有 waiting 中的线程，于是创建新的线程。
• ⑤ 中新建的线程从队列中获取任务，虽然当前 oversubscribed，但由于状态是 stalled，于是不受规则 4.1 限制，而此时 (active+waiting = 2 <= 1+oversubscribed=2)，也不受规则 4.2 限制，于是获取任务并执行。看到 ⑥ 中执行命令，此时距离接受到命令过去了 1s+，也因此整个请求是 4s+。

这个例子给出两个信息：

1. waiting 和 active 的负载对 thread group 调度来说是有差异的
2. 由于创建线程的滞后性（由 check stall 定时线程），任务执行会有延迟，且延迟不低

能不能简化？

Percona 实际的线程模型显然没有我们假想模型简单，那能不能简化呢？

例如为什么不用全职 listener，listener 完成不参与执行任务？这是因为直接让listener 处理任务效率更高，listener 刚从等待网络中被唤醒，不需要从再唤醒一个worker，减少线程切换。但 listener 擅离职守会造成后续任务的延时，因此 listener一方面只在当前任务队列为空时才转为 worker，另一方面有定时的check_stall 线程来保底。但如差异二中看到的，还是会造成任务执行的延时^[2]。

再例如能不能只用一个 queue？早期的实现其实就没有区分高优低优队列，Percona 后来实现优先队列是为了缩短服务端内部的 XA 事务^[3]。对表锁的高优操作也是后来才添加的。

还有为什么不在创建线程时就限制总数不能超过 oversubscribed？（以下是猜想）oversubscribe 从设计来看不应该是一个硬限制，它要达到的目的是在全局限制线程数的前提下，防止某个 thread group 疯狂创建吃掉所有限额，造成其它 group 创建不了线程的情况。但是适当允许某个 group 创建超过 oversubscribe 的线程数是有助于提高整体效率的。而且绝对限死线程数也更可能造成 group 内的死锁，保持弹性能应对更多异常的情况。

参考

• Percona thread pool 对 threadpool 的行为有一些说明，但并不是很全面
• 添加 debug 信息的源码文件，有兴趣的可以自己编译验证
  • threadpool_unix.cc
  • sql_parse.cc
• MySQL 线程组实现文档 oversubscribed 实现不同，但其它的如线程创建方面可以参考

另外关于 Percona 线程机制的描述一搜一大把，可以结合本文案例理解。

Einstein notation

你没有看错，Einstein notation是那位著名的爱因斯坦发明的，用来对线性代数中求和的表示做的约定。我们还是看个例子^[1]：

$$y = \sum_{i=1}^{3}{c_i x_i} = c_1 x_1 + c_2 x_2 + c_3 x_3$$

如果一个下标（如 $i$）在公式中出现两次，则隐式地认为需要遍历它的所有可能性。上面公式可以简化成：

$$y = c_i x_i$$

于是矩阵乘法中，每个输出元素可以这样表示：

$$c_{ij} = \sum_{k}{a_{ik} b_{kj}} \implies c_{ik} = a_{ik} b_{kj}$$

Einsum

在 Numpy和 Pytorch 中都实现了类似的机制。einsum 函数的第一个参数就是把上节公式中的各个下标按a,b->c 的格式写下来：

代码实例如下：

In [6]: a = np.asarray(range(1,9)).reshape(2,4)

In [7]: a
Out[7]:
array([[1, 2, 3, 4],
       [5, 6, 7, 8]])

In [8]: b = np.asarray(range(1,9)).reshape(4,2)

In [9]: b
Out[9]:
array([[1, 2],
       [3, 4],
       [5, 6],
       [7, 8]])

In [10]: a @ b
Out[10]:
array([[ 50,  60],
       [114, 140]])

In [11]: np.einsum('ik,kj->ij', a, b)
Out[11]:
array([[ 50,  60],
       [114, 140]])

复杂应用

例如在 CNN 求卷积时，输入是 (n, c, ih, iw) 的矩阵，卷积权重是 (C, c, h, w)（这里大写字母代表输出维度）。可以把原图像按卷积大小的各个子图求出，得到 (n, c, H, W, h, w) 的输入矩阵，于是可以使用 einsum 直接求结果：

def conv2d(x, weight, stride=(1,1), padding=(0,0)):
    # x is a 4d matrix (n, c, ih, iw), where n is batchsize, c is input channel
    # weight is a 4d matrix (oc, c, h, w), where o_c is output channel
    # out is (n, oc, h, w)
    if padding != (0,0):
        p_h, p_w = padding
        x_padded = np.pad(x, ((0,0), (0,0), (p_h,p_h), (p_w, p_w)))
    else
        x_padded = x

    i_n, i_c, i_h, i_w = x_padded.shape
    s_h, s_w = stride
    wo_c, wi_c, w_h, w_w = weight.shape

    o_n = i_n
    o_c = wo_c
    o_h = (i_h - w_h) // s_h + 1
    o_w = (i_w - w_w) // s_w + 1

    view_shape = (i_n, i_c, o_h, o_w, w_h, w_w)
    view_strides = np.array([(i_c*i_h*i_w), (i_h*i_w) , s_h*i_w, s_w, i_w, 1]) * x_padded.itemsize
    submatrix = np.lib.stride_tricks.as_strided(x_padded, view_shape, view_strides)
    return np.einsum('ncHWhw,Cchw->nCHW', submatrix, weight, optimize=True)

看到 ncHWhw,Cchw->nCHW 的输入中，c, h, w 下标是重复的，按约定要遍历所有三个下标的元素相乘，要是裸写代码的话，类似下面这样：

n,c,H,W,h,w = submatrix.shape
C,c,h,w = weight.shape

result = 0
for cc in range(c):
    for hh in range(h):
        for ww in range(w):
            result += a[n, cc, H, W, hh, ww] * b[C, cc, hh, ww]
out[n, C, H, W] = result

性能

一般比裸写 for 循环是要快不少的（比如上面的卷积，比我自己裸写的快 3x~5x）。但比专门优化的肯定还是不能比的（pytorch 的 conv2d 是用 C++ 专门优化的，比相同的einsum 快 10x）。

另外 这篇文章 建议无脑开 Numpy 中的optimize 参数。

目标

如果有函数 $f(x_1, \cdots, x_n)$，我们要使用链式法则计算函数 $f$ 对所有输入$x_i$ 的偏导。我们记中间函数为 $v$，记 $\bar{v_i} = \frac{\partial f}{\partialv_i}$，则最核心的计算公式为：

$$\bar{v_i} = \frac{\partial f}{\partial v_i}= \sum_{j \in next(i)}{\overline{v_{i\to j}}}= \sum_{j \in next(i)}{\frac{\partial f}{\partial v_{j}} \frac{\partial v_{j}}{\partial v_i}= \sum_{j \in next(i)}{\overline{v_j} \frac{\partial v_{j}}{\partial v_i}}}$$

大家可以配合算法篇的图来理解：

整体思路

首先需要允许用户构建计算图，很自然地关心 3 个部分：

1. 节点。计算图中的节点代表了计算，如 sin 这样的函数，我们把它叫作算子(operator)。在 AD 的场景下，算子既要关心前向计算，也需要关心后向求导
2. 边。需要有办法找到算子的上下游，在 AD 中我们使用邻接表来表示^[1]
3. 边上流转的数据。边上流转的有正向的计算数据，逆向的偏导数据，我们会统一用张量(Tensor)来表示

要注意的是为了符合用户的使用习惯，我们并不是要求用户直接给出一个“节点” List，再给出一个“边” List。计算图是隐式构建的。因此实际上是 数据 --(来源)--> 节点 --(输入)--> 数据 这样的引用关系。

计算图构建好之后，需要有遍历引擎，按拓扑排序顺序，正向地、逆向地遍历所有节点，正向计算输出，逆向计算偏导。这里的执行引擎其实有很多可以优化的空间，比如多线程计算，多节点合并计算等，但本文里就是简单地走流程。

最终希望怎么使用呢？

x1 = Tensor(np.array([0.5]), requires_grad=True)
x2 = Tensor(np.array([0.5]), requires_grad=True)
v3 = sin(x1)
v4 = mul(x1, x2)
v5 = add(v3, v4)
grad = np.array([1])
v5.backward(grad)
print(x1.grad) # expected to be 1.37758
print(x2.grad) # expected to be 0.5

框架实现

Tensor

我们用张量 Tensor 来定义数据部分。代码如下：

class Tensor(object):
    """tensor"""
    def __init__(self, ndarray: NDArray, requires_grad=False, grad_fn=None):
        super(Tensor, self).__init__()
        self.ndarray = ndarray            # ①
        self.requires_grad = requires_grad  # ②
        self.grad_fn = grad_fn            # ③
        self.grad = None                  # ④
        self._grad_accmulator = None

    def is_leaf(self) -> bool:
        return self.requires_grad and self.grad_fn is None

    def backward(self, output_grad):
        if self.grad_fn is None:
            raise "backward could not be called if grad_fn is None"
        execute_graph(self.grad_fn, output_grad)

    def __str__(self):
        grad_info = f' grad_fn={self.grad_fn}' if self.grad_fn is not None else ''
        return f'tensor({self.ndarray}{grad_info})'

    def __repr__(self):
        return self.__str__()

① 中使用 numpy.ndarray 保存前向数据，直接使用 numpy 来减少复杂度，毕竟我们只关心 AD 部分

③ 的 grad_fn 可以理解成保存的是 Tensor 的来源算子。实际上当 Tensor 生成时，对应的数据就计算完成了，记录它的来源也没有意义，但由于后续还要反向计算偏导，才需要记录来源来反查。因此只有在 ② requires_grad = True 时才有记录的必要。

④ 的 grad 就是偏导的结果，即 $\bar{v_i}$ 的值。

Operator

首先算子既需要管前向计算，也需要关心后向求导，于是框架性的定义如下：

# 注意 Operator 里计算的都是 Tensor 内部的数据，即 NDArray
class Operator(object):
    def __init__(self):
        super(Operator, self).__init__()
        self.next_ops = [] # ①

    def forward(self, *args: Tuple[NDArray]) -> NDArray:
        raise NotImplementedError("Should be override by subclass")

    def backward(self, output_grad: Tuple[NDArray]) -> Union[NDArray, Tuple[NDArray]]:
        raise NotImplementedError("Should be override by subclass")

forward 代表前向计算，可以有多个输入。backward 则相反，给定输出的偏导，需要为每个输入输出一个偏导。即如果 $op = f(x, y)$，则 forward 输出的是$f(x, y)$ 的值，而 backward 输出为 $[\frac{\partial op}{\partial x}, \frac{\partial op}{\partial y}]$

但仅有两个计算方法是不够的，在 forward 计算时，算子还需要维护“边”的信息，在后向计算偏导时使用。①中的 next_ops 就是用来计算边的信息的，例如样例代码中，执行完 v5 = add(v3, v4) 后，内部信息如下图：

但我们不希望建图的操作在每个算子中都实现一遍，因此我们在父类上实现 __call__函数，在使用时用户不应该直接调用 forward 函数，而应该直接调用 __call__ 函数，实现如下：

def __call__(self, *args: Tuple[Tensor]) -> Tensor:
    grad_fn = None
    requires_grad = any((t.requires_grad for t in args)) # ①

    if requires_grad:
        # add edges
        for input in args:
            if input.is_leaf(): # ②
                if input._grad_accmulator is None:
                    input._grad_accmulator = OpAccumulate(input)
                self.next_ops.append(input._grad_accmulator)
            else:
                self.next_ops.append(input.grad_fn) # ③
        grad_fn = self

    inputs = [t.ndarray for t in args]
    output = self.forward(*inputs) # ④
    return Tensor(output, requires_grad=requires_grad, grad_fn=grad_fn) # ⑤

其中 ① 会将输入 Tensor 的 requires_grad 值传染给输出，算子任意输入 Tensor 中，只要有一个需要算梯度，则输出的 Tensor 也需要计算梯度。另外④中可以看出__call__ 就是 forward 方法的包装。注意到 forward 的输出是 ndarray，而因为算子输出也需要是 Tensor，因此在 ⑤ 中做了封装。

在构造计算图时，会将 input.grad_fn 指向的算子，加入 next_ops 中，如 ③ 所示。只有②的例外，如果输入本身就是叶子节点，则它的 grad_fn 没有指向任何节点，因此这里构造了一个特殊的 OpAccumulate 算子来累加并设置梯度，如下所示：

class OpAccumulate(Operator):
    def __init__(self, tensor):
        super(OpAccumulate, self).__init__()
        self.tensor = tensor
    def backward(self, grad):
        self.tensor.grad = Tensor(grad)
        return grad

计算图遍历

计算图是一个有向无环图（简称 DAG），DAG 遍历的重点是需要按拓扑排序遍历，在一个算子的所有输入都被满足时才能执行该算子的 backward 方法。于是我们先搞个辅助函数，按拓扑的顺序，统计每个算子依赖的输入个数。

def compute_dependencies(root):
    # deps: {op: num}
    deps = {}
    q = deque()
    traversed = {root}
    q.append(root)
    while len(q) != 0:
        cur = q.pop()
        if len(cur.next_ops) == 0:
            continue
        for next in cur.next_ops:
            deps[next] = deps.get(next, 0) + 1
            if next not in traversed:
                q.append(next)
                traversed.add(next)
    return deps

在样例代码里，最终会以 root = op:+ 来调用，因此它会返回类似如下信息（当然key 会是各个实例化的算子，而不是字符串）：

{
  "op:+": 1,
  "op:sin": 1,
  "op:*": 1,
  "op:acc|x1": 2,
  "op:acc|x2": 1,
}

接下来我们会遍历整个图：

def execute_graph(root, output_grad):
    deps = compute_dependencies(root)
    inputs = {root: output_grad}  # ①

    q = deque()
    q.append(root)
    while len(q) != 0:
        task = q.pop()
        input = inputs[task]
        outputs = task.backward(input)
        if not isinstance(outputs, collections.abc.Sequence):
            outputs = [outputs]

        for next_op, output in zip(task.next_ops, outputs):
            if next_op is None:
                continue

            # accumulate the "inputs" for next_op # ②
            op_input = inputs.get(next_op, 0)
            inputs[next_op] = op_input + output

            deps[next_op] -= 1
            if deps[next_op] == 0: # ③
                q.append(next_op)

这个遍历过程可说的内容也不多，就是将 ready 的算子一个个放进队列 q 中，一个个执行它们的 backward 方法。其中比较关键的是，如果算子 backward 的输入如果有多个，则需要在 ① 中缓存部分输入，并且在 ② 中当新的输入到来需要进行累加，这里对应了开头公式 $\bar{v_i} = \sum_{j \in next(i)}{\overline{v_{i\to j}}}$ 的部分。最后在 ③ 中，要确保目标算子的所有输入都计算完成，才认为目标算子 ready 了。

如此，所有“框架”层面的内容均实现完毕。

具体算子

有了框架还不够，还需要实现算子，而实现算子最关键的是可能需要在 forward 过程中记录输入信息，在 backward 中用来计算偏导。例如文章开头的样例中 $\bar{v_2}= \bar{v_4} v_1$ 就需要在 forward 时记录 $v_1$ 的值。下面补齐示例中需要的几个算子

另外注意下面的代码中除了实现算子，我们还实现了诸如 add, mul 等函数，方便对Tensor 构建计算图。

元素加法

class OpEWiseAdd(Operator):
    # func: y = a + b
    # deri: y'/a' = 1
    # deri: y'/b' = 1
    def forward(self, a: NDArray, b: NDArray):
        return a + b
    def backward(self, grad: NDArray):
        ret = grad, grad
        return ret

def add(a, b):
    return OpEWiseAdd()(a, b)

元素乘法

class OpEWiseMul(Operator):
    # func: y = a * b
    # deri: y'/a' = b
    # deri: y'/b' = a
    def forward(self, a: NDArray, b: NDArray):
        self.a = a
        self.b = b
        return a * b
    def backward(self, grad: NDArray):
        return self.b * grad, self.a * grad

def mul(a, b):
    return OpEWiseMul()(a, b)

sin

class OpSin(Operator):
    # func: y = sin(x)
    # deri: y'/x' = cos(x)
    def forward(self, x: NDArray):
        self.x = x
        return np.sin(x)
    def backward(self, grad: NDArray):
        ret = np.cos(self.x) * grad
        return ret

def sin(x):
    return OpSin()(x)

向量与实验

样例实跑

>>> x1 = Tensor(np.array([0.5]), requires_grad=True)
>>> x2 = Tensor(np.array([0.5]), requires_grad=True)
>>> v3 = sin(x1)
>>> v4 = mul(x1, x2)
>>> v5 = add(v3, v4)
>>> grad = np.array([1])
>>> v5.backward(grad)
>>> print(x1.grad)
tensor([1.37758256])
>>> print(x2.grad)
tensor([0.5])

大家可以算算，跟公式算出来是一样的

扩展到向量

如果 $x_1, x_2$ 是向量呢？这里关系到向量的求导到底要怎么算，但整体来说，咱们实现的框架还是成立的。例如上面例子中的 +, *, sin，如果都只考虑是按元素的操作（不涉及矩阵乘法），则上面的算子定义依旧适用，下面我们对应在 Pytorch 运行的结果和我们刚实现的框架的结果：

#------------------- torch -------------------------|====================== Ours ========================
>>> import torch
>>> x1 = torch.tensor([0.0140, 0.5773, 0.0469],      >>> x1 = Tensor(np.array([0.0140, 0.5773, 0.0469]),
        requires_grad=True)                                  requires_grad=True)
>>> x2 = torch.tensor([0.3232, 0.4903, 0.9395],      >>> x2 = Tensor(np.array([0.3232, 0.4903, 0.9395]),
        requires_grad=True)                                  requires_grad=True)
>>> v3 = torch.sin(x1)                               >>> v3 = sin(x1)
>>> v4 = torch.mul(x1, x2)                           >>> v4 = mul(x1, x2)
>>> v5 = torch.add(v3, v4)                           >>> v5 = add(v3, v4)
>>> grad = torch.tensor([0.4948, 0.8746, 0.7076])    >>> grad = np.array([0.4948, 0.8746, 0.7076])
>>> v5.backward(grad)                                >>> v5.backward(grad)
>>> print(x1.grad)                                   >>> print(x1.grad)
tensor([0.6547, 1.1617, 1.3716])                     tensor([0.65467087 1.16167806 1.37161212])
>>> print(x2.grad)                                   >>> print(x2.grad)
tensor([0.0069, 0.5049, 0.0332])                     tensor([0.0069272  0.50490658 0.03318644])

小结

本文中我们实现了一个自动微分（Automatic Differentiation）的框架。主要是 Tensor、Operator 的定义，以及后向计算的引擎。

整体的实现和 PyTorch 的实现是比较类似的，但为了示例简单也做了些取舍。如Pytorch 中 Operator 的第一个参数是 ctx，也鼓励算子把信息记录在 ctx 中，但我们是直接用 self.x 来记录；再如 PyTorch 中在计算结束后会把计算图销毁，我们没有做；再有 PyTorch 在 Tensor 中重载了一些基本操作（如 + - * /），方便操作，但我们直接额外定义了 add, mul 等函数。等等等等。

总的来说，希望通过 AD 的简单实现，让大家认识到机器学习背后的一些原理，实际上也并没有特别复杂。当然我们也要认识到，能 Work 距离能在工业上使用，中间还隔了个太平洋。

参考

• CMU 的课程 Deep Learning Systems: Algorithms and Implementation
  • Lecture 4 - Automatic Differentiation AD 算法讲解
  • Lecture 5 - Automatic Differentiation Implementation AD 算法实现，着重讲解了诸如 Tensor，Operator 部分，图遍历的部分留做作业了。
  • 5_automatic_differentiation_implementation.ipynb Lecture 5 的部分代码
• PyTorch源码浅析(4)：Autograd PyTorch 源码解析，版本比较老，但整体逻辑依旧适用

点积(dot product)

向量点积

Dot product 运算仅定义在两个向量上，输出一个标量，也称为 “scalar product”。

坐标定义：假设有两个向量$\color {red}{\mathbf {a} =[a_{1},a_{2},\cdots ,a_{n}]}$ 和$\color {blue}{\mathbf {b} =[b_{1},b_{2},\cdots ,b_{n}]}$ ^[1]，则 dot product 定义为：

$$\mathbf {\color {red}{a}} \cdot \mathbf {\color {blue}{b}}= \sum_{i=1}^{n}{\color {red}{a}_{i} \color {blue}{b}_{i}}={\color {red} a_1}{\color {blue}b_1}+{\color {red}a_2}{\color {blue}b_2}+\cdots +{\color {red}a_n}{\color {blue}b_n}$$

还可以把 dot product 理解成是矩阵的线性变换，写成矩阵乘法，此时$\color{red}{\mathbf{a}}$ 与 $\color{blue}{\mathbf{b}}$ 都是列向量，定义如下：

$$\mathbf {\color {red}{a}} \cdot \mathbf {\color {blue}{b}}= \begin{bmatrix}{\color{red} a_1} \\\vdots \\ {\color{red}a_n} \end{bmatrix}\cdot\begin{bmatrix}{\color{blue}b_1} \\\vdots \\ {\color{blue}b_n} \end{bmatrix}= {\color {red} a_1}{\color {blue}b_1}+{\color {red}a_2}{\color {blue}b_2}+\cdots +{\color {red}a_n}{\color {blue}b_n}= \begin{bmatrix}{\color{red}a_1} &\cdots & {\color{red}a_n}\end{bmatrix}\begin{bmatrix}{\color{blue}b_1} \\\vdots \\ {\color{blue}b_n} \end{bmatrix}= \mathbf {\color {red}{a}} ^T \mathbf {\color {blue}{b}}$$

矩阵与矩阵点积

严格来说，点积的输入只能是两个向量，不存在矩阵和矩阵，矩阵和向量的点积，但矩阵计算方便，人们扩充了定义。先看矩阵和矩阵，可以认为矩阵就是列向量的集合，因此点积就是列向量分别做点积^[2]。

$${\color{red}\mathbf {A}} \cdot {\color{blue}\mathbf {B}}= \begin{bmatrix}{\color{red}\mathbf{a}_1} & \cdots & {\color{red}\mathbf{a}_n}\end{bmatrix}\cdot\begin{bmatrix}{\color{blue}\mathbf{b}_1} &\cdots & {\color{blue}\mathbf{b}_n}\end{bmatrix}= \begin{bmatrix}{\color{red}\mathbf{a}_1} \cdot {\color{blue}\mathbf{b}_1}& \cdots& {\color{red}\mathbf{a}_n} \cdot {\color{blue}\mathbf{b}_n}\end{bmatrix}$$

上式中的 $\mathbf{a}_i, \mathbf{b}_i$ 都是列向量。另外注意由于 $\mathbf{a}_i\cdot \mathbf{b}_i$ 的结果是标量，所以最终的结果是一个行向量。

矩阵与向量点积

矩阵和向量的点积本质上是将向量扩充再当成矩阵和矩阵的点积，定义如下：

$${\color{red}\mathbf {A}_{n \times m}} \cdot {\color{blue}\mathbf {v}}= \begin{bmatrix}{\color{red}\mathbf{a}_1} & \cdots & {\color{red}\mathbf{a}_m}\end{bmatrix}\cdot\begin{bmatrix}{\color{blue}\mathbf{v}} &\cdots & {\color{blue}\mathbf{v}}\end{bmatrix}= \begin{bmatrix}{\color{red}\mathbf{a}_1} \cdot {\color{blue}\mathbf{v}}& \cdots& {\color{red}\mathbf{a}_m} \cdot {\color{blue}\mathbf{v}}\end{bmatrix}$$

此时结果为行向量。考虑到向量的点积也可以写成矩阵乘法的形式${\color{red}\mathbf{a}} \cdot {\color{blue}\mathbf{b}} = {\color{red}\mathbf{a}^T} {\color{blue}\mathbf{b}}$，因此有

$$\begin{align}{\color{red}\mathbf {A}_{n \times m}} \cdot {\color{blue}\mathbf {v}}&= \begin{bmatrix}{\color{red}\mathbf{a}_1} \cdot {\color{blue}\mathbf{v}}& \cdots& {\color{red}\mathbf{a}_m} \cdot {\color{blue}\mathbf{v}}\end{bmatrix}= \begin{bmatrix}{\color{red}\mathbf{a}_1^T} {\color{blue}\mathbf{v}}& \cdots& {\color{red}\mathbf{a}_m^T} {\color{blue}\mathbf{v}}\end{bmatrix}\\({\color{red}\mathbf {A}} \cdot {\color{blue}\mathbf {v}})^T&= \begin{bmatrix}{\color{red}\mathbf{a}_1^T} {\color{blue}\mathbf{v}}\\ \cdots\\ {\color{red}\mathbf{a}_m^T} {\color{blue}\mathbf{v}}\end{bmatrix}=\begin{bmatrix}{\color{red}\mathbf{a}_1^T} \\ \vdots \\ {\color{red}\mathbf{a}_m^T}\end{bmatrix} {\color{blue}\mathbf{v}}= {\color{red}(\mathbf{A}^T)_{m \times n}}{\color{blue}\mathbf{v}_{n \times 1}}\\{\color{red}\mathbf {A}} \cdot {\color{blue}\mathbf {v}}&= \big({\color{red}\mathbf{A}^T}{\color{blue}\mathbf{v}}\big)^T= {\color{blue}\mathbf{v}^T}{\color{red}\mathbf{A}}\end{align}$$

当然，上述式子中，我们严格按数学上的定义：向量就是“列”向量。这个假设不太方便，因为输入 $\mathbf{x}$ 是列向量，但输出 $\mathbf{A} \cdot \mathbf{x}$ 却是行向量。但实际上为了方便，我们也可以按“列”来排列输出结果，例如在深度学习中，单个输出 $y_i = \mathbf{w_i} \cdot \mathbf{x} + b$，则结果列向量：

$$\mathbf{y}_{m \times 1}= \begin{bmatrix} y_1 \\ y_2 \\ \vdots \\ y_m \end{bmatrix}= \begin{bmatrix}\mathbf{w_1} \cdot \mathbf{x} + b \\\mathbf{w_2} \cdot \mathbf{x} + b \\\vdots \\\mathbf{w_m} \cdot \mathbf{x} + b\end{bmatrix}= \begin{bmatrix}\mathbf{w_1}^T \mathbf{x} + b \\\mathbf{w_2}^T \mathbf{x} + b \\\vdots \\\mathbf{w_m}^T \mathbf{x} + b\end{bmatrix}= \begin{bmatrix}\mathbf{w_1}^T \\\mathbf{w_2}^T \\\vdots \\\mathbf{w_m}^T\end{bmatrix} \mathbf{x} + b= \mathbf{W}^T \mathbf{x} + b= (\mathbf{W}^T)_{m \times n} \mathbf{x}_{n \times 1} + b$$

因此不管是按行还是按列切片，关键在于点积 dot product 可以转换成矩阵乘法的形式。

矩阵乘法

矩阵与向量相乘

考虑一个矩阵 $\mathbf{A} \in \mathbb{R}^{m \times n}$ 和向量 $\mathbf{x} \in \mathbb{R}^n$，矩阵和向量的乘法定义为^[3]：

$$\mathbf{A}\mathbf{x} = x_1 \mathbf{a}_{*,1} + x_2 \mathbf{a}_{*,2} + \cdots + x_n \mathbf{a}_{*,n}$$

其中 $\mathbf{a}_{*, i}$ 代表矩阵 $\mathbf{A}$ 的第 $i$ 个列向量。

矩阵乘法有几种不同的理解方式，其中一种理解方式是“线性变换”^[4]，即向量 $\mathbf{x}$ 所在空间的基坐标，分别经过变换后，$\mathbf{x}$ 所在的坐标。因此，跟上述的定义基本一致：

$$\mathbf{A}{\color{brown}\mathbf{x}} =\begin{bmatrix}{\color {red} a_{11}} & {\color {blue} a_{12}} & \cdots & {\color {green}a_{1n}} \\{\color {red} a_{21}} & {\color {blue} a_{22}} & \cdots & {\color {green}a_{2n}} \\\vdots & \vdots & \ddots & \vdots \\{\color {red} a_{m1}} & {\color {blue} a_{m2}} & \cdots & {\color {green}a_{mn}}\end{bmatrix}\begin{bmatrix}{\color {brown} x_1} \\ {\color {brown} x_2} \\ \vdots \\ {\color {brown} x_n}\end{bmatrix}= {\color{brown} x_1} \begin{bmatrix}{\color {red} a_{11}} \\{\color {red} a_{21}} \\\vdots \\{\color {red} a_{m1}}\end{bmatrix} +{\color{brown} x_2} \begin{bmatrix}{\color {blue} a_{12}} \\{\color {blue} a_{22}} \\\vdots \\{\color {blue} a_{m2}}\end{bmatrix} + \cdots +{\color{brown} x_n} \begin{bmatrix}{\color {green} a_{1n}} \\{\color {green} a_{2n}} \\\vdots \\{\color {green} a_{mn}}\end{bmatrix}=\begin{bmatrix}{\color{brown}x_1}{\color {red} a_{11}} + {\color{brown}x_2}{\color {blue} a_{12}} + \cdots + {\color{brown}x_n}{\color {green}a_{1n}} \\{\color{brown}x_1}{\color {red} a_{21}} + {\color{brown}x_2}{\color {blue} a_{22}} + \cdots + {\color{brown}x_n}{\color {green}a_{2n}} \\\vdots \\{\color{brown}x_1}{\color {red} a_{m1}} + {\color{brown}x_2}{\color {blue} a_{m2}} + \cdots + {\color{brown}x_n}{\color {green}a_{mn}}\end{bmatrix}$$

还有一种理解和上面的“点积”类似，把矩阵看作是 $m$ 个行向量，每个向量都和 $x$ 作点积。如下：

$$\mathbf{A}{\color{brown}\mathbf{x}} =\begin{bmatrix}{\color {red} a_{11}} & {\color {red} a_{12}} & \cdots & {\color {red}a_{1n}} \\{\color {blue} a_{21}} & {\color {blue} a_{22}} & \cdots & {\color {blue}a_{2n}} \\\vdots & \vdots & \ddots & \vdots \\{\color {green} a_{m1}} & {\color {green} a_{m2}} & \cdots & {\color {green}a_{mn}}\end{bmatrix}\begin{bmatrix}{\color {brown} x_1} \\ {\color {brown} x_2} \\ \vdots \\ {\color {brown} x_n}\end{bmatrix}= \begin{bmatrix}{\color {red} \mathbf{a}_{1,*}} \cdot {\color{brown}\mathbf{x}} \\{\color {blue} \mathbf{a}_{2,*}} \cdot {\color{brown}\mathbf{x}} \\\vdots \\{\color {green} \mathbf{a}_{m,*}} \cdot {\color{brown}\mathbf{x}}\end{bmatrix}$$

矩阵与矩阵乘法

延续线性变换的观点，矩阵和矩阵的相乘，可以看作变换的组合。矩阵 $\mathbf{B}$ 的每个列可以认为是变换后的基向量，而 $\mathbf{A} \mathbf{B}$ 可以认为是把每个基向量再做一次线性变换 $\mathbf{A}$。如下：

$$\mathbf{A}\mathbf{B}= \mathbf{A} \begin{bmatrix}{\color{red}\mathbf{b}_{*,1}} & {\color{blue}\mathbf{b}_{*,2}} & \cdots & {\color{green}\mathbf{b}_{*,n}}\end{bmatrix}= \begin{bmatrix}\mathbf{A} {\color{red}\mathbf{b}_{*,1}}& \mathbf{A} {\color{blue}\mathbf{b}_{*,2}}& \cdots& \mathbf{A} {\color{green}\mathbf{b}_{*,n}}\end{bmatrix}$$

如果再用上面的“点积”观点展开，则会是这样：

$$\mathbf{A}\mathbf{B}= \mathbf{A} \begin{bmatrix}\mathbf{b}_{*,1} & \mathbf{b}_{*,2} & \cdots & \mathbf{b}_{*,n}\end{bmatrix}= \begin{bmatrix}\begin{bmatrix}\mathbf{a}_{1,*} \\\mathbf{a}_{2,*} \\\vdots \\\mathbf{a}_{m,*}\end{bmatrix}\mathbf{b}_{*,1}&\begin{bmatrix}\mathbf{a}_{1,*} \\\mathbf{a}_{2,*} \\\vdots \\\mathbf{a}_{m,*}\end{bmatrix}\mathbf{b}_{*,2}& \cdots&\begin{bmatrix}\mathbf{a}_{1,*} \\\mathbf{a}_{2,*} \\\vdots \\\mathbf{a}_{m,*}\end{bmatrix}\mathbf{b}_{*,n}\end{bmatrix}= \begin{bmatrix}\mathbf{a}_{1,*} \cdot \mathbf{b}_{*,1} & \mathbf{a}_{1,*} \cdot \mathbf{b}_{*,2} & \cdots & \mathbf{a}_{1,*} \cdot \mathbf{b}_{*,n} \\\mathbf{a}_{2,*} \cdot \mathbf{b}_{*,1} & \mathbf{a}_{2,*} \cdot \mathbf{b}_{*,2} & \cdots & \mathbf{a}_{2,*} \cdot \mathbf{b}_{*,n} \\\vdots & \vdots & \ddots & \vdots \\\mathbf{a}_{m,*} \cdot \mathbf{b}_{*,1} & \mathbf{a}_{m,*} \cdot \mathbf{b}_{*,2} & \cdots & \mathbf{a}_{m,*} \cdot \mathbf{b}_{*,n}\end{bmatrix}$$

这个也就是我们熟悉的，每个元素等于行乘列的形式：

$$(\mathbf{A}\mathbf{B})_{ij}= {\color{red}\mathbf{a}_{i,*}} \cdot {\color{blue}\mathbf{b}_{*,j}}= \begin{bmatrix}\cdots & \cdots & \cdots & \cdots \\{\color{red}a_{i1}} & {\color{red}a_{i2}} & \cdots & {\color{red}a_{in}} \\\cdots & \cdots & \cdots & \cdots\end{bmatrix}\begin{bmatrix}\vdots & {\color{blue}b_{1,j}} & \vdots \\\vdots & {\color{blue}b_{2,j}} & \vdots \\\vdots & \vdots & \vdots \\\vdots & {\color{blue}b_{n,j}} & \vdots\end{bmatrix}$$

按元素操作

即将两个矩阵对应位置上的元素做相应的操作。也称为 element-wise operation.

按元素乘法

按元素乘法有个特殊的名字，叫 Hadamard product，一般记作 $A \circ B$ 或 $A \odot B$，即将相同位置的元素相乘

$$\begin{bmatrix}a_{11} & a_{12} & \cdots & a_{1n} \\a_{21} & {\color{red}a_{22}} & \cdots & a_{2n} \\\vdots & \vdots & \ddots & \vdots \\a_{m1} & a_{m2} & \cdots & a_{mn}\end{bmatrix}\circ\begin{bmatrix}b_{11} & b_{12} & \cdots & b_{1n} \\b_{21} & {\color{blue}b_{22}} & \cdots & b_{2n} \\\vdots & \vdots & \ddots & \vdots \\b_{m1} & b_{m2} & \cdots & b_{mn}\end{bmatrix}= \begin{bmatrix}a_{11}b_{11} & a_{12}b_{12} & \cdots & a_{1n}b_{1n} \\a_{21}b_{21} & {\color{red}a_{22}}{\color{blue}b_{22}} & \cdots & a_{2n}b_{2n} \\\vdots & \vdots & \ddots & \vdots \\a_{m1}b_{m1} & a_{m2}b_{m2} & \cdots & a_{mn}b_{mn}\end{bmatrix}$$

其它操作也类似，如加法减法等

广播 broadcast

数学上，按元素操作只在矩阵的形状相同时才有效。但在实际应用中，可以尝试把“低维”的元素复制 N 份做填充，这就是 broadcast 机制。其实在介绍矩阵和向量的点积是已经用过这个操作了。示例如下：

$$\mathbf{A} \circ \mathbf{v}= \mathbf{A} \circ \begin{bmatrix}\mathbf{v} & \cdots & \mathbf{v} \end{bmatrix}$$

另外注意，这里依旧是以“列”为第一维，“行”为第二维。而在 numpy 中，第 0 维是行：

>>> w = np.array([[1,2,3],[4,5,6],[7,8,9]])
>>> w
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])
>>> x = np.array([1,2,3])
>>> x
array([1, 2, 3])
>>> w * x
array([[ 1,  4,  9],
       [ 4, 10, 18],
       [ 7, 16, 27]])

这个特性可以一直往高维推广，具体机制参考 numpy broadcast。

求导

深度学习中最重要的数学知识就是对矩阵求导了，The Matrix Calculus You Need For Deep Learning这篇论文针对性地做了综述。下面的知识算是摘录其中一些部分加深记忆^[5]。矩阵求导更复杂的内容，参考 wiki: Matrix calculus

向量点积求导

考虑 $y = \mathbf{w} \cdot \mathbf{x}$，因为有多个输入，于是导数为偏导向量，这里我们用行向量表示：

$$\frac{\partial y}{\partial \mathbf{x}}= \begin{bmatrix}\frac{\partial y}{\partial x_1} &\cdots &\frac{\partial y}{\partial x_n}\end{bmatrix}= \begin{bmatrix}\frac{\partial (w_1 x_1 + \cdots + w_n x_n)}{\partial x_1} &\cdots &\frac{\partial (w_1 x_1 + \cdots + w_n x_n)}{\partial x_n}\end{bmatrix}= \begin{bmatrix} w_1 & \cdots & w_n \end{bmatrix}= \mathbf{w}^T$$

同理对 $\mathbf{w}$ 求导的值为：

$$\frac{\partial y}{\partial \mathbf{w}}= \mathbf{x}^T$$

Jacobian Matrix

我们知道“导数”要求的是“变化”，即如果输入 $x$ 有微小的变化 $\Delta x$ 时，输入$y$ 的变化 $\Delta y$。那么如果有多个输入 $x_1, \cdots, x_n$ 和多个输出 $y_1 =f_1(\mathbf{x}), \cdots, y_m = f_m(\mathbf{x})$，则任意输入 $x_i$ 有变化，任意输出 $y_j$ 就有可能有变化。于是它们间的偏导关系是一个矩阵，记做：

$$\mathbf {J}=\begin{bmatrix}\frac{\partial y_1}{\partial \mathbf{x}} \\ \vdots \\ \frac{\partial y_m}{\partial \mathbf{x}} \end{bmatrix}=\begin{bmatrix}{\dfrac {\partial f_{1}}{\partial x_{1}}}&\cdots &{\dfrac {\partial f_{1}}{\partial x_{n}}}\\\vdots &\ddots &\vdots \\{\dfrac {\partial f_{m}}{\partial x_{1}}}&\cdots &{\dfrac {\partial f_{m}}{\partial x_{n}}}\end{bmatrix}$$

矩阵与向量点积求导

如果有 $m$ 个输出，$y_j = \mathbf{w_j} \cdot \mathbf{x}$（注意 $\mathbf{w_j}$本身是 $n$ 维的向量，向量个数是 $m$）。对 $\mathbf{x}$ 的求导比较直观：

$$\frac{\partial \mathbf{y}}{\partial \mathbf{x}}= \begin{bmatrix}\frac{\partial y_1}{\partial \mathbf{x}} \\\frac{\partial y_2}{\partial \mathbf{x}} \\\vdots \\\frac{\partial y_m}{\partial \mathbf{x}}\end{bmatrix}= \begin{bmatrix}\mathbf{w}_1^T \\\mathbf{w}_2^T \\\vdots \\\mathbf{w}_m^T\end{bmatrix}$$

但是，如果对 $\mathbf{w}$ 求导，$\mathbf{w}$ 有 $m \times n$ 个元素，因此求导的结果是一个 $m \times (m \times n)$ 的 Jacobian 矩阵。特别复杂。所幸，在深度学习中，求导的目的是为了做梯度下降，所以为 0 的导数实际上也没用。而通过$y_j$ 的定义，我们知道如果 $i \ne j$ 则 $\frac{\partial y_i}{\partial \mathbf{w}_j} = 0$。于是我们去除这些为 $0$ 的项，保留：

$$\frac{\partial \mathbf{y}}{\partial \mathbf{w}}= \begin{bmatrix}\frac{\partial y_1}{\partial \mathbf{w}_1} \\\frac{\partial y_2}{\partial \mathbf{w}_2} \\\vdots \\\frac{\partial y_m}{\partial \mathbf{w}_m}\end{bmatrix}= \begin{bmatrix}\mathbf{x}^T \\\mathbf{x}^T \\\vdots \\\mathbf{x}^T\end{bmatrix}$$

如果我们把 $\mathbf{w}$ 写成矩阵形式 $\mathbf{W} = [\mathbf{w}_1 \cdots,\mathbf{w}_m]$，此时 $\mathbf{y} = \mathbf{W} \cdot \mathbf{x} = \mathbf{W}^T \mathbf{x}$，则上面的结论可以写成：

$$\begin{eqnarray}\frac{\partial \mathbf{y}}{\partial \mathbf{x}} = \mathbf{W}^T\end{eqnarray}$$

$$\begin{eqnarray}\frac{\partial \mathbf{y}}{\partial \mathbf{W}}= \begin{bmatrix}\mathbf{x}^T \\\mathbf{x}^T \\\vdots \\\mathbf{x}^T\end{bmatrix}= \begin{bmatrix}\mathbf{x} \cdots \mathbf{x} \end{bmatrix}^T\end{eqnarray}$$

矩阵与向量乘法求导

其实上一节的矩阵与向量点积的求导已经得出结论了。当 $\mathbf{y} = \mathbf{A} \mathbf{x}$ 时，有

$$\begin{align}\frac{\partial \mathbf{y}}{\partial \mathbf{x}} &= \mathbf{A} \\\frac{\partial \mathbf{y}}{\partial \mathbf{A}} &= \begin{bmatrix}\mathbf{x} \cdots \mathbf{x} \end{bmatrix} ^T\end{align}$$

这里我们再从数值的角度分析一下，已知 $y_i = \sum_{j}{a_{ik}x_j}$，则有：

$$\frac{\partial y_i}{\partial x_j}= \frac{a_{i1}x_1+\cdots+a_{mn}{x_n}}{\partial x_j} = a_{ij}$$

再次

$$\frac{\partial y_i}{\partial w_{ij}}= \frac{a_{i1}x_1+\cdots+a_{mn}{x_n}}{\partial w_{ij}} = x_j$$

写成矩阵形式就是结论部分。

矩阵与矩阵乘法求导

这部分过于复杂，且符号也没有统一，后续如果有用到再进行补充。

按元素(element-wise)操作

将按元素操作记为 $\unicode{x2D54}$，考虑 $\mathbf{y} = \mathbf{f(u)} \unicode{x2D54} \mathbf{g(v)}$，且向量 $\mathbf{u}, \mathbf{v}, \mathbf{y}$ 有相同的维度。写成如下形式：

$$\begin{bmatrix}y_1 \\ y_2 \\ \vdots \\ y_n\end{bmatrix}= \begin{bmatrix}f_1(\mathbf{u}) \unicode{x2D54} g_1(\mathbf{v}) \\f_2(\mathbf{u}) \unicode{x2D54} g_2(\mathbf{v}) \\\vdots \\f_n(\mathbf{u}) \unicode{x2D54} g_n(\mathbf{v})\end{bmatrix}$$

于是偏导则变成了 Jacobian 矩阵的形式：

$$\mathbf{J_u}= \frac{\partial \mathbf{y}}{\partial \mathbf{u}}= \begin{bmatrix}\frac{\partial y_1}{\partial \mathbf{u}} \\\frac{\partial y_2}{\partial \mathbf{u}} \\\vdots \\\frac{\partial y_n}{\partial \mathbf{u}} \\\end{bmatrix}= \begin{bmatrix}\frac{\partial}{\partial u_1} f_1(\mathbf{u}) \unicode{x2D54} g_1(\mathbf{v})& \frac{\partial}{\partial u_2} f_1(\mathbf{u}) \unicode{x2D54} g_1(\mathbf{v})& \cdots& \frac{\partial}{\partial u_n} f_1(\mathbf{u}) \unicode{x2D54} g_1(\mathbf{v})\\\frac{\partial}{\partial u_1} f_2(\mathbf{u}) \unicode{x2D54} g_2(\mathbf{v})& \frac{\partial}{\partial u_2} f_2(\mathbf{u}) \unicode{x2D54} g_2(\mathbf{v})& \cdots& \frac{\partial}{\partial u_n} f_2(\mathbf{u}) \unicode{x2D54} g_2(\mathbf{v})\\\vdots & \vdots & \ddots & \vdots\\\frac{\partial}{\partial u_1} f_n(\mathbf{u}) \unicode{x2D54} g_n(\mathbf{v})& \frac{\partial}{\partial u_2} f_n(\mathbf{u}) \unicode{x2D54} g_n(\mathbf{v})& \cdots& \frac{\partial}{\partial u_n} f_n(\mathbf{u}) \unicode{x2D54} g_n(\mathbf{v})\end{bmatrix}$$

但是注意到 $\unicode{x2D54}$ 是 element-wise 操作，于是 $y_i$ 只跟$f_i(\mathbf{u})$ 和 $g_i(\mathbf{v})$ 相关，换句话说，对于 $i \ne j$ 的情况，有 $\frac{\partial y_i}{\partial u_j} = 0$。更进一步，element-wise 操作代表着 $f_i(\mathbf{u})$ 可以退化成 $f_i(u_i)$，而跟其它所有 $u_k (k \ne i)$ 无关。

$$\mathbf{J_u}= \begin{bmatrix}\frac{\partial}{\partial u_1} (f_1(u_1) \unicode{x2D54} g_1(v_1))& 0& \cdots& 0\\0& \frac{\partial}{\partial u_2} (f_2(u_2) \unicode{x2D54} g_2(v_2))& \cdots& 0\\\vdots & \vdots & \ddots & \vdots\\0& 0& \cdots& \frac{\partial}{\partial u_n} (f_n(u_n) \unicode{x2D54} g_n(v_n))\end{bmatrix}$$

注意到只有对象元素有值，是对角矩阵，于是写成下式：

$$\frac{\partial \mathbf{y}}{\partial \mathbf{u}} = \mathbf{J_u}= diag \left(\frac{\partial}{\partial u_1} \left(f_1(u_1) \unicode{x2D54} g_1(v_1)\right),\cdots,\frac{\partial}{\partial u_n} \left(f_n(u_n) \unicode{x2D54} g_n(v_n)\right)\right)$$

再进一步，深度学习中一般 $f(u_i) = u_i$ 和 $g(v_i) = v_i$，所以还能简化：

$$\frac{\partial \mathbf{y}}{\partial \mathbf{u}} = \mathbf{J_u}= diag \left(\frac{\partial}{\partial u_1} \left(u_1 \unicode{x2D54} v_1\right),\cdots,\frac{\partial}{\partial u_n} \left(u_n \unicode{x2D54} v_n\right)\right)$$

于是常见的 element-wise 操作及其导数如下

由于一般拿导数是用来更新向量的，对角矩阵经常也直接当成向量来使用。

小结

主要回顾了 dot product、矩阵乘法与 element-wise 乘法的关系，以及这些操作求偏导的矩阵形式。

本篇只介绍算法的基础知识，实现部分请参考实现篇。

AD 能干什么？

AD 能用来求偏导值的。

例如有一个 $\mathbb{R}^2 \mapsto \mathbb{R}$ 的函数（函数有 2 个输入，1 个输出）：$f(x, y)$ ，对于 $x$、$y$ 的偏导分别计为$\frac{\partial f}{\partial x}$ 和 $\frac{\partial f}{\partial y}$。通常我们不关心偏导的解析式，只关心具体某个 $x_i$, $y_i$ 取值下偏导$\frac{\partial f}{\partial x} \vert_{x=x_i,y=y_i}$ 和$\frac{\partial f}{\partial y} \vert_{x=x_i,y=y_i}$ 的值。

另外注意在神经网络在使用“梯度下降”学习时，我们关心的是“参数 $w$”的偏导。而不是“输入 $x$”的偏导。假设有 $f(x) = ax^2 + b$ 这样的神经网络，损失函数是 $l(f(x), y)$，现在给了一个样本标签对$(x_0, y_0)$，我们要计算的是$\frac{\partial l}{\partial a}\vert_{x=x_0,y=y_0,a=a_0,b=b_0}$ 和$\frac{\partial l}{\partial b}\vert_{x=x_0,y=y_0,a=a_0,b=b_0}$。在对号入座时要牢记这点。

为什么用 AD？

求偏导有很多做法，例如 symbolic differentiation使用“符号计算” 得到准确的偏导解析式，但对于复杂的函数，偏导解析式会特别复杂，占用大量内存且计算慢，并且通常应用也不需要解析式；再比如numerical differentiation通过引入很小的位移 $h$，计算 $\frac{f(x+h) - f(h)}{h}$ 得到偏导，这种方法编码容易，但受 float 误差影响大，且计算慢（有几个输入就要算几次 $f$）。

AD 认为所有的计算最终都可以拆解成基础操作（如加减乘除，exp, log, sin,cos 等基本函数）的组合。然后通过链式法则逐步计算偏导。这样使用方只需要正常组合基础操作，就能自动计算偏导，且不受 float误差的影响，还可以复用一些中间结果来减少计算量（等价于动态规划）。

链式法则回顾

AD 的数学基础就是链式法则(chain rule)：

对于函数 $z = h(x)$，如果有子函数 $y = f(x)$，满足 $z = h(x) = g(y) = g(f(x))$，则求偏导有如下关系：

$$h’(x) = g’(f(x))f’(x)\iff\frac{\partial z}{\partial x} \bigg\vert_{x_0} = \frac{\partial z}{\partial y}\bigg\vert_{y=f(x_0)} \frac{\partial y}{\partial x} \bigg\vert_{x_0}$$

上述两种写法是一致的。另外如果涉及多个变量，例如 $z = f(x, y)$，而 $x = g(t),y = h(t)$，则有：

$$\frac{\partial z}{\partial t} = \frac{\partial z}{\partial x}\frac{\partial x}{\partial t} +\frac{\partial z}{\partial y}\frac{\partial y}{\partial t}$$

上面的式子叫 multivariable case：多变量的链式法则。也可以认为是Total Derivative全微分的链式法则。

AD 具体是怎么做的？

AD 其实就是链式法则的具体实现。它有两种模式：前向模式(Forward accumulation)和反向模式(Reverse accumulation)，我们只考虑反向模式。那么具体是怎么工作的呢？考虑下面的复杂函数^[1]

$$\begin{aligned}y &= f(x_{1},x_{2})\\&= \sin x_{1} + x_{1}x_{2}\\&= \sin v_{1} + v_{1}v_{2}\\&= v_{3}+v_{4}\\&= v_{5}\end{aligned}$$

上述公式中，我们用了一些子函数来简化整个函数，画成图如下左图：

于是为了求偏导 $\frac{\partial f}{\partial x_1}$ 与 $\frac{\partial f}{\partial x_2}$的值，我们可以先定义中间值 $\bar{v_i} = \frac{\partial f}{\partial v_i}$，根据链式法则，有

$$\bar{v_i} = \frac{\partial f}{\partial v_i} = \frac{\partial f}{\partial v_{i+1}} \frac{\partial v_{i+1}}{\partial v_i} = \bar{v_{i+1}} \frac{\partial v_{i+1}}{\partial v_i}$$

于是计算时需要先“前向”计算一次，得到 $v_1, v_2, \cdots, v_5$ 的值，之后再“后向”计算 $\bar{v_5}, \bar{v_4}, \cdots, \bar{v_1}$ 的值（参考上右图），最终得到的$\bar{v_1}, \bar{v_2}$ 就是我们要计算的结果。而需要先“前向”计算一次，是因为后向计算时会用到前向的值，例如 $\bar{v_2} = \bar{v_4} v_1$ 就需要用到前向的$v_1$。

注意图里 $\bar{v_1}$ 的计算依赖了链式法则中多变量的情况，等于它所有后继节点偏导（即图中的 $\bar{v_1^a}, \bar{v_1^b}$）的和。当计算图中存在$v_i$ 指向 $v_j$ 的箭头时，我们记 $\overline{v_{i \to j}}$ 为 $f$ 从 $v_j$ 方向对 $v_i$ 的偏导，则公式可以扩充如下：

$$\bar{v_i} = \frac{\partial f}{\partial v_i}= \sum_{j \in next(i)}{\overline{v_{i\to j}}}= \sum_{j \in next(i)}{\frac{\partial f}{\partial v_{j}} \frac{\partial v_{j}}{\partial v_i}= \sum_{j \in next(i)}{\overline{v_j} \frac{\partial v_{j}}{\partial v_i}}}$$

多输出情形

多输出的情况偏理论，跳过也影响不大。神经网络的输出，在训练时最终都会接入损失函数，得到 loss 值，一般都是一个标量，可以认为神经网络的学习总是单输出的。

在多输出的情况下，链式法则依然生效。

刚才都假设函数是 $\mathbb{R}^n \mapsto \mathbb{R}$，即 n 个输入，1 个输出。考虑 m 个输出，即 $\mathbb{R}^n \mapsto \mathbb{R}^m$ 的情况。假设输入是$x_1, x_2, \cdots, x_n$，而输出是$f_1(x_1, \cdots, x_n), f_2(x_1, \cdots, x_n), \cdots, f_m(x_1, \cdots, x_n)$。此时我们要计算的偏导就不是 n 个值了，而是一个 m×n 的矩阵^[2]，每个元素 $J_{ij} = \frac{\partial f_i}{\partial x_j}$。这个矩阵一般称为Jacobian Matrix：

$$\mathbf {J_{m\times n}} =\begin{bmatrix}{\dfrac {\partial \mathbf {f} }{\partial x_{1}}}&\cdots &{\dfrac {\partial \mathbf {f} }{\partial x_{n}}}\end{bmatrix}=\begin{bmatrix}\nabla ^{\mathrm {T} }f_{1}\\\vdots \\\nabla ^{\mathrm {T} }f_{m}\end{bmatrix}=\begin{bmatrix}{\dfrac {\partial f_{1}}{\partial x_{1}}}&\cdots &{\dfrac {\partial f_{1}}{\partial x_{n}}}\\\vdots &\ddots &\vdots \\{\dfrac {\partial f_{m}}{\partial x_{1}}}&\cdots &{\dfrac {\partial f_{m}}{\partial x_{n}}}\end{bmatrix}$$

其中 $\nabla^{\mathrm{T}}f_i$ 代表 $f_i$ 对于所有输入的偏导（行向量）的转置。

考虑函数 $g: \mathbb{R}^n \mapsto\mathbb{R}^k$，$h: \mathbb{R}^k \mapsto \mathbb{R}^m$，而函数 $f$ 是二者的组合：$f(x) = h \circ g(x) = h(g(x))$，则有

$$J = J_{h \circ g} = J_h(g(x)) \cdot J_g(x)$$

此时 $\mathbf{J}$ 中的每个元素：

$$J_{ij} = \frac{\partial f_i}{\partial x_j}= \sum_{l = 1}^{k}{\frac{\partial h_i}{\partial g_l} \frac{\partial g_l}{\partial x_j}}= \begin{bmatrix}{\dfrac {\partial h_i}{\partial g_{1}}}&\cdots &{\dfrac {\partial h_i }{\partial g_{k}}}\end{bmatrix}\begin{bmatrix}{\dfrac {\partial g_1}{\partial x_{j}}} \\ \vdots \\ {\dfrac {\partial g_k }{\partial x_{j}}}\end{bmatrix}$$

可以看到和 $J_h \cdot J_g$ 的结果是一致的。不过这些性质其实都是链式法则的内容，这里也只是扩充视野。

小结

AD 把复杂的函数看成是许多小函数的组合，再利用链式法则来计算偏导。它有不同的模式，其中“后向模式”在计算偏导时先“前向”计算得到一些中间结果，之后再“反向”计算偏导。从工程的视角看，由于中间的偏导可以重复利用，能减少许多计算量。深度学习的反向传播算法（BP）是 AD 的一种特例。

所以回过头来，什么是 AD？AD 就是利用链式法则算偏导的一种实现。

参考

• A Review of automatic differentiation and its efficient implementation 一篇综述，对 AD “是什么”、“为什么”的描述比较清晰
• What is Automatic Differentiation? Youtube 视频，回过头来看它介绍了 AD 的各个方面，但第一次直接看还是比较懵的，视频也有对应的综述论文，也是比较好的补充材料
• Lecture 4 - Automatic Differentiation 一个 DL 的课程，前面的内容和其它材料差不多，最后通过扩展计算图来计算 AD 的方式对理解一些框架的具体实现很有帮助

环境信息：

• Spark 3.1.2, 运行采用 Spark on k8s operator
• 数据存储在华为 FusionInsight 6.5.1, 并开启了联邦

简单排查都没问题

Spark driver 看到的报错如下(省略了一些不重要的):

Caused by: java.io.IOException: Failed on local exception: java.io.IOException: org.apache.hadoop.security.AccessControlException: Client cannot authenticate via:[TOKEN, KERBEROS]; Host Details : local host is: "spark-sql-xcjgz-43aabd8581a44cae-exec-4/10.244.55.49"; destination host is: "xxx-hdp02":25019; 
at org.apache.hadoop.net.NetUtils.wrapException(NetUtils.java:776)
        ...
Caused by: java.io.IOException: org.apache.hadoop.security.AccessControlException: Client cannot authenticate via:[TOKEN, KERBEROS]
at org.apache.hadoop.ipc.Client$Connection$1.run(Client.java:688)
at java.security.AccessController.doPrivileged(Native Method)
at javax.security.auth.Subject.doAs(Subject.java:422)
at org.apache.hadoop.security.UserGroupInformation.doAs(UserGroupInformation.java:1746)
at org.apache.hadoop.ipc.Client$Connection.handleSaslConnectionFailure(Client.java:651)
at org.apache.hadoop.ipc.Client$Connection.setupIOstreams(Client.java:738)
at org.apache.hadoop.ipc.Client$Connection.access$2900(Client.java:376)
at org.apache.hadoop.ipc.Client.getConnection(Client.java:1529)
at org.apache.hadoop.ipc.Client.call(Client.java:1452)
... 39 more
Caused by: org.apache.hadoop.security.AccessControlException: Client cannot authenticate via:[TOKEN, KERBEROS]
at org.apache.hadoop.security.SaslRpcClient.selectSaslClient(SaslRpcClient.java:172)
at org.apache.hadoop.security.SaslRpcClient.saslConnect(SaslRpcClient.java:396)
at org.apache.hadoop.ipc.Client$Connection.setupSaslConnection(Client.java:561)
at org.apache.hadoop.ipc.Client$Connection.access$1900(Client.java:376)
at org.apache.hadoop.ipc.Client$Connection$2.run(Client.java:730)
at org.apache.hadoop.ipc.Client$Connection$2.run(Client.java:726)
at java.security.AccessController.doPrivileged(Native Method)
at javax.security.auth.Subject.doAs(Subject.java:422)
at org.apache.hadoop.security.UserGroupInformation.doAs(UserGroupInformation.java:1746)
at org.apache.hadoop.ipc.Client$Connection.setupIOstreams(Client.java:726)
... 42 more

由于环境之前都是正常的，大概率没有人动过环境。基于之前的经验，先排查一些常见的问题：

• 确定 Spark Operator 需要的（k get mutatingwebhookconfiguration）存在。之前遇到过未知原因导致 webhook 消失，driver 没有 mount 上认证信息导致鉴权失败
• 确认 namespace 加了 label，可以被 webhook 监测到。之前遇到因为 namespace 没有相关 label，webhook 失效 mount 缺少认证信息，鉴权失败
• 尝试重启 Spark Operator Controller，无效。之前遇到过 controller 本身问题导致mount 失效，可重启解决
• 确认 driver pod mount 进了 hadoop-conf, krb5.conf, keytab

说明问题应该不是由 Spark Operator 引起的。然后怀疑是不是 mount 的鉴权信息有过变动：

• 由于 driver 里没有 kinit，尝试把鉴权相关信息(core-site.xml,hdfs-site.xml, krb5.conf, keytab) 放到另一台机器上，kinit 能成功。
• 在 kinit 的基础上，尝试 hdfs mkdir 和 hdfs rmdir 都能成功。

初步认定 Hadoop 集群没问题^[1]。

接着再排除代码问题，使用 driver 镜像在另一个环境（连的另一个 CDH 集群）能正常运行。

于是环境、集群、代码看起来都没有问题，那问题在哪呢？

知识不够，Debug 来凑

最开始还是怀疑环境有问题，怀疑 spark operator 依赖环境的某些东西被修改了。但实在不知道从何查起，于是考虑远程 debug。走了一些弯路，最后是这么操作的：

1. 修改提交的 spark application，加上参数spark.driver.extraJavaOptions="-agentlib:jdwp=transport=dt_socket,server=y, suspend=y,address=5005"，这样 driver 启动后就会开启 5005 端口等待 debug
2. 之后通过 kubectl port-forward --address 0.0.0.0 <driver pod> 5005: 5005 来开启宿主机到 pod 的流量转发
3. 开启 IDEA 远程 debug，连上宿主机的 5005 端口。由于报错的是 hadoop 相关的，随便开了一个项目引入 hadoop 依赖，打断点就能用了。

断点打在 UserGroupInformation.doAs 上。发现 driver 调用时的用户信息都正常。再通过添加 spark.executor.extracJavaOptions 参数来 debug executor（记得把executor数调成 1）。结果发现 executor 调用 doAs 时，使用的用户名是root（预期是 work），鉴权模式是 SIMPLE（预期是 KERBEROS）。

这妥妥的是 executor pod 创建的问题呀。于是开始排查 executor，发现 executor 的环境变量 SPARK_USER=root，同时它没有mount krb5.conf 和 keytab，难道发现了root cause?

可惜几番折腾后都不生效。最后对比运行成功的环境，发现运行成功的 exeucotr 环境变量也是一样的，也没有 mount 任何鉴权相关的信息。

无从下手，恶补知识

既然没有任何鉴权相关的信息，executor 里是怎么鉴权的？涉及到知识盲区，怎么办？

下载 Spark 3.1.1 代码，但代码太多又无从看起，于是上网搜索相关 Feature 的PR，找到下面几个信息：

• https://github.com/apache/spark/pull/21669 增加 kerberos support for k8s
• https://docs.google.com/document/d/1RBnXD9jMDjGonOdKJ2bA1lN4AAV_1RwpU_ewFuCNWKg/edit spark kerberos support 的设计文档
• https://github.com/apache/spark/pull/22911 增加了 client mode 的支持

最重要的是最后这个链接，这个 PR 的留言里有这样的信息：

In either of those cases, the driver code will handle delegation tokens: incluster mode by creating a secret and stashing them, in client mode by usingexisting mechanisms to send DTs to executors.

说明鉴权都是 driver 做的，而 executor 会从 driver 拿到 delegation token。

不过 “delegation token” 又是啥玩意？大概搜到它是 Hadoop 发的 token，目标是减少鉴权压力，一般在 Map Reduce, Spark 这些有多个 worker 要访问 HDFS 的时候使用。但这些信息并没有本质帮助。

好在这样有了 debug 的头绪。一开始尝试使用 spark 的源码进行 remote debug，发现有许多问题搞不定。于是尝试直接下载 driver 里的所有 jar 包，导入到一个空project 中，由于是 scala 直接依赖 jar 包不好单步，于是再下载对应 jar 包的sources.jar，就可以在 IDEA 里打断点单步执行了。

同时 debug 成功和失败环境里的 executor，在对比一些步骤的变量后，终于发现问题所在：executor 在执行下面代码时，失败的环境里获取到的 token 是空。

cfg.hadoopDelegationCreds.foreach { tokens =>
  SparkHadoopUtil.get.addDelegationTokens(tokens, driverConf)
}

考虑到 DT 是从 driver 获取的，看来是 driver 里存储的 delegation token 本来就是空的。另外回过头来发现 driver 的日志里有这么一句日志：

23/01/08 21:03:31 INFO DFSClient: Cannot get delegation token from work

为什么 driver token 是空？继续 debug driver 发现 driver 在获取 delegationtoken 时返回的是 null: FileSystem.collectDelegationTokens。最终缩小到最小的复现代码：

var conf = new org.apache.hadoop.conf.Configuration();
conf.addResource(new Path("./core-site.xml"));
conf.addResource(new Path("./hdfs-site.xml"));
System.setProperty("java.security.krb5.conf", "./krb5.conf");
UserGroupInformation.setConfiguration(conf);
UserGroupInformation.loginUserFromKeytab("work", "./user.keytab");
var fs = FileSystem.get(URI.create("hdfs:///"), conf);
var creds = new Credentials();
fs.addDelegationTokens("work", creds);
Assertions.assertFalse(creds.getAllTokens().isEmpty());

Client or Server?

是 hadoop client 有 BUG？还是 hadoop server 有问题？

查代码看到 delegation token 是 namenode 创建的。于是上集群的 namenode，用arthas 监控 FSNamesystem.getDelegationToken 方法。

结果……执行复现代码，发现没有输出，尽管同时监控了所有的 4 个 namenode，没有任何一个有输出。然而在正常的环境里是有输出的。难道是 Client 有 BUG 没把请求发出去？

开始用 wireshark 抓包，发现还是有请求包发出的，当然因为是 RPC，内容看不出来，但 namenode 的 arthas 就是没有输出……最后在对比正常和错误环境的请求包，突然发现错误环境连接的是 25019 端口，这又是啥端口？

在集群上通过 netstat -natp 找到了进程，进程的命令显示它是集群的 router 角色。不管是 arthas 还是它的日志（如下），都发现它才是罪魁祸首：

2023-01-09 20:09:59,312 | WARN  | IPC Server handler 43 on 25019 | trying to get DT with no secret manager running | RouterSecurityManager.java:124
2023-01-09 20:09:59,391 | WARN  | IPC Server handler 39 on 25019 | trying to get DT with no secret manager running | RouterSecurityManager.java:124

其实之前在看 hadoop 相关代码时就注意到如果 server 出错应该要有这个日志，但在namenode 日志里没有找到。现在看到这个日志，基本确定就是它的问题。最后重启router 之后发现世界和平了。

为什么 router 会出错？

Router 生成 getDelegationToken的逻辑如下所示，通过 arthas 发现是因为 dtSecretManager.isRunning 判断失败。再追代码发现 isRunning 失败的唯一可能就是调用了AbstractDelegationTokenSecretManager::stopThreads，但是 stopThreads 只有停止 router的时候才会调用，与当前的现象不相符。

public Token<DelegationTokenIdentifier> getDelegationToken(Text renewer)
    throws IOException {
  LOG.debug("Generate delegation token with renewer " + renewer);
  final String operationName = "getDelegationToken";
  boolean success = false;
  String tokenId = "";
  Token<DelegationTokenIdentifier> token;
  try {
    if (!isAllowedDelegationTokenOp()) {
      throw new IOException(
          "Delegation Token can be issued only " +
              "with kerberos or web authentication");
    }
    if (dtSecretManager == null || !dtSecretManager.isRunning()) {
      LOG.warn("trying to get DT with no secret manager running");
      return null;
    }
    ...

可惜的是 stopThreads 的调用链路并不会输出日志。只能尝试人肉看一看从最后一次重启，到出问题之间的日志，开头结尾如下：

2022-12-29 18:54:18,187 | INFO  | pool-1-thread-1 | Stopping security manager | RouterSecurityManager.java:62
2022-12-29 21:14:37,456 | WARN  | IPC Server handler 5 on 25019 | trying to get DT with no secret manager running | RouterSecurityManager.java:124

看了不久发现，在创建 secret manager，调用 startThreads 时因为 ZK 的原因有报错：

2022-12-29 21:13:06,167 | ERROR | main | Error starting threads for zkDelegationTokens  | ZKDelegationTokenSecretManagerImpl.java:48
java.io.IOException: Could not start Sequence Counter
at org.apache.hadoop.security.token.delegation.ZKDelegationTokenSecretManager.startThreads(ZKDelegationTokenSecretManager.java:324)
at org.apache.hadoop.hdfs.server.federation.router.security.token.ZKDelegationTokenSecretManagerImpl.<init>(ZKDelegationTokenSecretManagerImpl.java:46)
at sun.reflect.NativeConstructorAccessorImpl.newInstance0(Native Method)
at sun.reflect.NativeConstructorAccessorImpl.newInstance(NativeConstructorAccessorImpl.java:62)
at sun.reflect.DelegatingConstructorAccessorImpl.newInstance(DelegatingConstructorAccessorImpl.java:45)
at java.lang.reflect.Constructor.newInstance(Constructor.java:423)
at org.apache.hadoop.hdfs.server.federation.router.FederationUtil.newSecretManager(FederationUtil.java:214)
at org.apache.hadoop.hdfs.server.federation.router.security.RouterSecurityManager.<init>(RouterSecurityManager.java:54)
at org.apache.hadoop.hdfs.server.federation.router.RouterRpcServer.<init>(RouterRpcServer.java:262)
at org.apache.hadoop.hdfs.server.federation.router.Router.createRpcServer(Router.java:385)
at org.apache.hadoop.hdfs.server.federation.router.Router.serviceInit(Router.java:194)
at org.apache.hadoop.service.AbstractService.init(AbstractService.java:164)
at org.apache.hadoop.hdfs.server.federation.router.DFSRouter.main(DFSRouter.java:69)
Caused by: org.apache.zookeeper.KeeperException$ConnectionLossException: KeeperErrorCode = ConnectionLoss for /hdfs-federation/RouterZkDtSecretManager/ZKDTSMRoot/ZKDTSMSeqNumRoot
at org.apache.zookeeper.KeeperException.create(KeeperException.java:104)
at org.apache.zookeeper.KeeperException.create(KeeperException.java:51)
at org.apache.zookeeper.ZooKeeper.create(ZooKeeper.java:1480)
at org.apache.curator.framework.imps.CreateBuilderImpl$11.call(CreateBuilderImpl.java:740)
at org.apache.curator.framework.imps.CreateBuilderImpl$11.call(CreateBuilderImpl.java:723)
at org.apache.curator.RetryLoop.callWithRetry(RetryLoop.java:109)
at org.apache.curator.framework.imps.CreateBuilderImpl.pathInForeground(CreateBuilderImpl.java:720)
at org.apache.curator.framework.imps.CreateBuilderImpl.protectedPathInForeground(CreateBuilderImpl.java:484)
at org.apache.curator.framework.imps.CreateBuilderImpl.forPath(CreateBuilderImpl.java:474)
at org.apache.curator.framework.imps.CreateBuilderImpl$4.forPath(CreateBuilderImpl.java:349)
at org.apache.curator.framework.imps.CreateBuilderImpl$4.forPath(CreateBuilderImpl.java:291)
at org.apache.curator.framework.recipes.shared.SharedValue.start(SharedValue.java:229)
at org.apache.curator.framework.recipes.shared.SharedCount.start(SharedCount.java:155)
at org.apache.hadoop.security.token.delegation.ZKDelegationTokenSecretManager.startThreads(ZKDelegationTokenSecretManager.java:321)
... 12 more
2022-12-29 21:13:06,168 | INFO  | main | Secret manager instance created | FederationUtil.java:215

但如果创建的时候就失败了，按代码逻辑，启动的时候应该失败：

public RouterSecurityManager(Configuration conf) throws IOException {
  AuthenticationMethod authMethodConfigured =
      SecurityUtil.getAuthenticationMethod(conf);
  AuthenticationMethod authMethodToInit =
      AuthenticationMethod.KERBEROS;
  if (authMethodConfigured.equals(authMethodToInit)) {
    this.dtSecretManager = FederationUtil.newSecretManager(conf);
    if (this.dtSecretManager == null || !this.dtSecretManager.isRunning()) {
      throw new IOException("Failed to create SecretManager");
    }
  }
}

最后发现看的代码和集群的版本不一致，FI 6.5.1 是基于 hadoop 3.1.1 版本，但是3.1.1 版本代码里并没有 RouterSecurityManager，于是拉到集群里的 jar 包反编译，发现，FI 实现的 RouterSecurityManager 没有任何的校验：

而开源的 hadoop 在创建完后是有校验的，也找到了相关的修改commit 。

小结

最终的问题链路是：

1. FI 集群 Router 重启，启动过程需要创建 secret manager，具体的实现是基于 ZK的
2. 由于某些原因在启动 ZKDelegationTokenSecretManagerImpl 连接 ZK 失败，导致相关的服务启动失败
3. 但是 RouterSecurityManager 并没有对 secret manager 的启动状态做校验（后续版本修复），仍然继续运行
4. 在实际创建 delegation token 时，由于 secret manager 的状态异常，创建的请求也失败
5. Spark Driver 在获取 delegation token 时得到一个空值
6. Spark Executor 需要访问 HDFS 时，会从 Driver 的 Resource Manager 中获取profile，其中包含 driver 获取的 token，由于 token 是空的，于是鉴权失败。

另外几点感想：

• 所幸问题是发生在测试环境，要发生在客户环境，估计一万年都查不出来（当然任务肯定也会失败）。
• 另外虽然看代码很重要，但还是应该补齐相关模块的理论知识，这样事半功倍。
• 重启大法好

参考

• https://blog.cloudera.com/hadoop-delegation-tokens-explained/ Hadoopdelegation token 介绍，比较偏使用而非实现原理
• https://docs.google.com/document/d/1RBnXD9jMDjGonOdKJ2bA1lN4AAV_1RwpU_ewFuCNWKg/edit spark on k8s kerberos 支持的设计文档，跟实现还是有点差距
• https://medium.com/agile-lab-engineering/spark-remote-debugging-371a1a8c44a8spark driver/executor 远程 debug 的方法

1. K8S 中使用 ServiceAccount 时，内部本质上是用 JWT 做校验
2. JWT 中的 nbf 字段代表 token 的“开始时间”。开始时间不得早于“机器当前时间”，实际允许有 1min 偏差
3. 如果集群节点的时钟偏差(clock skew)超过 1min，可能出现 A 节点签发的 token 开始时间过早，导致 token 在 B节点校验失败

排查过程

显示 SA 没权限，但 SA 配置都正确

在 k8s 上启动的任务，会通过 fabric8.io java client 创建 SparkApplication 的Custom Resource(CR)。然而某一天开始，测试环境提交的任务全都失败，报下面的错误：

Exception in thread "main" io.fabric8.kubernetes.client.KubernetesClientException: Failure executing: GET at: https://10.233.0.1/api/v1/namespaces/.../pods/xxx-pod. Message: Unauthorized! Configured service account doesn't have access. Service account may have been revoked. Unauthorized.
        at io.fabric8.kubernetes.client.dsl.base.OperationSupport.requestFailure(OperationSupport.java:682)
        at io.fabric8.kubernetes.client.dsl.base.OperationSupport.requestFailure(OperationSupport.java:661)
        at io.fabric8.kubernetes.client.dsl.base.OperationSupport.assertResponseCode(OperationSupport.java:610)
        at io.fabric8.kubernetes.client.dsl.base.OperationSupport.handleResponse(OperationSupport.java:555)
        at io.fabric8.kubernetes.client.dsl.base.OperationSupport.handleResponse(OperationSupport.java:518)
        at io.fabric8.kubernetes.client.dsl.base.OperationSupport.handleGet(OperationSupport.java:487)
        at io.fabric8.kubernetes.client.dsl.base.OperationSupport.handleGet(OperationSupport.java:457)
        at io.fabric8.kubernetes.client.dsl.base.BaseOperation.handleGet(BaseOperation.java:698)
        at io.fabric8.kubernetes.client.dsl.base.BaseOperation.getMandatory(BaseOperation.java:184)
        at io.fabric8.kubernetes.client.dsl.base.BaseOperation.get(BaseOperation.java:151)
        at io.fabric8.kubernetes.client.dsl.base.BaseOperation.get(BaseOperation.java:83)
        ...

由于近期刚做过部署操作，开始怀疑是不是 SA 配置错了。于是人肉检查 SA

• get pod 检查 serviceAccount 和 serviceAccountName 的值都是符合预期的。
• get clusterrole 和 get rolebinding -n <namespace> 检查都是正确的

发现跟节点有关系

SA 检查无误，于是想先抓包看看是不是网络相关的问题，在 get pods -o wide 时发现所有任务都调度到 node2 这个节点。于是先把 node2 下掉，直接搜了个命令：

kubectl taint nodes node2 key1=value1:NoSchedule

新起的任务调度到 node1 后发现任务都 OK。因为暂时还在搞其它事情，把这个问题汇报给 SRE 同事，就暂停了。

sleep 40s 后能跑过

SRE 同事做了一些尝试，发现有即使在 node2 提交，偶尔也是能通过的。期间有两个怀疑：

1. SA 是不是过期了，但搜了搜发现一般时间还是挺长的，应该不是这个问题
2. 由于任务是由 argo workflow 提交的，开始怀疑是不是 argo 的问题（命令是argoexec 运行的）

另外 SRE 同事试着在命令执行前增加 sleep 40s 发现就能提交通过。

修改系统时间确认相关，但解释不通

在查资料时有提到是不是有同步问题，于是灵光一闪会不会跟系统时间有关，一查，两台节点的时间差是大约是 1min30s。于是把时间拔到 1min 内，发现提交的任务正常了。于是确定是和系统时间相关。但是具体的机制搞不清楚。

sudo date $(date +%m%d%H%M%Y.%S -d '-1 minutes')

更多测试

期间已经把测试的内容把成单纯的 curl:

       find / -name "*.crt"
       TOKEN=$(cat /var/run/secrets/kubernetes.io/serviceaccount/token)
1. 401 curl -vik --header "Authorization: Bearer ${TOKEN}" -k https://node1:6443/api/v1/namespaces/.../pods/..
2. 404 curl -vik --header "Authorization: Bearer ${TOKEN}" -k https://node2:6443/api/v1/namespaces/.../pods/..
       sleep 30
3. 404 curl -vik --header "Authorization: Bearer ${TOKEN}" -k https://node1:6443/api/v1/namespaces/.../pods/..
4. 404 curl -vik --header "Authorization: Bearer ${TOKEN}" -k https://node2:6443/api/v1/namespaces/.../pods/..
       find / -name "*.crt"

这里有一个失误，这个 curl 用的 API 即使成功也是 404, 导致在测试的过程中会有误判，实际上测试结果 #2 几乎都是 404, 但有时候会看成是 401. 也尝试人工进入 pod执行 curl，都是通的，想不通开始的 30s 究竟触发了什么机制导致认证失败。

401 的 curl 如下所示：

* ALPN, offering h2
* ALPN, offering http/1.1
* successfully set certificate verify locations:
*  CAfile: /etc/ssl/certs/ca-certificates.crt
*  CApath: /etc/ssl/certs
} [5 bytes data]
* TLSv1.3 (OUT), TLS handshake, Client hello (1):
} [512 bytes data]
* TLSv1.3 (IN), TLS handshake, Server hello (2):
{ [122 bytes data]
* TLSv1.3 (IN), TLS handshake, Encrypted Extensions (8):
{ [15 bytes data]
* TLSv1.3 (IN), TLS handshake, Request CERT (13):
{ [105 bytes data]
* TLSv1.3 (IN), TLS handshake, Certificate (11):
{ [1002 bytes data]
* TLSv1.3 (IN), TLS handshake, CERT verify (15):
{ [264 bytes data]
* TLSv1.3 (IN), TLS handshake, Finished (20):
{ [52 bytes data]
* TLSv1.3 (OUT), TLS change cipher, Change cipher spec (1):
} [1 bytes data]
* TLSv1.3 (OUT), TLS handshake, Certificate (11):
} [8 bytes data]
* TLSv1.3 (OUT), TLS handshake, Finished (20):
} [52 bytes data]
* SSL connection using TLSv1.3 / TLS_AES_256_GCM_SHA384
* ALPN, server accepted to use h2
* Server certificate:
*  subject: CN=kube-apiserver
*  start date: Nov  7 08:38:44 2022 GMT
*  expire date: Nov  7 08:38:45 2023 GMT
*  issuer: CN=kubernetes
*  SSL certificate verify result: unable to get local issuer certificate (20), continuing anyway.
* Using HTTP2, server supports multi-use
* Connection state changed (HTTP/2 confirmed)
* Copying HTTP/2 data in stream buffer to connection buffer after upgrade: len=0
} [5 bytes data]
* Using Stream ID: 1 (easy handle 0x55975bd06560)
} [5 bytes data]
> GET /api/v1/namespaces/argo-run/pods/4622-4622-import-529855050 HTTP/2
> Host: 172.27.128.212:6443
> user-agent: curl/7.74.0
> accept: */*
> authorization: Bearer eyJhbGciOiJSUzI1NiIsImtpZCI6ImhGdk15S3F6REdtUkdXMUprelhzRTF0RHJuT2kwdlhrQWktdnphclJSSG8ifQ.eyJhdWQiOlsiaHR0cHM6Ly9rdWJlcm5ldGVzLmRlZmF1bHQuc3ZjLmt1YmVybmV0ZXMiXSwiZXhwIjoxNzAxMDcyODQxLCJpYXQiOjE2Njk1MzY4NDEsImlzcyI6Imh0dHBzOi8va3ViZXJuZXRlcy5kZWZhdWx0LnN2Yy5rdWJlcm5ldGVzIiwia3ViZXJuZXRlcy5pbyI6eyJuYW1lc3BhY2UiOiJhcmdvLXJ1biIsInBvZCI6eyJuYW1lIjoianotMjAtNTY0My01NjQzLWltcG9ydC0zOTY2MDk2NDU0IiwidWlkIjoiOTM1NzY3M2UtNWNiMi00YjVkLTk1NDAtYTM4NDY1YTE5YWQ2In0sInNlcnZpY2VhY2NvdW50Ijp7Im5hbWUiOiJsb29mYWgiLCJ1aWQiOiI0ZTM1ZDIwZi01OGI3LTQxNWItOGZhMS01YTk5MjlkM2YyZWEifSwid2FybmFmdGVyIjoxNjY5NTQwNDQ4fSwibmJmIjoxNjY5NTM2ODQxLCJzdWIiOiJzeXN0ZW06c2VydmljZWFjY291bnQ6YXJnby1ydW46bG9vZmFoIn0.bfxySo3e2rT3mToSEmSN5Pmi4YI4X2kE4aXA_BVIPyrg8DKc9pDUFEo_kvS608pm0u5b7e7wG3A48upBUjtm2uAMwEYiDqSLning7kCdycXT1-_aXVQjeASio4dZL6w3ddi_JGyFoZA76e9cQVfaWB9PGenKlg2uJXe5xFNJA12EuCvXgTLC7rXrNZIPksI0ZR6bRBt2ENWf_aaYPLTE7H7g8TJlYfP__H5DBaBr6sRkO15q8mCKpEyIqCx-t9mf6pCWfJ3D2KOBMc01n8g55EUvlaPDFngn5eV3izfMxuJADB4QqrVt_-mIgpPbJr3j3H5wYHmzcCSvTVg_Cp32Zg
>
{ [5 bytes data]
* TLSv1.3 (IN), TLS handshake, Newsession Ticket (4):
{ [146 bytes data]
* Connection state changed (MAX_CONCURRENT_STREAMS == 250)!
} [5 bytes data]
< HTTP/2 401
< audit-id: 1bdd5211-54ea-4b3e-ac14-d7716730166a
< cache-control: no-cache, private
< content-type: application/json
< content-length: 165
< date: Sun, 27 Nov 2022 08:12:15 GMT
<
{ [5 bytes data]
100   165  100   165    0     0  11785      0 --:--:-- --:--:-- --:--:-- 12692
* Connection #0 to host 172.27.128.212 left intact
HTTP/2 401
audit-id: 1bdd5211-54ea-4b3e-ac14-d7716730166a
cache-control: no-cache, private
content-type: application/json
content-length: 165
date: Sun, 27 Nov 2022 08:12:15 GMT

{
  "kind": "Status",
  "apiVersion": "v1",
  "metadata": {

  },
  "status": "Failure",
  "message": "Unauthorized",
  "reason": "Unauthorized",
  "code": 401
}

接下来走了个弯路，因为看到日志时的 HTTP/2 401，扫了一眼看到是中间 Debug 日志输出，就以为是 TLS 握手的过程中出的错。中途 Debug 了很久 TLS 相关的内容。后来看 k8s apiserver 的日志才恍然大悟这个 401 是 apiserver 给出来的。

E1123 11:53:33.385964       1 claims.go:126] unexpected validation error: *errors.errorString
E1123 11:53:33.386042       1 authentication.go:63] "Unable to authenticate the request" err="[invalid bearer token, Token could not be validated.]"

JWT 有 1min 差异，但时间是在哪定义的？

找到是 K8S 的问题，就去找 k8s 的日志，找了很长时间找到了

• 错误位置 显示它是在校验 JWT 的 public claims 里的时间字段，和当前字段是否一致
• 在校验时 默认会有 1min 的余地

并且看代码它会对比 JWT 里的 NotBeforeTime 字段。于是通过 get secrets 拿到token，并在 jwt.io 里解析，奇怪的是并没有看到时间相关的字段

k -n <ns> get secret <SA-secret-name> -o jsonpath='{.data.token}' | base64 --decode

payload 如下

{
  "iss": "kubernetes/serviceaccount",
  "kubernetes.io/serviceaccount/namespace": "...",
  "kubernetes.io/serviceaccount/secret.name": "...",
  "kubernetes.io/serviceaccount/service-account.name": "...",
  "kubernetes.io/serviceaccount/service-account.uid": "4e35d20f-58b7-415b-8fa1-5a9929d3f2ea",
  "sub": "system:serviceaccount:argo-run:loofah"
}

这里其实又有个失误，其实很早之前就已经在 pod 里打印出 pod 里读到的 token，但一直以为 pod 里拿到的 token 和 get secrets 的结果是一样的。对比了半天才发现它们不一样，终于找到时间字段：

{
  "aud": [
    "https://kubernetes.default.svc.kubernetes"
  ],
  "exp": 1701074189,
  "iat": 1669538189,
  "iss": "https://kubernetes.default.svc.kubernetes",
  "kubernetes.io": {
    "namespace": "argo-run",
    "pod": {
      "name": "...",
      "uid": "d7f59189-7050-4922-804f-075e5411b950"
    },
    "serviceaccount": {
      "name": "...",
      "uid": "4e35d20f-58b7-415b-8fa1-5a9929d3f2ea"
    },
    "warnafter": 1669541796
  },
  "nbf": 1669538189,
  "sub": "system:serviceaccount:..."
}

通过查 JWT 的说明知道 nbf 字段就是 NotBefore 的时间。对比实际的值也发现和node2 运行任务的时间非常接近。终于破案。

情况复盘

1. B 节点的时间比 A 节点快 1min30s
2. 任务被调度到 B 节点，B 节点的 kubelet 为 Pod 生成 SA token，token 的 nbf时间为 B 节点的当前时间。（这里应该是创建 token 的请求会发往 B 的 apiserver，目前没找到方法验证）
3. B 节点里需要访问 apiserver，会访问 kubernetes.default，请求被路由到节点 A
4. A 节点在校验 JWT 时发现 token 的 nbf 在 A 节点当前时间+ 1min 之后，拒绝请求

小结

这个问题从断断续续排查了近一周，中间还是有不少失误

1. 构造测试用例，结果的判断最好清晰明确，这次排查依赖看结果是 401 还是 404, 看错了好几次，影响判断
2. 有条件的话，一些现象要相互印证。中途跑去怀疑 TLS 握手浪费了不少时间
3. 数据和信息要贴源，因为没有识别出 get secrets 和 pod 里 token 的区别，又浪费了半天时间
4. 底层机制是会咬人的，从方案和运维的视角，要会机制上防止出现相关问题，太难查了

常在想，自己工作里遇到的代码，为什么有那么多屎山？不管是加入前已经存在的，还是加入后新写的；不管是别人写的还是自己写的，仿佛不是屎山就是在变成屎山。今天主要在长短期决策上吐吐糟。

生存永远大于发展。这个功能如果这样做，未来修改的代价比较大。没事，我们先做成这样，如果没有成果，说不准明年就没有我们了。真的活到第二年了，似乎当时的困难就不存在了，要求继续全速前进。

未来的未来再说。“这期先这样设计，下个迭代我们再优化”，第二个迭代到来，“这个迭代这些需求优先级比较高，优化放下一个迭代吧”，子子孙孙无穷尽也。

先看看怎么跑通。这期能跑通就不错了，哪顾得上代码优雅不优雅，反正后面屎山维护不了，大不了跑路呗？我又不和公司和团队共存亡。

贪心算法和动态规划，我们知道通常贪心算法得不到全局最优。软件开发上，如果总是选择现在的利益，忽略未来，则注定会走向死亡。但一来不这么搞我现在就没了，二来说不准未来锅不是我背呢？再来谁知道现做的准备未来能用上呢？去 TM 的全局最优。

Iptables 是一门配置语言

它是一门配置语言。用来在网络处理的各个环节里加 Hook。常见的用途是做防火墙，做流量的转发等等。

像学习其它语言一样，语言本身有语法，语法之外还需要学习库函数。iptables 的语法大概如下：

iptables [-t table] {-I | -A | -D | -R} chain rule_specification

iptables 里有 table 和 chain 的概念，代表机器处理网络包的各个阶段，因此在指定配置时需要先指定配置在哪个阶段生效。之后是配置处理的规则，规则语法如下：

rule-specification = [matches...] [target]

match = -m matchname [per-match-options]

target = -j targetname [per-target-options]

一个规则可以有多个 match 匹配条件，以及一个 target 作为目标。它表明当一个网络包命中这些规则时，执行 target 目标。另外，iptables 是可（由其它模块）扩展的，扩展会提供新的 match 和新的 target。我们先看一个典型示例：

iptables -t nat -A PREROUTING -p tcp -j REDIRECT --to-ports 7892

这个规则的作用是将所有的 tcp 流量，全部转发到 7892 端口。这里的 -p tcp条件选中 tcp 流量是 iptables 默认支持的，但 REDIRECT 转发操作是扩展提供的。

Table and Chain

要学习语言，要先了解语言背后的执行模型（类比栈、指针等），iptables 的作用是在各个环节里增加 hook，那有哪些 hook 可以用呢？先看下图^[1]：

--->PRE------>[ROUTE]--->FWD---------->POST------>
                 |                ^
                 |                |
                 |             [ROUTE]
                 v                |
                 IN              OUT
                 |                ^
                 v                |

一个包从左侧进入系统，先到 PRE 环节。接着进入 [ROUTE] 阶段做路由，来决定包的去向。如果本机是目标地址则接收，否则尝试转发，亦或者丢弃。

对于本机接收的包，触发 IN 环节后交给对应的应用程序；转发的包在触发 FWD环节后尝试向外发包。外出的包最后还会经过 POST 环节，做最后的处理后发往网卡。

本机应用程序发出的包，会先经过 OUT 环节处理，之后经过 [ROUTE] 决定去向^[2]，最终再经过 POST 环节后发出。

在这些 hook 的基础上，iptables 用 “table” 的概念来组织常见的包修改需求。例如：

• Filter: 来做包过度
• Nat: 做地址转换
• Mangle: 其它的通用的包修改
• Raw: 处理一些 connection track 生效之前的修改

# modified from https://www.netfilter.org/documentation/HOWTO/netfilter-hacking-HOWTO.txt

--->PRE------>[ROUTE]--->FWD---------->POST------>
    Raw          |       Mangle   ^    Mangle
    ConnTrack    |       Filter   |    NAT (Src)
    Mangle       |                |
    NAT (Dst)    |             [ROUTE]
                 v                |
                 IN Mangle       OUT Filter
                 |  NAT           ^  NAT (Dst)
                 |  Filter        |  Mangle
                 |                |  ConnTrack
                 v                |  Raw

具体使用时，先决定要做的修改是什么内容，决定 table 名，然后找到 hook 的时机，决定 chain 的名字。当然 iptables 允许用户增加自己的 chain，但用户增加的 chain并不能决定 hook 的时机。

例如下面的例子里，我们要把所有流量转发到 7892 端口，我们通过 man iptables-extensions 查到，它只能加到 nat 表的 PREROUTING 或 OUTPUT链，由于我们要转发入口流量，所以修改的是 PREROUTING chain。

iptables -t nat -A PREROUTING -p tcp -j REDIRECT --to-ports 7892

REDIRECT 的限制也很容易理解，转发需要支持源、目标地址的改写，因此属于 nat表的范畴，而它需要在路由之前做修改（否则改了也发不出去），所以只能在PREROUTING 和 OUTPUT hook 里处理。

Rule 执行顺序

上面我们提到 iptables 是通过 table, chain 来组织切入点的，一个 chain 上可以配置多条规则，用户还可以自己创建 chain 来管理规则。那么 iptables 在是如何使用这些规则的呢？

正常情况下规则会一条条向下匹配，iptables 有一些特殊的 target 也提供了一些特殊的操作来在规则中跳转的能力（可以类比编程语言中的 continue, break），如下图：

• JUMP(-j <chain>)：跳转到自定义的 chain 里
• ACCEPT：流量通过当前 table + chain，不再匹配任何规则
• RETURN：从当前 chain 跳出，回到上一个 chain 跳转的位置
• DROP：丢弃流量，不再匹配任何 table 任何 chain

此外也得注意一些扩展 target 的语义，如 REDIRECT 相当于 ACCEPT；如 REJECT相当于 DROP，会在发送终止包后丢弃数据包。实操如果发现有问题，要注意是不是规则顺序引起的。

常用的扩展

上面我们了解了 iptables 的语法和执行顺序，接下来要学习“库函数”，表面上学习库函数就是学习“扩展”提供了哪些 match 和 target，但真正的难点是学习它们背后的网络处理机制。这里我们简单提几个。

fwmark

Firewall Mark(fwmark) 可以理解成一个 iptables 的扩展，它提供了 MARK 和CONNMARK 的 target，允许我们把一个数据包或一个连接打上标记。之后在其它地方可以使用这个标记。

典型的使用方式是让有某个标记的流量走某个特殊的路由表^[3]，例如：

ip rule add fwmark 1 table 100
ip route add local default dev lo table 100

iptables -t mangle -A OUTPUT -p udp -d 198.18.0.0/16 -j MARK --set-mark 1

其中的 ip rule add fwmark 1 table 100 是创建了一张名为 100 的路由表，并指定当 fwmark 为 1 时才查这张表。而下面的规则指定了 -p udp 匹配 UDP 流量，且目标地址为 -d 198.18.0.0/16 时执行 -j MARK 操作，把数据包打上 --set-mark 1 这个标记。

成果是目标地址为 198.18.0.0/16 的 UDP 流量会查 100 路由表。

NAT: SNAT, DNAT, MASQUERADE

Network Address Translation 的变种比较多，但思路还是容易理解的。在网络隔离的情况下，如果想两个网段里交换网络包，则需要在路由器（能同时访问两个网段）里对包做地址转换，如下所示：

SNAT 是换了源 IP 字段，所以一般用于出口流量；DNAT 换了目标 IP 字段，所以一般用于做“端口映射”来穿透内网。可以看到不论是 SNAT 还是 DNAT 都需要提供目标的 IP 地址。而 MASQUERADE 可以理解成 SNAT 的变种，它可以自动填写对应网卡的 IP，不需要手工指定了，一般用于路由器流量内外网转发。

另外从图里看到，无论是 SNAT 还是 DNAT，都需要维护一张 NAT 映射表，可以通过conntrack -L 看到。如果在路由器的 SNAT 里，--to-source IP 不是本机会怎么样呢？连接会建立失败，路由还是正常记录了 NAT 映射表，但 ACK 包会直接发到--to-source IP 上，被丢弃。

额外的，TCP 流量只有在连接建立时会查 iptables NAT 表，同个连接后续的包会沿用建立连接时的规则。

参考

• A Deep Dive into Iptables and Netfilter Architecture 介绍了 hook 和 table 的作用
• Iptables Flow 一个简化但容易理解的 iptables 流程图
• Packet flow in Netfilter and General Networking 一张复杂但全面的流程图
• man iptables-extensions 各种扩展支持的 match, target 都有说明

因果关系与物理时钟

假设你发了朋友圈，有两个朋友评论：

• A 说：“这是在北京吧”
• B 回复 A 说：“应该不是，看着像上海”

我们人肉能识别出两句话之间的因果关系：#A 是因，#B 是果，但是计算机怎么判断呢？

一种思路是给评论加上生成时间，比如 #A_10:01, #B_10:02，系统按时间对评论排序，就能判断 #B 发生成 #A 之后。这个方法逻辑上没问题，但现实中没有一种可靠的方法，能准确地同步各个机器上的时间（也称为物理时间）。于是可能出现下面的情况：

处理 B 评论的机器时钟慢了，导致 B 评论的时间戳更小，系统排序时把 #B 放在了前面，因果错乱。

Lamport 时钟

Lamport 时钟^[1]是一种逻辑上的机制，用来给各个事件打标签，保证如果事件 A 发生于 B 之前，则 A 的标签 L(A) 一定小于 B 的标签 L(B)。

具体要怎么做呢？每个机器各自维护一个计数器 t，然后：

1. 初始化时，每个机器都把 t 置为 0
2. 本机产生一个事件时，先执行 t = t+1，再用自增后的 t 来标记事件
3. 要发送一个事件时，执行 t = t+1，并发送 (t, m)，即把计数器和事件都发出去
4. 接收到一个事件 (t', m) 时，则需要更新本地的计数器 t = max(t, t') + 1，并把 m 发送到本地

于是如果使用这个算法，则上面朋友圈的例子就变成了：

可以看到事件 (4, 这是北京吧) 发生在 (6, 应该不是)之前，它们的标签 t 能反映出这一点。

Lamport 时钟的局限

为什么 Lamport 时钟能体现事件发生的“因果”关系？如果两个事件有“因果”，它们一定是有“同步”的操作，而 Lamport 时钟则是在“同步”时（第 #4 点），通过 max(t, t')同步了二者的逻辑时间。

由于 A-Before 的事件满足 t <= T，而 B-After 的事件满足 t >= T+1，所以能保证 A-before <= T < T+1 <= B-after，而 B-after 中的事件逻辑上是发生成A-before 的事件之后的，且标签 t 也满足先后关系，因此保证了因果顺序。

但是在上图中，我们虽然推出 A-before < B-after，但其它几个区域发生的事件就没法有确定的对比结论了。例如所有 B-before 中的事件，一定发生成 A-after 中的事件之前吗？（B-before < A-after），细想一下会发现并没有办法得出这个结论。明确可比的有这几个区域：

• A-before < A-after，A 机事件发生的先后决定
• B-before < B-after，B 机事件发生的先后决定
• A-before < B-after，A、B 之间的因果性决定

从另一个角度看，Lamport 时钟可以保证如果事件 a < b（a 发生在 b 之前），就可以推出它们的标签满足 L(a) < L(b)。但反过来，如果看到两个标签 L(a) < L(b)，能反推出 a < b 吗？其实是不行的，因为我们能判定的只有 A-before 和B-after 两个区域的事件，但只看 L(a) 和 L(b) 我们并不知道 a 和 b 落在哪个区域，因此无法判断 a 和 b 发生的先后。这就是 Lamport 时钟的局限性，

Vector 时钟

vector 时钟可以解决这个问题：如果两个事件落在可比较的区域，则通过对比 vector时钟产生的标记，可以得出对应事件发生的先后顺序，即通过 L(a) < L(b) 可以得出a < b 的结论。那 vector 时钟是怎么做到的？

1. 假设有 N 台机器，记为 N[1], N[2], ..N[n]
2. 每台机器需要维护一个 N 维向量作为计数器，记为 T = <t1, t2, ..., tn>
3. N[i] 本机产生一个事件时，就把本机向量里的 ti 递增，即 T[i]++
4. 机器 N[i] 发送消息 m 时，先执行 T[i]++，再发送 (T, m)
5. 机器 N[j] 收到消息 (T', m) 时，执行 T = max(T, T')，再执行 T[j]++

这些规则看起来很复杂，但实际上它和 Lamport 时钟的“同步逻辑”一样，只是每个节点都保存了其它所有节点，最后一次同步过的计数器。执行起来如下图^[2]：

Vector 时钟最后的标签有多维，如何比较呢？vector 时钟要求，如果每一维上，都有T[i] < T'[i]，则认为 T < T'；如果每一维都有 T[i] = T'[i]，则认为 T = T'；其它情况，都认为 T 和 T' 不可比。

条件 a < b 推出 T(a) < T(b) 的结论是比较简单的，与 Lamport 时钟类似，这里给个图，不多说明了：

从 T(a) < T(b) 反推 a < b 呢？其实从 T 的定义来看，可以理解成 T 代表的是当前事件及之前发生的所有事件的集合，而 T(a) < T(b) 可以等价于集合的从属关系，那么事件 a 一定包含在 T(b) 里，因此 a < b^[3]。

小结

Lamport 时钟解决的是分布式系统下的因果一致性问题，方式是在多机有交互时求计数器的 max。它的局限是无法从计数器的大小反推事件的先后顺序。

Vector 时钟基本思路和 Lamport 时钟一样，但它在每个机器上都维护了最后看到的，其它机器的计数器。

缓存行（cache line）

首先需要知道一个概念：cache line（缓存行）。缓存从内存加载数据时，并不是只加载我们请求的那部分，而是会多加载一些，例如我们想访问一个 int，只有 4 字节，但缓存会一次性加载 64 字段（不同机器不同）。缓存每次处理的这一“块”数据，就叫 cacheline。

为什么缓存要多加载数据呢？为了利用空间局部性（space locality）来提高性能^[1]。相当于是缓存做了猜想，后续地址的数据，通常就是接下来马上要访问的数据，提前加载能提高性能。

缓存一致性与 MESI 协议

现代 CPU 体系中，一般每个核都单独配备了自己的（L1）缓存，为了保证多个CPU 在读写缓存时保证整体数据的一致性，通常需要使用缓存一致性协议，MESI 就是其中一种（可以参考博主的这篇文章）。

MESI 协议可以简单理解为“踢人协议”，如果一个 CPU 写数据到缓存里，则需要“踢”掉其它缓存里的副本。

上图中，第 ⑦ 步就是“踢人”的操作。同时要注意如果 CPU 对缓存只做“读”操作，缓存也是需要同步的，如上图的第 ⑤ 步，只是它的开销更小。

伪共享（False Sharing）

缓存一致性说的是一个 cache line 在不同缓存间的同步操作。那如果一个 cache line上存了两个变量，并且两个变量分别被不同的线程写入呢？

可以看到，虽然 CPU A 和 CPU B 各自在写自己关心的变量 x 和 y，但由于它们存在于同一个 cache line，每次写入都会造成另一个 CPU 的缓存失效。造成严重的性能问题。

常见场景与解法

我们看到伪共享的发生有两个条件：

1. 两个变量在同一个缓存行里。通常是一个类/结构体的两个 field，或是同一个数组的相邻元素
2. 不同线程同时对两个地址读写^[2]。因此通常发生在高并发的场景下。

由于 #2 条件多线程处理一般是业务要求，解法通常是打破 #1 条件：加 padding，让一个cache line 里只保留一个变量^[3]。例如 int 只占4 字节，可以在后面加 15个没用的 int 变量，撑满 64 字节^[4]。而Java 专门提供了@Contended^[5] 来简化这种情形。

False Sharing 的影响有多大？

JMH 有一个测 False Sharing 的Benchmark，在我的机器上（20c、Java 11）运行结果^[6]如下：

Benchmark                          Mode  Cnt      Score     Error   Units
BenchmarkRunner.baseline          thrpt   25  10145.377 ± 240.354  ops/us
BenchmarkRunner.baseline:reader   thrpt   25   1305.421 ± 120.908  ops/us
BenchmarkRunner.baseline:writer   thrpt   25   8839.956 ± 211.739  ops/us
BenchmarkRunner.contended         thrpt   25  11329.372 ± 105.857  ops/us
BenchmarkRunner.contended:reader  thrpt   25   2845.015 ±  48.006  ops/us
BenchmarkRunner.contended:writer  thrpt   25   8484.357 ± 112.052  ops/us
BenchmarkRunner.hierarchy         thrpt   25  11373.481 ±  39.691  ops/us
BenchmarkRunner.hierarchy:reader  thrpt   25   2885.091 ±  56.386  ops/us
BenchmarkRunner.hierarchy:writer  thrpt   25   8488.389 ±  78.959  ops/us
BenchmarkRunner.padded            thrpt   25  11338.519 ±  49.043  ops/us
BenchmarkRunner.padded:reader     thrpt   25   2868.776 ±  51.762  ops/us
BenchmarkRunner.padded:writer     thrpt   25   8469.743 ±  78.427  ops/us
BenchmarkRunner.sparse            thrpt   25   9288.740 ±  63.073  ops/us
BenchmarkRunner.sparse:reader     thrpt   25   2582.364 ±  26.045  ops/us
BenchmarkRunner.sparse:writer     thrpt   25   6706.376 ±  77.000  ops/us

baseline、contended 及 padded 吞吐上的差别大概 10%（网上一些文章差异在 2 倍、3 倍，和我的结果出入这么大的原因还没找到）。我们再用 perf 对比 cache misses：

---------------------------- Baseline ------------------------------------
      1,012,415.83 msec task-clock                #   19.258 CPUs utilized
            39,279      context-switches          #    0.039 K/sec
             1,813      cpu-migrations            #    0.002 K/sec
            82,632      page-faults               #    0.082 K/sec
 4,855,609,024,017      cycles                    #    4.796 GHz                      (50.03%)
 4,682,761,843,004      instructions              #    0.96  insn per cycle           (62.52%)
   685,213,954,685      branches                  #  676.811 M/sec                    (62.53%)
       160,820,719      branch-misses             #    0.02% of all branches          (62.51%)
 2,551,078,768,362      L1-dcache-loads           # 2519.793 M/sec                    (62.47%)
    23,872,018,958      L1-dcache-load-misses     #    0.94% of all L1-dcache hits    (62.47%)
    12,410,125,606      LLC-loads                 #   12.258 M/sec                    (49.98%)
         2,810,037      LLC-load-misses           #    0.02% of all LL-cache hits     (50.01%)

---------------------------- Contended -----------------------------------
      1,011,118.07 msec task-clock                #   19.219 CPUs utilized
            38,773      context-switches          #    0.038 K/sec
             1,830      cpu-migrations            #    0.002 K/sec
            82,441      page-faults               #    0.082 K/sec
 4,849,385,107,175      cycles                    #    4.796 GHz                      (49.99%)
 6,794,835,895,672      instructions              #    1.40  insn per cycle           (62.48%)
 1,006,635,368,787      branches                  #  995.567 M/sec                    (62.49%)
       147,370,270      branch-misses             #    0.01% of all branches          (62.49%)
 3,585,031,069,557      L1-dcache-loads           # 3545.611 M/sec                    (62.50%)
       615,324,166      L1-dcache-load-misses     #    0.02% of all L1-dcache hits    (62.51%)
       167,043,519      LLC-loads                 #    0.165 M/sec                    (50.01%)
         2,434,845      LLC-load-misses           #    1.46% of all LL-cache hits     (50.01%)

对比其中的 L1-dcache-load-misses，可以看出，加了 @Contended 的 cache miss只有 baseline 的 3%。

小结

缓存的加载写入以 cache line 为单位，典型的大小为 64B。为了保证多 CPU 下缓存数据的一致性，需要使用一些缓存一致性协议，MESI 是其中的一个经典协议，写入缓存行时会“踢掉”其它 CPU 上的缓存。如果两个变量在同一个 cache line 中，且多线程频繁读写这两个变量，会导致多 CPU “互踢”对方的 cache line，导致性能下降。在博主的机器上 False Sharing 实测大概慢 10%，而 cache miss 大概是正常的 33 倍。

参考

• A Guide to False Sharing and @Contended里面基本把 False Sharing 涉及的知识都说清楚了，推荐阅读
• @Contended (a.k.a. JEP 142) PPT 介绍 @Contended 功能及实现
• JVM series: Contend annotation and false-sharing 其中的实验显示 padding 版本的吞吐大概是没有 padding 版本的 2 倍

The MESI protocol is an Invalidate-based cache coherence protocol, andis one of the most common protocols that support write-back caches.

发现其实只要能理解什么是 “Invalidate-based”，MESI 协议就很容易理解了。在这之前先补充些相关知识。

Write-Back Cache 写回

当一份内存的数据存储在缓存时，我们有必要保证两者是一致的。假设我们修改了缓存上的数据，这份数据要如何同步回内存呢？常见的有两种方法^[1]：

1. Write-Though（直写），每次修改都同步更新到缓存和内存中
2. Write-Back（写回），修改先更新到缓存上，缓存快失效时才更新回内存中

它们的核心区别在于更新操作是“同步”还是“异步”。显然异步的写入性能更高。

Cache-Coherence 缓存一致性

“一致性”这个词的含义深挖的话还挺深奥的，类似的内容可以参考博主的另一篇文章：什么是顺序一致性。这里举一个可能容易理解但不太准确的例子：

假设没有缓存，多个 CPU 对同一个内存地址做读写，逻辑上，我们会认为这些操作是原子的，有顺序的。假设当前内存的值是 0，CPU1 先发出写操作 W(1), CPU2 再发出读操作 R，则逻辑上我们理解 CPU2 一定要读到 1 这个值。

现在假设两个 CPU 都有自己的缓存，CPU1 先发出 W(1) 写到自己的缓存，因为使用了Write-Back 技术，还没有更新到内存，此时 CPU2 发出 R，读到的是自己的缓存（或者缓存不存在从内存加载），读到的还是 0，和我们上面说的预期不一致。

缓存一致性是指：通过在缓存之间做同步，达到仿佛系统不存在缓存时的行为。一般有如下要求：

• Write Propagation（写传播）：即写入一个缓存要让其它缓存能看到
• Transaction Serialization（事务顺序化）：即不同 CPU 对同一个地址发出读写指令，不管这些指令最终的先后顺序如何，不同 CPU 看到的顺序要一样。

这也对应我们一般说的可见性和顺序性。

Invalidate-Based 基于缓存失效

一份数据，缓存 A 有副本，缓存 B 也有副本，这时如果对 A 有修改，那 A、B 就不一致了，怎么办？Invalidate-based 的思路是，对 A 有修改，就想办法让其它副本都失效，只剩下 A 这么一个副本，不就没有“不一致”的情况了？

那其它缓存要再读数据时怎么办？简单，让剩下的那个副本把数据写回到内存，再从内存里把最新的数据捞到缓存即可。

MESI 就是用 4 个状态实现了状态机，实现了这个逻辑，我喜欢把它叫作“踢人”逻辑。

MESI 逻辑简述

MESI 的状态机包含了 4 个状态，也是名字的由来：

• (M)odified: 单副本 + 脏数据（即缓存改变，未写回内存）
• (E)xclusive: 单副本 + 干净数据
• (S)hared: 多副本 + 干净数据
• (I)nvalid: 数据未加载或缓存已失效

CPU 会有读写操作，记为 PrRd 和 PrWr，缓存接收到操作后需要与其它缓存同步并更新状态，同步的信息通过总线传递，同步信号有 5 种：BusRd, BusRdX,BusUpgr, Flush, FlushOpt，不用记具体的含义，我们只需要知道，这些信号的作用和目的，就是为了在自己接收到写入操作时，把其它缓存踢掉。

考虑缓存 A 和缓存 B 都有一个副本，都处于 Shared 状态，此时 A 接收到写入操作PrRd，则有如下变化：

1. A 会向总线发出 BusUpgr，代表自己要更新缓存上的数据
2. A 发出信号后，状态变为 Modified（单副本＋脏数据），这就需要 B 的配合了
3. B 处于 Shared 状态，在接收到总线上的 BusUpgr 信号后，主动把状态变为 Invalid
4. 于是只剩下 A 一个副本了

MESI 与内存屏障

MESI 如果简单粗暴地实现，会有两个很明显的性能问题：

1. 当尝试写入一个 Invalid 缓存行时，需要等待从其它处理器或主存中读取最新数据，有较长的延时
2. 将 cache line 置为 Invalid 状态也很慢

因此 CPU 在实现时一般会通过 Store Buffer 和 Invalidate Queue 机制来做优化。

Store buffer

在写入 Invalid 状态的缓存时，CPU 会先发出 read-invalid（这样其它 CPU 的缓存行会写入更改并变成 Invalid 的状态），然后把要写入的内容先放在 Store buffer 上，等收到其它 CPU 或内存发送过来的缓存行，做合并后才真正完成写入操作。

这会导致虽然 CPU 以为某个修改写入缓存了，但其实还在 Store buffer 里。此时如果要读数据，则需要先扫描 Store buffer，此外，其它 CPU 在数据真正写入缓存之前是看不到这次写入的。

Invalidate Queue

当收到 Invalidate 申请时（如 Shared 状态收到 BusUpgr），CPU 会将申请记录到内部的Invalidate Queue，并立马返回/响应。缓存会尽快处理这些请求，但不保证“立马完成”。此时 CPU 可能以为缓存已经失效，但真的尝试读取时，缓存还没有置为 Invalid状态，于是读到旧的数据。

内存屏障

这些优化的存在，要求我们在代码里使用内存屏障，插入 store barrier 会强制将store buffer 的数据写到缓存中，这样保证数据写到了所有的缓存里；插入 readbarrier 会保证 invalidate queue 的请求都已经被处理，这样其它 CPU 的修改都已经对当前 CPU可见。

MESI 与 MSI 的区别

不做相关工作也不用太深入。大概就是如果 CPU 要读的数据在其它 CPU 中都不存在，则对于 MSI 来说需要通过 2 个总线事务才能捞到数据，但 MESI 只需要一次。

小结

本文所有内容均来源于 MESI 的 wiki。文章的核心想是指出要理解 MESI 协议，关键在于理解它是一个“基于缓存失效”的协议，理解了这点，就能理解 MESI 的状态机为什么要这么做。

另外简单讨论了 MESI 之下为什么还需要内存屏障，以及 MESI 和同类 MSI 的区别。

博主做的是上层的应用开发，点到为止已经够用了。

协程：可暂停可恢复

正常函数调用的控制流是“单入单出”，从调用开始，正常或异常返回后结束，调用的栈帧也随之销毁。而异步编程要求在函数执行到一半时，“暂停”控制流，在未来的某个时刻再“恢复”。由于控制流尚未结束，因此调用链路上的栈帧还不能被销毁，这些信息需要以某种形式保存。可暂停可恢复的控制流，加上它所保存的信息，就可以称为“协程”。

栈帧（Stack Frame）

函数调用过程中使用的临时变量会记录到栈上，这些信息是与某个函数的某次调用绑定的，调用结束后就被废弃，这些数据就是栈帧。物理形态上，通常栈帧是“叠”在一起的，例如函数 A 中调用了函数 B，而 B 又调用了 C，则在 C 运行中，栈的状态类似下图：

|    ...     |
| Frame of C |
+------------+
| Frame of B |
+------------+
| Frame of A |
+------------+

Python 记录栈帧

Python coroutine^[1] 的处理方式是直接保存栈帧。调用的最内层通过 yield暂停控制流，中间层通过 yield from 或 await^[2] 将内层的coroutine 一路往外传，需要恢复时，再使用 send 方法恢复执行^[3]：

def inner():
    print('pause inner')
    yield
    print('resumed inner')
    return 10

def middle():
    print('pause middle')
    value = (yield from inner())
    print('resumed middle')
    return value * 2

coro = middle()   # 因为 yield from 的机制，coro 指向 inner 的状态
coro.send(None)
# pause middle
# pause inner

x.send(None)      # 可以看到是从 inner 开始恢复的
# resumed inner
# resumed middle
# ---------------------------------------------------------------------------
# StopIteration                             Traceback (most recent call last)
# <ipython-input-19-9cc02a983a52> in <module>
# ----> 1 coro.send(None)
#
# StopIteration: 20

注意在 coroutine 中，最终的返回值是通过 StopIteration 带出来的。

此外，外层拿到的 coro 其实包含了最内层 inner 的栈帧（需要了解yield from 的机制），因此第二次调用coro.send(None) 时，会从 inner 函数 yield 处恢复执行。

Rust 编译成状态机

对于缺少 GC 的语言来说，移动、复制栈帧是个原理可行，实际几乎不可行的操作。这些语言里手工创建的指针，可以指向栈上分配的内存，指针还可能被其它线程引用。栈帧移动时，这些指针都需要“修复”；栈帧复制时，数据多了份引用，内存释放又成问题。

Rust 使用了“状态机”的方式来实现控制流的暂停、恢复的能力^[4]。

首先是最内层的暂停逻辑，与 Python 不同，内层没有专门的暂停机制，只约定了接口，如果（因为资源未就绪）要暂停，则返回一个特殊值（Poll::Pending），由调用方来决定是否真的暂停和处理恢复。

pub trait Future {
    type Output;
    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output>;
}

pub enum Poll<T> {
    Ready(T),
    Pending,
}

Python 的中间层会通过 yield from 向外传递栈帧^[5]，那 Rust的中间层如何对外层提供暂停、恢复的能力呢？Rust 里提供了 await 关键词来表达等待内层的 future^[6]：

fn inner1() -> impl Future<Output = u32> {
    future::ready(1) // 返回的是 Future 的一个具体实现，这里省略
}

fn inner2() -> impl Future<Output = u32> {
    future::ready(2)
}

async fn middle() -> usize {
    let x = inner1().await; // await 代表等待内层的 future
    let y = inner2().await;
    x + y
}

那么 async/await 底层发生了什么？Rust 编译器会做这么几件事：

1. 遇到 async fn 定义时，会把 middle 方法的返回改为 Future<Output=...>
2. 将代码逻辑以 await 为拆分点，拆成状态机的 N 个状态，每个状态存储下个await 可见的变量和 future
3. 将两个 await 之间的代码，转换成状态机的转移逻辑

上面的例子编译器会编译成类似下面的这些代码^[7]：

// 状态存储
struct StartState {}
struct WaitingInner1State {
    inner1_future: impl Future<Output = usize>,
}
struct WaitingInner2State {
    x: usize,
    inner2_future: impl Future<Output = usize>,
}
struct EndState {}

// 状态机
enum StateMachine {
    Start(StartState),
    WaitingInner1(WaitingInner1State),
    WaitingInner2(WaitingInner2State),
    End(EndState),
}

// 转移逻辑
impl Future for StateMachine {
    type Output = usize; // return type of `middle`

    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        loop {
            match self { // TODO: handle pinning
                StateMachine::Start(state) => {         // 开始到第一个 await
                    inner1_future = inner1();
                    let state = WaitingInner1State {inner1_future};
                    *self = StateMachine::WaitingInner1(state);
                }
                StateMachine::WaitingInner1(state) => { // 第一个 await 到第二个 await
                    match state.inner1_future.poll(cx) => {
                        Poll::Pending => return Poll::Pending,
                        Poll::Ready(x) => {
                            inner2_future = inner1();
                            let state = WaitingInner2State {x, inner2_future};
                            *self = StateMachine::WaitingInner2(state);
                }}}
                StateMachine::WaitingInner2(state) => { // 第二个 await 到结束
                    match state.inner2_future.poll(cx) => {
                        Poll::Pending => return Poll::Pending,
                        Poll::Ready(y) => {
                            let ret = state.x + y;
                            *self = StateMachine::End(EndState);
                            return Poll::Ready(ret)
                }}}
                StateMachine::End(state) => {
                    panic!("poll called after Poll::Ready was returned");
                }
}}}}

// async def 编译成了返回 Future
fn middle() -> impl Future<Output = usize> {
    StateMachine::Start(StartState{})
}

可以看到中间层返回的 StateMachine 本身记录了内部调用的 Future 所处的状态。最外层的调用方如果需要恢复执行，只需再调用 middle 返回 future 的 poll 方法即可，middle 会根据当前状态决定去 poll 哪个内层 future。

轮询与中断/回调

异步编程的特征之一，是当资源未就绪时，先暂停当前控制流，先执行其它可推进的逻辑，等资源就绪时，再恢复之前暂停的控制流。那什么时候才知道资源就绪呢？一般有两种方法：轮询与中断。

轮询

轮询很好理解，就是外围调用方不断调用 poll 方法去查看当前资源的状态是否就绪：

future = middle();
loop {
    match future.poll() {
        Poll::Pending => {}
        Poll::Ready(ret_val) => {
            // 执行逻辑
}}}

但如果是这么做，资源未就绪前会不断执行 future.poll，浪费 CPU。此时空闲的 CPU可以用来处理其它就绪的 future，于是可以把所有需要轮询的协程添加到一个队列里，这样一个线程就可以处理 N 个协程。伪代码如下：

loop {
    future = waiting_queue.pop_front() // 队列存放所有 future
    match future.poll() {
        Poll::Pending => {
            waiting_queue.push_back(future)
        }
        Poll::Ready(ret_val) => {
            // ① 执行正常逻辑
        }}}

这会引申出一个问题：在 ① 中，如果 middle future 的结果就绪了，接下来需要执行哪部分代码呢？显然需要从 future 暂停的地方接着执行（即 outer 的后续逻辑），但我们怎么找到外层的逻辑？

一种想法是把外层逻辑也封装成一个 future^[8]，队列里直接存放outerfuture 而不是 middle future，恢复时只要执行 outer future 的 poll方法即可。这就是异步编程的传染性，只要内部有一处异步，它的每个调用方都需要是异步的，一直到顶层的 main 函数^[9]。

于是就像有多个线程一样，我们的队列里可以存放 N 个顶层的 future，可以类比成轮询N 个 main 函数。这个不断从队列中获取新的协程并调用 poll 的角色在 Rust 里叫executor^[10]。

loop {
    main_future = waiting_queue.pop_front()
    match main_future.poll() {
        Poll::Pending => {
            // ② 处理队列中的下一个
            waiting_queue.push_back(future)
        }
        Poll::Ready() => {
            // future 完成退出
        }}}

② 中的逻辑会不断把未就绪的 future 放入队列，这样每轮轮询时都会 poll 所有future，这样依旧会浪费很多资源（CPU & IO），最理想的方式是每次 poll 时只poll 那些“很有希望 ready”的 future。这就是我们下面要说的“中断”的模式，当资源就绪时，再把future 加入队列。

中断与 waker

我们希望 future 只在资源就绪时才被重新放回队列^[11]，于是executor 需要提供如下方法（伪代码）：

let mut ready_queue = Queue::new();
let mut futures: HashMap<usize, RefCell<Future<Output=()>>> = HashMap::new();
let mut num: usize = 0;

// ① 监听 ready_queue 并对其中的元素进行 poll
fn run() {
    loop {
        let _ = ready_queue.pop_front().poll();
    }
}

// ② 提供方法监听新的 future，需要将其加入 ready_queue 进行首次 poll
fn add_future(future: RefCell<Future<Output=()>>) {
    ready_queue.push_back(future.clone())
    num += 1;
    futures.insert(num)
}

// ③ 提供机制在 future 就绪时将其加入 ready_queue 中，等待下次 poll
fn wake_up(n: usize) {
    let future = futures.remove(&n).unwrap();
    ready_queue.push_back(future);
}

现在的问题是：“谁”负责在“什么时候”调用 wake_up 方法？

先来看“谁”的问题，唤醒的条件是资源就绪，那必然是资源的拥有者来唤醒，而只有“最内层”的协程才知道它等待的是什么资源，因此需要最内层的协程（通过注册回调函数）来触发。但是 wake_up 唤醒的时候得唤醒最外层的协程，即上面伪代码的参数 n，于是每次调用 poll 都需要把 n 一路下传到最内层：

fn run() {
    loop {
        let (future, index) = ready_queue.pop_front();
        future.poll(index);
    }
}

当然，伪代码里用 future 的序号 n 来唤醒外层 future 是一个实现细节。回过头来看 rust Future 接口，它包含了一个 Context 的引用，cx.waker() 可以获得“唤醒器”，再调用wake 方法即可唤醒对应的最外层的协程。与 n 一样，每次对poll 的调用，都需要把 cx 一路下传到最内层。

pub trait Future {
    type Output;
    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output>;
}

另一个问题是“什么时候”调用，显然是“资源就绪”时。那怎么知道资源什么时候就绪？这就需要资源的提供方来通知了。通常异步编程多是在处理 IO，对于 IO 一般是操作系统通过 select 或者 epoll 等等机制提供了异步通知的能力。代码里需要在等待资源时加上回调函数。整体逻辑如下图：

其中的 reactor 会监听所有在等待的资源，如果某个资源就绪了，同步的 poll 会返回就绪的资源，reactor 会调用它们的回调函数（即 wake 方法来唤醒）。Rust 里一般把 executor 和 reactor 合起来称为 Runtime。

Python 实现

前文的描述都是以 Rust 为样例，这是因为 Rust 里的角色分得相对更清楚一些。像 executor 和 reactor 的能力，在 Python 里都囊括在event loop里了，能监听什么资源，也被安排得明明白白了。

Python 里也经常用到Future，但它的概念和 Rust 里的不太一样，Python 中的 Future 本身是一个协程（实现了await方法），另外有一个 set_result 方法能设置最终结果，结果设置后，协程就能正常返回了（类似Rust里返回 Poll::Ready）。

Python 里的一个典型协程工作流如下所示：

图里包含了比较多的细节，整体逻辑和 Rust 类似，注意几点：

1. inner 注册监听事件时，Python 的做法是创建一个 future、注册事件，await future
2. 事件的注册最终都是调用 loop 的 API 来完成，也说明 Python 的 loop 包含了多个角色
3. 几乎所有的操作都是异步的，包括注册，也是通过 loop.call_soon 延迟执行的
4. future.set_result 之后，也是通过 call_soon 延迟唤醒协程
5. 唤醒后的协程，是直接从断点处恢复的（通过栈帧机制），与 Rust 不同
6. Event Loop 直接操作的是 task 而不是 coroutine，它是一个包装类，提供了取消、唤醒等功能

小结

异步编程的优势主要是节省线程数量（从而节省线程占用的栈等资源），也有说减少线程切换来节省 CPU 消耗。但总的来说，异步的最大作用和目标是提高吞吐而非降低延时。

但是，异步编程的缺点也很明显，最关键的是它的“传染性”，只要有一处要异步，所有地方都需要异步。另一个是“隔离性”，它的生态和同步的方法天然不通，一般为了支持异步，几乎所有同步的标准库都需要重写一个异步版本的。我甚至认为如果“高吞吐”不是产品的核心特性（如网关），就不应该使用异步框架。

本文尝试挖掘 Rust 和 Python 实现异步框架的模式，让我们对异步的底层实现建立一个概念，希望借助这些概念，去理解、解决编程中遇到的异步相关问题。文章主要讲解了三方面的内容：

1. 协程的核心是控制流的中断和恢复，Python 为代表的 GC 语言用的是存储栈帧的方式，而以 Rust 为代表的非 GC 语言使用了编译成状态机的方式。
2. 异步的优势想要体现，需要满足一个线程可以处理多个协程的能力。轮询的想法引导我们创建了 executor 处理协程队列的思路；中断的想法引导我们理清 reactor 的作用以及上下层需要传递的信息。
3. 最后是过程中列举了 Rust 和 Python 典型的协程工作流，可以从实现上相互印证两种具体的实现思路。但在编程的使用方来看二者的 API 又没有太大的差异。

Kubernetes gives Pods their own IP addresses and a single DNS name for a setof Pods, and can load-balance across them.

K8s Service会为每个 Pod 都设置一个它自己的 IP，并为一组 Pod 提供一个统一的 DNS 域名，还可以提供在它们间做负载均衡的能力。这篇文章会对 kube-proxy 的 iptables 模式内部的机制做一个验证。大体上涉及的内容如下：

实验配置

创建一个 Service，配置如下：

apiVersion: v1
kind: Service
metadata:
  creationTimestamp: "2022-01-23T02:32:38Z"
  name: spring-test
  namespace: default
  resourceVersion: "94418"
  uid: cdaab6bc-a518-4235-a161-a4cae6f564cf
spec:
  clusterIP: 10.1.68.7
  clusterIPs:
  - 10.1.68.7
  externalTrafficPolicy: Cluster
  internalTrafficPolicy: Cluster
  ipFamilies:
  - IPv4
  ipFamilyPolicy: SingleStack
  ports:
  - nodePort: 31080
    port: 8080
    protocol: TCP
    targetPort: 8080
  selector:
    app: spring-test
  sessionAffinity: None
  type: NodePort
status:
  loadBalancer: {}

创建后的 service 如下：

$ k get svc -o wide -A
NAMESPACE     NAME          TYPE        CLUSTER-IP     EXTERNAL-IP   PORT(S)                  AGE     SELECTOR
default       kubernetes    ClusterIP   10.1.0.1       <none>        443/TCP                  2d22h   <none>
default       sender        NodePort    10.1.177.169   <none>        8081:31081/TCP           46h     app=sender
default       spring-test   NodePort    10.1.68.7      <none>        8080:31080/TCP           2d4h    app=spring-test
kube-system   kube-dns      ClusterIP   10.1.0.10      <none>        53/UDP,53/TCP,9153/TCP   2d22h   k8s-app=kube-dns

注意其中的 spring-test 和 kube-dns 两项，后面会用到。另外 service 对应的 podIP 如下：

$ k get ep
NAME          ENDPOINTS                         AGE
kubernetes    192.168.50.48:6443                2d23h
sender        10.244.1.7:8080,10.244.2.7:8080   47h
spring-test   10.244.1.3:8080,10.244.2.3:8080   2d4h

DNS

K8s 会为 Service 创建一个 DNS域名，格式为 <svc>.<namespace>.svc.<cluster-domain>，例如我们创建的spring-test Service 则会有spring-test.default.svc.cluster.local^[1] 域名。

我们首先进入 pod，看一下 /etc/resolv.conf 文件，关于域名解析的配置：

nameserver 10.1.0.10
search default.svc.cluster.local svc.cluster.local cluster.local
options ndots:5

• 这里的 10.1.0.10 是 kube-dns service 的 cluster IP

• 文件中配置了多个 search 域，因此我们写 spring-test 或spring-test.default 或 spring-test.default.svc 都是可以解析的，另外注意解析后的 IP 也不是具体哪个 POD 的地址，而是为 Service 创建的虚拟地址ClusterIP。

  root@spring-test-77d9d6dcb5-m9mvr:/# nslookup spring-test Server: 10.1.0.10 Address: 10.1.0.10#53 Name: spring-test.default.svc.cluster.local Address: 10.1.68.7 root@spring-test-77d9d6dcb5-m9mvr:/# nslookup spring-test.default Server: 10.1.0.10 Address: 10.1.0.10#53 Name: spring-test.default.svc.cluster.local Address: 10.1.68.7 root@spring-test-77d9d6dcb5-m9mvr:/# nslookup spring-test.default.svc Server: 10.1.0.10 Address: 10.1.0.10#53 Name: spring-test.default.svc.cluster.local Address: 10.1.68.7

• ndots:5 指的是如果域名中的 . 大于等于 5 个，则不走 search 域，目的是减少常规域名的解析次数^[2]

iptables 转发

DNS 里创建的记录解决了域名到 ClusterIP 的转换问题，发送到 ClusterIP 的请求，如何转发到对应的 POD 里呢？K8s Service 有几种实现方式，这里验证的是 iptables 的实现方式：kube-proxy 会监听 etcd 中关于 k8s 的事件，并动态地对 iptables 做配置，最终由 iptables 来完成转发。先看看跟这个 Service 相关的规则如下：

0.  -A PREROUTING -j KUBE-SERVICES
1.  -A KUBE-NODEPORTS -p tcp -m tcp --dport 31080 -j KUBE-SVC-S
2.  -A KUBE-SEP-A -s 10.244.2.3/32 -j KUBE-MARK-MASQ
3.  -A KUBE-SEP-A -p tcp -m tcp -j DNAT --to-destination 10.244.2.3:8080
4.  -A KUBE-SEP-B -s 10.244.1.3/32 -j KUBE-MARK-MASQ
5.  -A KUBE-SEP-B -p tcp -m tcp -j DNAT --to-destination 10.244.1.3:8080
6.  -A KUBE-SERVICES -d 10.1.68.7/32 -p tcp -m tcp --dport 8080 -j KUBE-SVC-S
7.  -A KUBE-SVC-S ! -s 10.244.0.0/16 -d 10.1.68.7/32 -p tcp -m tcp --dport 8080 -j KUBE-MARK-MASQ
8.  -A KUBE-SVC-S -p tcp -m tcp --dport 31080 -j KUBE-MARK-MASQ
9.  -A KUBE-SVC-S -m statistic --mode random --probability 0.50000000000 -j KUBE-SEP-B
10. -A KUBE-SVC-S -j KUBE-SEP-A

我们先用 iptables-save 打印出所有的规则，筛选出和 spring-test service相关的规则，删除了一些 comment，并对名字做了简化。可以看到有这么几类：

• KUBE-NODEPORTS，这类规则用来将发送到 NodePort 的报文转到 KUBE-SVC-*
• KUBE-SERVICES：是识别目标地址为 ClusterIP(10.1.68.7)，命中的报文转到KUBE-SVC-* 做处理
• KUBE-SVC 的作用是做负载均衡，将请求分配到 KUBE-SEP 中
• KUBE-SEP 通过 DNAT 替换目标地址为 Pod IP，转发到具体的 POD 中

另外经常看到 -j KUBE-MARK-MASQ，它的作用是在请求里加上 mark，在POSTROUTING 规则中做 SNAT，这点后面再细说。

我们开启 iptables 的 trace 模式^[3]，并在其中一个 pod 发送一个请求，检查 TRACE 中规则的命中情况（由于输出特别多，这里挑选了重要的输出并做了精简）：

0:  nat:PREROUTING    IN=cni0 OUT=           SRC=10.244.1.7 DST=10.1.68.7  DPT=8080
6:  nat:KUBE-SERVICES IN=cni0 OUT=           SRC=10.244.1.7 DST=10.1.68.7  DPT=8080
10: nat:KUBE-SVC-S    IN=cni0 OUT=           SRC=10.244.1.7 DST=10.1.68.7  DPT=8080
3:  nat:KUBE-SEP-A    IN=cni0 OUT=           SRC=10.244.1.7 DST=10.1.68.7  DPT=8080
    mangle:FORWARD    IN=cni0 OUT=flannel.1  SRC=10.244.1.7 DST=10.244.2.3 DPT=8080 

• 在 PREROUTING 时，进入第 6 条进判定
• KUBE-SERVICES 判断目标地址为 10.1.68.7 且目标端口为 8080，于是跳转进入 KUBE-SVC-S 链的判断
• KUBE-SVC-S 有多条规则，从日志看最终是从第 10 条退出，进入 KUBE-SEP-A 链
• KUBE-SEP-A 最终命中第 3 条规则退出，但此时会进行 DNAT 转换目标地址
• 下一条日志显示，DST 目标地址已经变成 pod 地址 10.244.2.3 了

类似的，如果我们是通过 NodePort 来访问 Service，则 Trace 日志如下：

0: nat:PREROUTING:      IN=eth0 OUT=     SRC=192.168.50.135 DST=192.168.50.238 DPT=31080
6: nat:KUBE-SERVICES:   IN=eth0 OUT=     SRC=192.168.50.135 DST=192.168.50.238 DPT=31080
1: nat:KUBE-NODEPORTS:  IN=eth0 OUT=     SRC=192.168.50.135 DST=192.168.50.238 DPT=31080
9: nat:KUBE-SVC-S:      IN=eth0 OUT=     SRC=192.168.50.135 DST=192.168.50.238 DPT=31080
9: nat:KUBE-SVC-S:      IN=eth0 OUT=     SRC=192.168.50.135 DST=192.168.50.238 DPT=31080
5: nat:KUBE-SEP-A:      IN=eth0 OUT=     SRC=192.168.50.135 DST=192.168.50.238 DPT=31080
   mangle:FORWARD:      IN=eth0 OUT=cni0 SRC=192.168.50.135 DST=10.244.1.3     DPT=8080

iptables 负载均衡

上一节我们比较关注 iptables 转发的内容，那么如何做负载均衡？这部分是比较纯粹的iptables 知识^[4]:

首先：iptables 对于规则的解析是严格顺序的，所以如果只是单纯列出两个条目，则会永远命中第一条：

-A KUBE-SVC-S -j KUBE-SEP-A
-A KUBE-SVC-S -j KUBE-SEP-B

于是，我们需要第一条规则在某些条件下不命中。这样 iptables 就有机会执行后面的规则。iptables 提供了两种方法，第一种是有随机数，也是上一节我们看到的：

-A KUBE-SVC-S -m statistic --mode random --probability 0.50000000000 -j KUBE-SEP-B

这条规则在执行时，iptables 会随机生成一个数，并以 probability 的概率命中当前规则。换句话说，第一条命中的概率是 p，则第二条规则就是 1-p。如果有 3 个副本，则会类似下面这样的规则，大家可以计算下最后三个 Pod 是不是平均分配：

-A KUBE-SVC-S --mode random --probability 0.33333333349 -j KUBE-SEP-A
-A KUBE-SVC-S --mode random --probability 0.50000000000 -j KUBE-SEP-B
-A KUBE-SVC-S -j KUBE-SEP-C

另外一种模式是 round-robin，但是 kubernetes 的 iptables 模式不支持，这里就不细说了。猜想 kubernetes iptables 模式下不支持的原因是虽然单机 iptables 能支持round-robin，但多机模式下，无法做到全局的 round-robin。

SNAT

前面我们提到 KUBE 系列的规则经常看到 -j KUBE-MARK-MASQ，和它相关的规则有这些：

-A KUBE-MARK-MASQ -j MARK --set-xmark 0x4000/0x4000
-A KUBE-POSTROUTING -m mark ! --mark 0x4000/0x4000 -j RETURN
-A KUBE-POSTROUTING -j MARK --set-xmark 0x4000/0x0
-A KUBE-POSTROUTING -m comment --comment "kubernetes service traffic requiring SNAT" -j MASQUERADE

首先 KUBE-MARK-MASQ 的作用是把报文打上 0x4000/0x4000 的标记，在KUBE-POSTROUTING 时，如果报文中包含这个标记，会执行 -j MASQUERADE 操作，而这个操作的作用就是做源地址转换（SNAT）。那 SNAT 是什么，为什么要做 SNAT 呢？

这里引用这篇文章里的图做说明：

如果没有 SNAT，被转发到 POD 的请求返回时，会尝试把请求直接返回给 Client，我们知道一个 TCP 连接的依据是(src_ip, src_port, dst_ip, dst_port)，现在client 在等待 eIP/NP 返回的报文，等到的却是 pod IP 的返回，client 不认这个报文。换句话说，经过 proxy 的流量都正常情况下都应该原路返回才能工作。

在一些情况下可能希望关闭 SNAT，K8S 提供 externalTrafficPolicy: Local 的配置项，但流量的流转也会发生变化，这里不深入。

小结

这篇文章和上一篇Flannel 网络通信验证类似，都是尝试搭建环境，在学习 kube-proxy 工作机制的同时，对 kube-proxy 的产出iptables 做一些验证。文章中验证了这些内容：

1. 验证了 service ClusterIP 和 domain 的创建，及 pod 中 /etc/resolv.conf 中搜索域的设置
2. 验证了 kube-proxy 生成的 iptables 规则，并验证请求在这些规则中的流转
3. 学习了 iptables 负载均衡的工作机制
4. 了解了 SNAT 是什么，kube-proxy 需要做 SNAT 的原因

这篇文章的信息量不大，希望读者也撸起袖子，实打实地做一些验证，能让我们对kube-proxy 涉及的 iptables 的操作有更深刻的理解。

参考

• A Guide to the Kubernetes Networking Model 讲解了 K8S 的网络模型，有一些（动）图描述网络包的走向
• Deep Dive kube-proxy with iptables mode 深挖 kube-proxy 在 iptables 模式下的工作原理，比本文更深入
• Debug Service K8S 官方文档，讲解 Service 不工作时常见的 Debug 方法

Kubernetes 规定了网络模型，要求^[1]如下，flannel 只是其中一种实现。

1. 任意两个 pod 之间其实是可以直接通信的，无需经过显式地使用 NAT 来接收数据和地址的转换；
2. node 与 pod 之间是可以直接通信的，无需使用明显的地址转换；
3. pod 看到自己的 IP 跟别人看见它所用的 IP 是一样的，中间不能经过转换。

实验配置

使用 3 个虚拟机搭建的 Kubernetes 1.23 集群，其中 Flannel 版本为 0.16.1. 上面起了两个服务，分别为两副本。Pod 信息如下：

$ k get pods -o wide
NAME                           READY   STATUS    RESTARTS   AGE     IP           NODE       NOMINATED NODE   READINESS GATES
sender-779db554f9-d796q        1/1     Running   0          83s     10.244.2.7   centos73   <none>           <none>
sender-779db554f9-kr69b        1/1     Running   0          84s     10.244.1.7   centos72   <none>           <none>
spring-test-77d9d6dcb5-2cgs5   1/1     Running   0          5h28m   10.244.1.3   centos72   <none>           <none>
spring-test-77d9d6dcb5-m9mvr   1/1     Running   0          5h28m   10.244.2.3   centos73   <none>           <none>

实验里会尝试说明 sender-779db554f9-kr69b(10.244.1.7) 到spring-test-77d9d6dcb5-m9mvr(10.244.2.3)之间的网络通信。

Pod 与虚拟网卡

首先要说明的是 Pod 里看到的网卡，在宿主机上是如何实现的，这部分知识强烈推荐这篇文章：How Do Kubernetes and Docker Create IP Addresses?!。具体来说，是要确认下面这部分内容：

Pod 的网卡在哪？

首先，我们进入 sender-779db554f9-kr69b 所在 pod，看到网卡信息如下（省略了loopback）：

root@sender-779db554f9-kr69b:/# ifconfig
eth0: flags=4163<UP,BROADCAST,RUNNING,MULTICAST>  mtu 1450
        inet 10.244.1.7  netmask 255.255.255.0  broadcast 10.244.1.255
        ether 22:5e:27:43:63:fa  txqueuelen 0  (Ethernet)
...

注意 pod 的 IP 地址和 MAC 地址，之后我们在 centos71 机器上列出所有网卡信息：

[jinzhouz@centos72 ~]$ ip link
1: lo: <LOOPBACK,UP,LOWER_UP> mtu 65536 qdisc noqueue state UNKNOWN mode DEFAULT group default qlen 1000
    link/loopback 00:00:00:00:00:00 brd 00:00:00:00:00:00
2: eth0: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1500 qdisc pfifo_fast state UP mode DEFAULT group default qlen 1000
    link/ether 52:54:00:a0:f0:57 brd ff:ff:ff:ff:ff:ff
3: docker0: <NO-CARRIER,BROADCAST,MULTICAST,UP> mtu 1500 qdisc noqueue state DOWN mode DEFAULT group default
    link/ether 02:42:19:9e:c1:e1 brd ff:ff:ff:ff:ff:ff
4: flannel.1: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1450 qdisc noqueue state UNKNOWN mode DEFAULT group default
    link/ether f2:17:d1:67:5c:94 brd ff:ff:ff:ff:ff:ff
5: cni0: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1450 qdisc noqueue state UP mode DEFAULT group default qlen 1000
    link/ether 0a:07:55:0f:84:7f brd ff:ff:ff:ff:ff:ff
7: veth45885375@if3: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1450 qdisc noqueue master cni0 state UP mode DEFAULT group default
    link/ether a6:f6:90:57:ea:33 brd ff:ff:ff:ff:ff:ff link-netnsid 1
11: veth8360c992@if3: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1450 qdisc noqueue master cni0 state UP mode DEFAULT group default
    link/ether 26:73:c7:95:d2:0f brd ff:ff:ff:ff:ff:ff link-netnsid 2

并没有发现 Pod 里使用的这张虚拟网卡（MAC 地址没有匹配上的）。发现不了的原因是Kubernetes/Docker 等虚拟化方案，本质上是用 namespace/cgroups 对资源进行隔离，Pod 里使用的虚拟网卡，其实在另一个网络 namespace 下，那么如何确认这一点呢？参考这里 需要如下步骤：

1. 查找 pod 对应的 docker container id（这里找的是 k8s 起的 pause container）:

   $ sudo docker ps --format '{{.ID}} {{.Names}} {{.Image}}' 7f780a596b66 k8s_app_sender-779db554f9-kr69b_default_f8c7cac8-680a-45ad-a091-2b8ada73d289_0 baobao:5000/jz/sender d7226b120121 k8s_POD_sender-779db554f9-kr69b_default_f8c7cac8-680a-45ad-a091-2b8ada73d289_0 registry.aliyuncs.com/google_containers/pause:3.6 ...

2. 这里我们要找的是 k8s_POD 开头的镜像，然后查找它的 pid:

   $ sudo docker inspect --format '{{.State.Pid}}' d7226b120121 513

3. 查询 PID=513 进程对应的 veth 网卡

   [jinzhouz@centos72 ~]$ sudo nsenter -t 513 -n ip link ... 3: eth0@if11: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1450 qdisc noqueue state UP mode DEFAULT group default link/ether 22:5e:27:43:63:fa brd ff:ff:ff:ff:ff:ff link-netnsid 0

可以看到它的 MAC 地址和 POD 里看到的 MAC 地址是一样的。说明 POD 里使用的网卡就是这一张。

VET 虚拟网卡

上文提到每个 POD 的网卡是在自己的 namespace 下的，既然 namespace 是用来做网络隔离的，不同 namespace 下的网络自然是不通的。但是 k8s 又要求“node 与 pod 之间是可以直接通信”，于是我们需要打通两个 namespace，让宿主机和 POD 能直接通信。

这里使用的技术是 Virtual Ethernet(VETH)，VETH 是成对出现的，可以理解成创建了一条隧道，两端各是一张网卡，可以分别位于两个 namespace 之中，发往其中一端的包等价于发给另一端，这样就可以打通两个namespace。我们看 pod namespace 下的网卡：

[jinzhouz@centos72 ~]$ sudo nsenter -t 513 -n ip link
...
3: eth0@if11: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1450 qdisc noqueue state UP mode DEFAULT group default
    link/ether 22:5e:27:43:63:fa brd ff:ff:ff:ff:ff:ff link-netnsid 0

注意到网卡中的 @if11 字样，另一个关键信息是 link-netnsid 0，说明它关联的是ID 为 0 的 namespace 下的 ID 为 11 的网卡。我们首先确定namespace^[2]：

$ sudo ls /var/run/netns # docker 创建的 namespace 需要软链后才能查到
$ sudo ip netns list
c9e7f13179fa (id: 2)
5cc5ba76a35a (id: 1)
default

虽然没有直接展示，但 0 对应的是默认的 namespace，也就是宿主机的 namespace。再结合之前的输出：

11: veth8360c992@if3: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1450 qdisc noqueue master cni0 state UP mode DEFAULT group default
    link/ether 26:73:c7:95:d2:0f brd ff:ff:ff:ff:ff:ff link-netnsid 2

可以确认它关联的是 veth8360c992@if3 这个网卡。同理也可以反推 veth8360c992关联的是 netnsid = 2 的 id = 3 的网卡，也是符合预期的。

虚拟网卡与桥接

如果我们尝试通过 ip addr 查看 veth 网卡的 IP 地址，会发现它们是没有 IP 的：

[jinzhouz@centos72 ~]$ ip addr
5: cni0: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1450 qdisc noqueue state UP group default qlen 1000
    link/ether 0a:07:55:0f:84:7f brd ff:ff:ff:ff:ff:ff
    inet 10.244.1.1/24 brd 10.244.1.255 scope global cni0
       valid_lft forever preferred_lft forever
7: veth45885375@if3: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1450 qdisc noqueue master cni0 state UP group default
    link/ether a6:f6:90:57:ea:33 brd ff:ff:ff:ff:ff:ff link-netnsid 1
11: veth8360c992@if3: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1450 qdisc noqueue master cni0 state UP group default
    link/ether 26:73:c7:95:d2:0f brd ff:ff:ff:ff:ff:ff link-netnsid 2

这是因为对于每个 POD，宿主机上都会创建 veth 虚拟网卡，而为了更方便这些卡的管理，k8s 会创建一张桥接的网卡 cni0。可以通过下面的命令查看：

[jinzhouz@centos72 ~]$ brctl show cni0
bridge name     bridge id               STP enabled     interfaces
cni0            8000.0a07550f847f       no              veth45885375
                                                        veth8360c992

桥接(bridge)网卡可以认为是一个 2 层的交换机，当它收到一个报文时，会根据自己维护的 MAC 地址映射表将报文从不同的端口发出，如果没有找到 MAC 地址则会往所有端口都发一份。它的 MAC 映射表如下：

[jinzhouz@centos72 ~]$ brctl showmacs cni0
port no mac addr                is local?       ageing timer
  3     26:73:c7:95:d2:0f       yes                0.00
  3     26:73:c7:95:d2:0f       yes                0.00
  2     a6:f6:90:57:ea:33       yes                0.00
  2     a6:f6:90:57:ea:33       yes                0.00

对数据敏感一些会发现出现的两个 MAC 地址分别对应 veth45885375 和veth8360c992。

发送方

那么当 Pod 中向另一个宿主机上的 Pod 发请求时，会发生什么呢？整体流程如下：

1. 首先请求发到 Pod 内的 eth0 网卡，通过我们上面说的 VETH 的机制，相当于发送到cni0 网卡

2. 此时内核需要查路由表，决定发送到哪个网卡：

   [jinzhouz@centos72 ~]$ route Destination Gateway Genmask Flags Metric Ref Use Iface default RT-AC86U-D830 0.0.0.0 UG 0 0 0 eth0 10.244.0.0 10.244.0.0 255.255.255.0 UG 0 0 0 flannel.1 10.244.2.0 10.244.2.0 255.255.255.0 UG 0 0 0 flannel.1

   我们发现目标地址 10.244.2.3 命中 10.244.2.0 网段，于是发往 flannel.1 网卡

3. 接下去需要由 flannel.1 将报文通过 eth0 端口发到 centos73 机器上，这里涉及 vxlan 的工作机制，下面详细说。