教你在 Pisa-Proxy 中,如何利用 Rust 实现 MySQL

一、前言

背景

在 Database Mesh 中,Pisanix 是一套以数据库为中心的治理框架,为用户提供了诸多治理能力,例如:数据库流量治理,SQL 防火墙,负载均衡和审计等。在 Pisanix 中,Pisa-Proxy 是作为整个 Database Mesh 实现中数据平面的核心组件。Pisa-Proxy 服务本身需要具备 MySQL 协议感知,理解 SQL 语句,能对后端代理的数据库做一些特定的策略,SQL 并发控制和断路等功能。在这诸多特性当中,能够理解 MySQL 协议就尤为重要,本篇将主要介绍 MySQL 协议和在 Pisa-Proxy 中 MySQL 协议的 Rust 实现。

Why Rust

为什么要选用 Rust 语言呢?我们的考量有以下几个必要条件。

  • 安全性:首先作为数据库治理的核心组件,其语言的安全性是居首位的。Rust 中,类型安全实现内存安全,如所有权机制、借用、生命周期等特性避免了程序开发过程中的空指针、悬垂指针等问题,从而保证了服务在语言层面的安全性。
  • 优秀的性能表现:Rust 的目标在性能方面对标 C 语言,但在安全和生产力方面则比 C 更胜一筹。其无 GC,不需要开发人员手动分配内存等特性,极大程度地减少内存碎片,简化内存管理。
  • 低开销:从开发效率和可读可维护性上来说,有足够的抽象能力,并且这种抽象没有运行时开销(runtime cost)。零开销抽象,通过泛型和 Trait 在编译期展开并完成抽象解释。
  • 实用性:有优秀的包管理器工具 Crate、文档注释支持、详细的编译器提示、友好的错误处理等,在开发过程中能够高效帮助程序员快速开发出可靠、高性能的应用。

二、整体架构,模块设计

整体架构
如图 1,代理服务包含服务端和客户端的协议解析、SQL 解析、访问控制、连接池等模块。
背景1.png

工作流程
在图 1 中我们可以看出整个 Proxy 服务可以概括为以下几个阶段。

  1. 首先 Pisa-Proxy 支持 MySQL 协议,将自己伪装为数据库服务端,应用连接配置只需修改访问地址即可建连 Pisa-Proxy 通过读取应用发来的握手请求和数据包;
  2. 得到应用发来的 SQL 语句后对该 SQL 进行语法解析,并得到该 SQL 的 AST;
  3. 得到对应 AST 后,基于 AST 实现高级访问控制和 SQL 防火墙能力;
  4. 访问控制和防火墙通过后 SQL 提交执行,SQL 执行阶段的指标将采集为 Prometheus Metrics,最后根据负载均衡策略获取执行该语句的后端数据库连接;
  5. 如果连接池为空将根据配置建立和后端数据库的连接,SQL 将从该连接发往后端数据库;
  6. 最后读取 SQL 执行结果,组装后返回给客户端。

三、如何用 Rust 快速实现 MySQL 代理服务

如图 2,在整个代理服务中总体分为两部分:服务端和客户端,即代理服务作为服务端处理来自客户端的请求。和服务端,需要对服务端发起认证,并将客户端端命令发送给 MySQL 数据库。在这两部分中我们需要创建两套 Socket 来完成网络数据包的处理。

图片2.png

技术选型

介绍

在网络报处理和运行时处理上,我们选用了 Rust 实现的 Tokio(https://github.com/tokio-rs/tokio)框架。Tokio 框架是 Rust 编写的可靠、异步和精简应用程序的运行时。并且 Tokio 是一个事件驱动的非阻塞 I/O 平台,用于使用 Rust 编程语言编写异步应用程序。在高层次上,它提供了几个主要组件:

同时,Tokio 还提供了丰富的工具链,例如编解码工具包、分帧包等等,可以使我们更加方便快捷地处理 MySQL 中各种各样的数据报文。

项目地址: https://github.com/tokio-rs/tokio

优势

a. 快速:Tokio 的设计旨在使应用程序尽可能都快。

b. 零成本抽象:Tokio 以 Future 为基础。虽然 Future 并非 Rust 独创,但与其他语言的 Future 不同的是,Tokio Future 编译成了状态机,用 Future 实现常见的同步,不会增加额外开销成本。Tokio 的非阻塞 IO 可以充分发挥系统优势,例如实现类似 Linux Epoll 这种多路复用技术,在单个线程上的多路复用允许套接字并批量接收操作系统消息,从而减少系统调用,所有这些都可以减少应用程序的开销。

c. 可靠:Rust 的所有权模型和类型系统可以实现系统级应用程序,而不必担心内存不安全。

d. 轻量级:没有垃圾收集器,因为 Tokio 是基于 Rust 构建的,所以编译后的可执行文件包含最少的语言运行时。这意味着,没有垃圾收集器,没有虚拟机,没有 JIT 编译,也没有堆栈操作。这样在编写多线程并发的系统时,能够有效避免阻塞。

e. 模块化:每个组件都位于一个单独的库中。如果需要,应用程序可以挑选所需的组件,避免依赖不需要的组件。

代码实现

Rust 中数据包处理

#[derive(Debug)]
pub struct Packet {
    pub sequence: u8,
    pub conn: BufStream<LocalStream>,
    pub header: [u8; 4],
}

以上为 Proxy 数据包处理逻辑中核心的结构体,结构体中包含了三个字段分别为:

  • sequence:报文中的序列号 ID
  • conn:处理链接的 Socket
  • header:存储消息头报文的数组,长度为 4 字节

在包处理逻辑中主要定义了以下函数,在整个代理服务中网络数据交换都由以下方法来完成。

Pisa-Proxy 作为服务端

pub struct Connection {
    salt: Vec<u8>,
    status: u16,
    collation: CollationId,
    capability: u32,
    connection_id: u32,
    _charset: String,
    user: String,
    password: String,
    auth_data: BytesMut,
    pub auth_plugin_name: String,
    pub db: String,
    pub affected_rows: i64,
    pub pkt: Packet,
}

上面的结构体描述了 Pisa- Proxy 作为服务端处理来自于客户端请求时所包含的字段。例如,其中包含了和 MySQL 客户端进行认证信息,和所包含数据包处理逻辑的 Packet。

Pisa-Proxy 作为客户端

#[derive(Debug, Default)]
pub struct ClientConn {
    pub framed: Option<Box<ClientCodec>>,
    pub auth_info: Option<ClientAuth>,
    user: String,
    password: String,
    endpoint: String,
}

Tokio 提供的编解码器

在 Pisa-Proxy 中,大量使用了 Tokio 工具包中的编解码器,使用 codec Rust 会自动帮助开发者将原始字节转化为 Rust 的数据类型,方便开发者处理数据。使用编解码器,只需要在代码中为定义好的类型实现 Decoder 和 Encoder 两个 Trait,就可以通过 stream 和 sink 进行数据的读写。下面通过一个简单的示例来看一下 Tokio 编解码器的使用步骤。

使用 Tokio 编解码器一共分为三步:

  1. 首先要自定义一个错误类型,这里定义了一个 ProtocolError,并为它实现一个 from 方法,能够让它接收错误处理。
pub enum ProtocolError {
    Io(io::Error),
}

impl From<io::Error> for ProtocolError {
    fn from(err: io::Error) -> Self {
        ProtocolError::Io(err)
    }
}
  1. 定义一个数据类型,这里我们声明一个 Message 为 String,然后定义一个结构体,也就是我们要解析的是原始字节流要实际转换成的结构体,也即 Tokio 中 framed(帧的概念),这里定义一个空结构体 PisaProxy;
type Message = String;
struct PisaProxy;

接下来就是为 PisaProxy 分别实现 Encoder 和 Decoder 这两个 Trait。这里的示例,实现的功能为将数据转为 byte 数组,追加到 buf 中。在编码器中,我们首先要指定 Item 类型为 Message 和错误类型,编码处理逻辑这里将字符串进行拼接并返回给客户端。

在这里,Encoder 编码是指将用户自定义类型转换成 BytesMut 类型,写到 TcpStream 中,Decoder 解码指将读到的字节数据序列化为 Rust 的结构体。

impl Encoder<Message> for PisaProxy {
    type Error = ProtocolError;

    fn encode(&mut self, item: Message, dst: &mut BytesMut) ->Result<(), Self::Error> {
        dst.extend(item.as_bytes());
        Ok(())
    }
}

impl Decoder for PisaProxy {
    type Item = Message;
    type Error = ProtocolError;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        if src.is_empty() {
            return Ok(None);
        }
        let data = src.split();
        let mut buf = BytesMut::from(&b"hello:"[..]);
        buf.extend_from_slice(&data[..]);
        let data = String::from_utf8_lossy(&buf[..]).to_string();

        Ok(Some(data))
    }
}
  1. 当实例化 PisaProxy 结构体后,就可以调用 framed 方法,codec 的 framed 方法(codec.framed(socket))将 TcpStream 转换为 Framed<TcpStream,PisaProxy>,这个 Framed 就是实现了 tokio 中的 Stream 和 Sink 这两个 trait,实现的这两个 Trait 的实例就具有了接收(通过 Stream)和发送(通过 Sink)数据的功能,这样我们就可以调用 send 方法发送数据了。
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let addr = "127.0.0.1:9088";
    let listener = TcpListener::bind(addr).await?;
    println!("listen on: {:?}", addr);

    loop {
        let (socket, addr) = listener.accept().await?;

        println!("accepted connect from: {}", addr);

        tokio::spawn(async move {
            let codec = PisaProxy {};
            let mut conn = codec.framed(socket);
            loop {
                match conn.next().await {
                    Some(Ok(None)) => println("waiting for data..."),
                    Some(Ok(data)) => {
                        println!("data {:?}", data);
                        conn.send(data).await;
                    },
                    Some(Err(e)) => {
                    },
                    None => {},
                }
            }
        });
    }
}

四、MySQL 协议在 Pisa-Proxy 中的实现

MySQL 协议简介
MySQL 数据库本身是一个很典型的 C/S 结构的服务,客户端和服务端通信可以通过 Tcp 和 Unix Socket 的方式进行交互。在本篇中,我们主要说明通过网络 Tcp 的方式来实现 MySQL 代理。

MySQL 客户端和服务端的交互主要包含了两个重要的过程:1. 握手认证,2.发送命令。本篇会主要围绕这两个过程来介绍相关的实现。在客户端和服务端交互的过程中主要包含了以下几种类型报文:数据包、数据结束包、成功报告包以及错误消息包,在后面的章节中会为大家详细介绍这几种报文。

交互过程
MySQL 客户端和服务端在交互的过程中主要包含了两个过程,即握手认证阶段和执行命令阶段,当然在这两个过程之前首先要经历 TCP 三次握手的过程。在三次握手结束后首先进入握手认证阶段,在交换完信息并且客户端正确登录服务端后,进入执行命令阶段,图 3 完整描述了整个交互过程。

代码链接:https://github.com/database-mesh/pisanix/blob/master/pisa-proxy/protocol/mysql/src/server/conn.rs
图片3.png

握手认证

在握手认证阶段是 MySQL 客户端和服务端建联非常重要的阶段,该阶段发生在 TCP 三次握手之后。首先服务端会给客户端发送服务端信息,其中包括协议版本号、服务版本信息、挑战随机数、权能标志位等等。当客户端接收到服务端发来的响应之后,客户端开始发起认证请求。认证请求中,会携带客户端用户名、数据库名以及通过服务端响应中的挑战随机数将客户端密码加密后,一同发送给服务端进行校验。校验过程中,除了会对用户名密码进行校验,还会匹配客户端所使用的认证插件,如果不匹配则会发生插件的自动切换,以及判断客户端是否使用了加密链接。当以上阶段全部正常完成后,客户端则登录成功,服务端返回客户端 OK 数据报文。

上述过程分别在 runtime 和 protocol 的 server 中实现。在 runtime 中的 start 函数等待请求进入,Tcp 三次握手结束后,从 handshake 函数如图 3,开始握手阶段。在 handshake 中分别包含了三个过程,首先由 write_initial_handshake 给客户端发送服务端信息,然后在 read_handshake_response 客户端拿着服务端信息开始认证请求。

pub async fn handshake(&mut self) -> Result<(), ProtocolError> {
        match self.write_initial_handshake().await {
            Err(err) => return Err(err::ProtocolError::Io(err)),
            Ok(_) => debug!("it is ok"),
        }

        match self.read_handshake_response().await {
            Err(err) => {
                return Err(err);
            }
            Ok(_) => {
                self.pkt.write_ok().await?;
                debug!("handshake response ok")
            }
        }

        self.pkt.sequence = 0;

        Ok(())
    }

执行命令

当握手和认证阶段完成后,此时客户端才算是真正意义上的与服务端完成建立链接。那么此时则进入执行命令阶段。在 MySQL 中,能够发送命令的指令类型有很多种,我们会在下文中为大家进行介绍。

代码链接: https://github.com/database-mesh/pisanix/blob/master/pisa-proxy/runtime/mysql/src/server/server.rs

在下面的代码中可以看到,Pisa-Proxy 在这里会对不同类型的指令进行不同的逻辑处理。例如,初始化 db 的处理逻辑为 handle_init_db,处理查询的逻辑为 handle_query

match cmd {
          COM_INIT_DB => self.handle_init_db(&payload, true).await,
          COM_QUERY => self.handle_query(&payload).await,
          COM_FIELD_LIST => self.handle_field_list(&payload).await,
          COM_QUIT => self.handle_quit().await,
          COM_PING => self.handle_ok().await,
          COM_STMT_PREPARE => self.handle_prepare(&payload).await,
          COM_STMT_EXECUTE => self.handle_execute(&payload).await,
          COM_STMT_CLOSE => self.handle_stmt_close(&payload).await,
          COM_STMT_RESET => self.handle_ok().await,
          _ => self.handle_err(format!("command {} not support", cmd)).await,
}

MySQL 协议基本数据类型
在 MySQL 协议中主要有以下几种数据类型:

  • 整型值:MySQL 报文中整型值分别有 1、2、3、4、8 字节长度,使用小端传输。
  • 字符串(以 NULL 结尾)(Null-Terminated String):字符串长度不固定,当遇到’NULL’(0x00)字符时结束。
  • 二进制数据(长度编码)(Length Coded Binary)

参考链接:https://dev.mysql.com/doc/internals/en/basic-types.html

报文结构
在 MySQL 客户端和服务端所交互的数据最大长度为 16MByte,其基本数据报文结构类型如下:

图4.png
图 4
图5.png

图 5

在图 4 和图 5 中描述了 MySQL 报文的基本结构。报文包括了两部分,消息头和消息体。其中在消息头中,3 字节表示数据报文长度,1 字节存储序列号,消息体中存储实际的报文数据。

消息头

用于标记当前请求消息的实际数据长度值,以字节为单位,占用 3 个字节,最大值为 0xFFFFFF,即接 2^24-1。

序列号

序列 ID 从 0 开始随每个数据包递增,并且当进入新的命令时重置为 0。序列号 ID 有可能会发生回绕,当发生回绕时,需将序列号 ID 重置为 0,并且重新开始计数递增。

报文数据

消息体用于存放请求的内容及响应的数据,长度由消息头中的长度值决定。

客户端请求命令报文

该指令用于标识客户所要执行命令的类型,以字节为单位占用 1 个字节。请求命令报文常用到的有 Text 文本协议和 Binary 二进制协议。这里可以参考【执行命令】中的代码,代码中描述了对不同指令的处理逻辑。

在文本协议中常用到的有以下指令:

指令 功能
0x01 COM_QUIT 关闭连接
0x02 COM_INIT_DB 切换数据库
0x03 COM_QUERY SQL 查询请求
0x04 COM_FIELD_LIST 获取数据表字段信息
0x05 COM_CREATE_DB 创建数据库
0x06 COM_DROP_DB 删除数据库
0x08 COM_SHUTDOWN 停止服务器
0x0A COM_PROCESS_INFO 获取当前连接的列表
0x0B COM_CONNECT (内部线程状态)
0x0E COM_PING 测试连通性

更多请参考:https://dev.mysql.com/doc/internals/en/text-protocol.html

在二进制协议,即 Prepare Statement 中常用到的有以下指令:

指令 功能
0x16 COM_STMT_PREPARE 预处理 SQL 语句
0x17 COM_STMT_EXECUTE 执行预处理语句
0x18 CCOM_STMT_SEND_LONG_DATA 发送 BLOB 类型的数据
0x19 COM_STMT_CLOSE 销毁预处理语句
0x1A COM_STMT_RESET 清除预处理语句参数缓存

更多请参考:https://dev.mysql.com/doc/internals/en/prepared-statements.html

响应报文

  • OK 响应报文
  • Error 响应报文
  • Field 结构
  • EOF 结构
  • Row Data 结构

例如,下面代码中展示了如何给客户端写 OK 和 EOF 报文。

#[inline]
pub async fn write_ok(&mut self) -> Result<(), Error> {
  let mut data = [7, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0];
  self.set_seq_id(&mut data);
  self.write_buf(&data).await
}
  
#[inline]
pub async fn write_eof(&mut self) -> Result<(), Error> {
  let mut eof = [5, 0, 0, 0, 0xfe, 0, 0, 2, 0];
  self.set_seq_id(&mut eof);
  self.write_buf(&eof).await
}

com_query 请求流程

如图 6 为 com_query 的请求流程,com_query 指令可能会返回以下结果集:

  • ERR 报文
  • OK 报文
  • Protocol::LOCAL_INFILE_Request 报文
  • Resultset 报文

在 (https://github.com/database-mesh/pisanix/blob/master/pisa-proxy/protocol/mysql/src/client/resultset.rs) 文件中主要为对 ResultSet 结果集的处理,可以看到在这里定义了 ResultSet 结构,同样对 ResultSetCodec 分别实现了 Encoder 和 Decoder 两个 Trait,这样就可以通过编解码器来处理 ResultSet 的报文。

#[derive(Debug, Default)]
pub struct ResultsetCodec {
  pub next_state: DecodeResultsetState,
  pub col: u64,
  pub is_binary: bool,
  pub seq: u8,
  pub auth_info: Option<ClientAuth>,
}

图6.png
图 6

五、总结

以上,就是本篇文章的全部内容。在本篇文章中我们介绍了使用 Rust 实现 MySQL 代理的动机,介绍了 MySQL 协议中一些常用到的概念和 MySQL 中数据报文在网络中是如何进行交换数据的;并在最后介绍如何使用 Rust 去快速实现一个 MySQL 代理服务。更多实现细节请关注 Pisanix 代码仓库。

欢迎点击链接查看相关教学视频:https://www.bilibili.com/video/BV1B54y1o7EC?spm_id_from=333.999.0.0

六、相关链接

Pisanix
项目地址:https://github.com/database-mesh/pisanix

官网地址:https://www.pisanix.io/

Database Mesh:https://www.database-mesh.io/

Mini-Proxy:一个最小化的 MySQL Rust 代理实现

项目地址:https://github.com/wbtlb/mini-proxy

社区
目前 Pisanix 社区每两周都会组织线上讨论,详细安排如下,欢迎各位小伙伴一起参与进来。

讨论 地址
邮件列表 https://groups.google.com/g/database-mesh
英文社区双周会(2022年2月27日起),周三 9:00 AM PST https://meet.google.com/yhv-zrby-pyt
中文社区双周会(2022年4月27日起),周三 9:00 PM GMT+8 https://meeting.tencent.com/dm/6UXDMNsHBVQO
微信小助手 pisanix
Slack https://databasemesh.slack.com/
会议记录 https://bit.ly/39Fqt3x

数据运维技术 » 教你在 Pisa-Proxy 中,如何利用 Rust 实现 MySQL