TypeScript 开发者用 Rust 实现代码补全

我的文章《使用 Code Llama 直接替代 Copilot 代码补全》收到了很多积极反馈。此后，我又进行了一些尝试，希望进一步改进 Copilot 服务器。

在性能方面，我为 exllamav2 中的 Copilot 服务器提交了一个 PR，使用了 exllama2 速度极快的自定义 CUDA 内核。

为了提高补全质量，我还尝试了其他几个 LLM，例如 replit-code-v1_5-3b、long_llama_code_7b、CodeShell-7B 和 stablelm-3b-4e1t。

作为一名一线开发者，同时又是一个无法使用 H100 集群的 GPU 穷人，为了在这个 AI「西部荒野」时代做出贡献，我更感兴趣的是改善 Copilot 服务器的易用性。

Hugging Face 的 candle 是一个面向 Rust 的极简机器学习框架，看起来非常有趣。因此，我开始用 Rust 🦀 创建一个极简的 Copilot 服务器。

在继续阅读之前，请注意：使用 Python + CUDA 的 exllamav2 版本仍然比 Rust 版本快得多。本文主要面向那些对学习 Rust 感兴趣，并希望通过构建一个有趣的项目来学习编程的人。

本质上，这是一篇「用 Rust 🦀 构建你自己的 Copilot」教程，代码以教学为目的。如果你只想试用最终产品 oxpilot：

brew install chenhunghan/homebrew-formulae/oxpilot

然后启动 Copilot 服务器：

ox serve

ox hi in Japanese

我们将使用 axum 作为 Web 框架，使用 candle 进行文本推理，使用 clap 解析 CLI 参数，并使用 tokio 作为异步运行时。

用于调试的打印（console.log） 美化打印 性能测量

获得安全感 变量默认不可变 你不应该移动！所有权 所有权与作用域 借用

变量默认不可变

你不应该移动！所有权 所有权与作用域 借用

异步 Rust 并行 并发 任务（绿色线程） 异步运行时 所有权与异步 在异步程序中共享状态：Arc 和 Mutex

在异步程序中共享状态：Arc 和 Mutex

动手实践 Server-Sent Events（SSE）服务器 对端点进行 BDD 构建者模式 ::default() 与 ::new() 对函数参数使用 impl Into<String> 类型状态 共享内存 Arc<Mutex<_>> 通过通信共享内存：Actor

Server-Sent Events（SSE）服务器 对端点进行 BDD

构建者模式 ::default() 与 ::new() 对函数参数使用 impl Into<String>

::default() 与 ::new()

对函数参数使用 impl Into<String>

共享内存 Arc<Mutex<_>>

通过通信共享内存：Actor

（部分章节仍在编写中）

如果你已经在一定程度上熟悉 Rust，例如对 Rust 的所有权和借用机制已经比较熟悉，但对异步世界还不熟悉，我建议直接跳到异步章节。

如果你已经熟悉异步 Rust，可以直接前往动手实践部分，其中介绍了一些你可能会觉得有用的设计模式；也可以直接前往 GitHub 上的 oxpilot 项目，所有内容都已经开源。

请做好遇到一些甚至许多人为错误的准备。我记录了自己的学习过程，希望它能帮助互联网上那些喜欢通过构建令人兴奋的项目来学习新语言的人。

感谢 jihchi 审阅本文草稿。

本文是自包含的，也就是说，它应该涵盖阅读 oxpilot 源代码所需了解的全部内容。

不过，每一节都不可能面面俱到。我尽量在各节末尾提供参考资料；如果你还没有读过《The Rust Programming Language》，强烈建议阅读。

Tracing 是开箱即用的 console.log

console.log 是 TypeScript 中一个强大的工具，你可以打印任何想要的内容，因此 console.log 对调试非常有用。Rust 中与之对应的是 print!。如果你想快速上手使用 print!，Rust by Example 是一份非常优秀的文档。

不过，print! 会阻塞 stdio，而更好的做法是锁定 stdio，并手动解锁，但这非常繁琐。

幸运的是，我们还有其他选择。Tracing 是 tokio 团队开发的一个非常棒的项目。作为 TypeScript 开发者，使用 tracing 让我感觉非常熟悉。

info!("Hello! Rust!");
info!("Print var: {:?}", var);

你可能想知道代码块中的 {:?} 是什么。

info!("Print var: {:?}", var);

{:?} 用于打印结构体（类似 TypeScript 中的 Object）。也可以使用 {:#?} 进行美化打印，详情请参阅相关文档；可以把它理解为 console.log(JSON.stringify(object,null,2))。

Tracing 非常适合记录性能指标。例如，如果我想测量 awesome() 的执行时间：

async fn awesome() {}
awesome().instrument(tracing::info_span!("awesome")).await;

它会打印非常有用的信息，告诉我们何时开始调用 awesome()、调用发生在哪一行、位于哪个线程，以及函数执行了多长时间。

2023-10-22T09:01:13.128553Z INFO ThreadId(01) awesome src/main.rs:172: enter
2023-10-22T09:01:13.128569Z INFO ThreadId(01) awesome src/main.rs:172: close time.busy=15.3µs time.idle=3.96µs

Rust 默认是安全的，这里的安全通常指内存安全。不过根据我的经验，一旦代码成功编译，Rust 还会让你在将它发布到生产环境时感到安心。

如果你曾经写过 JavaScript，然后转而使用 TypeScript，那么你大概明白我所说的「感到安心」是什么意思。

TypeScript 能够在编译阶段保护我们，避免出现 TypeError: Cannot read property '' of undefined。Rust 就像开启了超严格模式的 TypeScript，在编译阶段保护我们这些开发者，避免犯错。

Rust 让拉取请求更容易审查，也增强了将代码发布到生产环境的信心。编译器错误信息一开始可能显得令人不知所措，就像刚开始接触 TypeScript 错误时一样。

不过，如果你曾经承受过恢复生产服务器的压力，就会知道：学习解决编译时错误，要好过处理运行时异常。

要拥抱 Rust 提供的安全网，不可变性和所有权是必须理解的两个关键概念。

变量默认不可变

「默认不可变」意味着数据一旦创建，就不能再被修改。大多数人都会认同，不可变数据能让代码变得更好。

例如，一名经验丰富的 TypeScript 开发者大概知道使用 const 的好处：当你尝试修改值时，const 能明确表达代码的意图。

const x = 5;
x = 1; // Cannot assign to 'x' because it is a constant.

在 Rust 中，变量默认不可变，只有显式声明为可变时才能修改。

fn main() {
    let x = 5; // this does not compile, 
    x = 6; // explicit `let mut x` to make mutation possible.
}

《The Rust Programming Language》中的 Variables and Mutability 一节对 Rust 的可变性进行了全面讲解。

你不应该移动！所有权

对于来自带垃圾回收器语言的开发者来说，下面的代码看起来非常自然：我们尝试通过引用 string1 来创建 string2：

fn main() {
  let s1 = String::from("hello");
  let s2 = s1;
  println!("{}", s1);
}

不过，这段代码无法编译。编译器说你已经移动了 s1。

11 |   let s1 = String::from("hello");
   |       -- move occurs because `s1` has type `String`, which does not implement the `Copy` trait
12 |   let s2 = s1;
   |            -- value moved here
13 |   println!("{}", s1);
   |                  ^^ value borrowed here after move

这可能是你刚开始学习 Rust 时遇到的第一个、并且会持续令人沮丧的编译器错误。

Rust 没有自带垃圾回收器，这意味着它无法在运行时知道：当你不再需要某个值时，应该何时将其从内存中释放。

为了实现这一目标，Rust 引入了所有权检查器，让开发者在其余代码不再需要某个值时对它进行标记。所有权检查器帮助你在编译阶段管理内存，因此不需要让代码携带一个垃圾回收器，在运行时收集并从内存中释放不再使用的值。

上面示例代码中的 value moved here 表示代码违反了所有权规则。这些规则是：

Rust 中的每个值都有一个所有者。

同一时间只能有一个所有者。

当所有者离开作用域时，该值将被丢弃。

编译器是在告诉你：嘿！s1 是 String::from("hello") 的所有者，但你已经将所有权从 s1 移动到了 s2。既然你不再需要 s1，编译器就丢弃了 s1，因此你不应该再次在 println! 中使用它！

fn main() {
  let s1 = String::from("hello");
  let s2 = s1; // ownership moved from s1 to s2
  println!("{}{}", s1); // s1 is dropped, why you are still using it?
}

如果你来自 TypeScript 世界（或者任何带垃圾回收器的语言），所有权可能看起来很陌生。不过，学习所有权检查器会让你更加关注内存分配。

让我们再次回顾所有权规则，并进一步深入理解第三条规则。

Rust 中的每个值都有一个所有者。

同一时间只能有一个所有者。

当所有者离开作用域（花括号 {}）时，该值将被丢弃。

在下面的示例中，编译器会在第二次调用 do_something() 时阻止我们，因为我们将 owner 移动进了 do_something，然后又尝试再次使用 owner，从而违反了所有权规则。

这段代码无法编译：

fn main() {
    let owner = String::from("value");
    // 我们从 `owner` 取得了 "value" 的所有权，并且
    // "value" 在 `do_something` 函数结束时被丢弃
    // 因此变量 `owner` 不再拥有它
    do_something(owner);
    // 使用已移动的值: `owner` 在移动后被使用
    print!("{}", owner);
}

fn do_something(_: String) {
    // 
}

print!("{}", owner) 违反了所有权规则，因为我们已经将 owner 移动到了 do_something(owner) 的作用域，因此在 do_something(owner) 执行完毕后，owner 超出了作用域，被丢弃，我们无法再使用它。

为了绕过所有权规则，需要借用来拯救。

借用是使用引用语法 (&) 让 Rust 编译器知道我们只是借用而不是取得所有权，借用使用引用来承诺我们只是临时借用这个值，不打算取得所有权，当不再需要时会返回这个值。

fn main() {
    let owner = String::from("value");
    // `do_something` 从 `owner` 借用 `"value"`
    do_something(&owner);
    // 没有错误了！
    print!("{}", owner);
}

fn do_something(_: &String) {
    // "value" 在函数结束时不会被丢弃
    // 因为我们只是借用 (`&String`) 而不是取得所有权
}

就像所有权一样，借用也有一套规则，这些规则就像你从别人那里借东西时签订的契约。

现实生活类比：你想从朋友那里借一本《Rust for Rustaceans》的书，为了保持友谊，你签订了一份契约（口头承诺："我会在一个月内把借来的书还给你"），契约需要遵守借用规则：

在任何时刻，你可以有任意多个不可变引用，但只能有一个可变引用。

引用必须指向一个有效的值（不允许指向一个已丢弃的值）。

如果所有权和借用仍然显得模糊，《Rust 权威指南》的所有权章节是最好的阅读资料，在经历几次数据传递和编译器的叫喊之后，你很快就会熟悉所有权规则。

如果你是一个忙碌的开发者，Let's Get Rusty 的《The Rust Survival Guide》是快速掌握所有权规则的好方法。

在我们开始本节之前，让我们固定术语的定义。

异步是一种编程语言中的特性，旨在为程序提供机会在等待另一个计算单元完成时执行一个计算单元。

并行是程序同时执行计算单元，即真正的并发执行，例如在 CPU 的两个不同核心中运行两个计算。

并发是程序依次处理计算单元，当一个单元取得进展时，迅速从一个单元让步到另一个单元。程序在单元之间快速让步，看起来像同时执行单元一样（但实际上不是真正的并行）。参考单线程的 Node.js 运行时。

任务是指在并行或并发系统中运行的某些计算。在本文中，"任务"一词指的是异步绿色线程，它不是操作系统线程，而是由异步运行时管理的执行单元。

运行时（任务运行器）

Node.js 是单线程的异步运行时，程序可以异步处理任务，但程序不是并行处理任务，因为 Node.js 是单线程的。

要在 Rust 中异步处理任务，开发者需要设置一个任务运行器。主函数（想象它像 index.ts），这是一个 Rust 程序的入口点，始终是同步的，开发者需要设置运行时来能够在 Rust 中运行异步任务。

下面的代码使用 futures::executor 作为异步任务运行器。

fn main() {
    // 异步任务运行器
    futures::executor::block_on(do_something());
}

// 一个异步任务
async fn do_something() {
    //
}

在 Rust 中，你可以自由选择任何异步运行时，就像在 TypeScript 中一样，我们有 node.js、bun 和 deno。在 Rust 中我们有 tokio、async-std、smol 和 futures，这些运行时可以是单线程的，像 node.js 一样并发运行任务，或多线程的真正并行。

你可能会发现这些视频有助于理解 Rust 中的异步/等待。

1 Hour Dive into Asynchronous Rust

Async/await in Rust: Introduction

在"你不应该移动！所有权"中，我们讨论了所有权规则，在"借用"部分，我们讨论了如何通过借用来绕过所有权规则。

在异步 Rust 中，无论你使用的是单线程并发绿色线程运行时，还是将计算分布到多个 OS 线程（并行），所有权规则总是适用的。在异步 Rust 中，所有权规则在 Rust 的并发编程或并行编程中防止数据竞争。（也被称为"无畏并发"）

还记得所有权规则吗？

• Rust 中的每个值都有一个所有者。
• 在任何时刻只能有一个所有者。
• 当所有者超出作用域时，值将被丢弃。

特别注意"在任何时刻"，这就是所有权如何帮助我们在同时运行计算时避免数据竞争的（= 异步）。

让我们再看一遍同步版本，在之前的例子中这无法编译...

fn main() {
    let owner = String::from("value");
    do_something(owner);
    // 使用已移动的值: `owner` 在移动后被使用
    print!("{}", owner);
}

fn do_something(_: String) {
    // 
}

playground ...因为代码不遵守所有权规则，即 do_something() 取得了 String::from("hello") 的所有权，但 Rust 编译器只允许在任何时刻有一个所有者。为了保护我们免于忘记释放内存，所有者被移动到了第一个 do_something(owner)，我们无法编译代码，因为会出现这个错误"使用已移动的值: owner 在移动后被使用"。

3 |     do_something(owner);
  |                  ----- value moved here
4 |     // use of moved value: `owner` value used here after move
5 |     print!("{}", owner);
  |                  ^^^^^ value borrowed here after move

我们可以通过借用(&) 来绕过这个问题

fn main() {
    let owner = String::from("value");
    // 使用 & 来引用 owner
    do_something(&owner);
    // 之后我们仍然可以使用 owner
    print!("{}", owner);
}

fn do_something(_: &String) {
    // 
}

相同的所有权规则也适用于异步 Rust，让我们看看并行版本，它生成 OS 线程同时运行代码：

use std::thread;

fn main() {
    let owner = String::from("value");
    thread::spawn(|| {
        do_something(&owner);
    });
}
fn do_something(_: &String) {
    // 
}

我们知道我们需要使用借用(&) 来在调用 do_something(&owner) 时避免取得所有权。然而，编译器仍然拒绝了，它说：

closure may outlive the current function, but it borrows owner, which is owned by the current function

这个编译器错误告诉我们 owner 的借用可能会在线程闭包外引用一个值，在某个时刻，当这个值被丢弃时，违反了我们在借用中讨论的规则。

引用必须指向一个有效的值（不允许指向一个已丢弃的值）。

为了给这个 outlive 错误更多的上下文，尝试在 playground 中运行这段代码。

use std::thread;

fn main() {
    thread::spawn(|| {
        print!("from thread");
    });
    print!("from main");
}

你可能会惊讶地发现控制台中只有 from main。这是因为 Rust 的 std 中的线程实现允许创建的线程的生命周期超过创建它们的线程，换句话说，父线程（在我们的例子中是 main()）创建了子线程，通过 thread::spawn 创建的子线程可能会超过父线程（main()）。

这就是你在控制台看到 from main 的原因，|| print!("from thread") 的执行生命周期超过了 main 的执行。

如果我们退一步，在更高的借用层面思考线程：

use std::thread;

fn main() {
    let owner = String::from("value");
    thread::spawn(|| {
        // 我们借用 owner，但被借用的值 (`owner`)
        // 可能在 main() 中被丢弃，这意味着 `&` 可能指向
        // 一个已丢弃的值
        do_something(&owner);
    });
}
fn do_something(_: &String) {
    // 
}

我们同时从 main 和一个线程运行代码，在同一时刻，编译器通过告诉我们"闭包（在线程中）的生命周期可能超过当前函数，但它借用了 owner，owner 由 main() 拥有"来阻止我们，我们不应该这样做，因为当 owner 在父线程（main()）中无效时，我们可能会引用它。

同样的 outlive 问题也可以在并发代码中观察到，即使在大多数并发运行时中，代码执行不是在 OS 线程中而是在任务中：

/*
[dependencies]
tokio = { version = "1.32.0", features = ["full"] }
*/

#[tokio::main]
async fn main() {
    let owner = String::from("value");

tokio::spawn(do_something(&owner));
}

async fn do_something(_: &String) {
    //
}

这个代码块失败了，显示类似错误信息"owner"生命周期不够长。

为了克服这个线程不能借用（Threads Don't Borrow）错误，即克服不允许从父线程/任务向子线程/任务引用值的所有权规则。我们有几个解决方案：

1. 将值移动到线程中（详细了解）

use std::thread;

fn main() {
    let owner = String::from("value");
    // 将所有者移动到生成的线程中
    thread::spawn(move || {
        do_something(&owner);
    });
}

fn do_something(_: &String) {
    // 
}

2. 使用 scoped 线程，它在父线程（main）退出之前退出。

use std::thread;

fn main() {
    let owner = String::from("value");
    // scoped 线程总是在主线程退出之前存在
    // 因此我们可以使用对 `owner` 的引用
    thread::scope(|_| {
        do_something(&owner);
    });
}

fn do_something(_: &String) {
    // 
}

"不要通过共享内存来通信；而是通过通信来共享内存"，如Go语言文档中所说。我们将在actor部分深入探讨这一点。

原子引用计数（Arc<T>）和互斥锁（Mutex<T>）。

我们将在下一部分深入探讨（Arc<T>）。

在异步程序中共享状态：Arc和Mutex

在异步程序中共享状态可能是一个挑战。所有权规则只允许一个值在同一时间有一个所有者。我们不能使用借用，因为编译器不知道线程/任务中的借用者是否会在某个时刻指向已被丢弃的值。

为了解决这个问题，我们可以使用Arc（原子引用计数）。

Arc可以安全地用于在多个线程/任务之间共享状态。将数据包装在Arc中以拥有相同数据的多个副本：

use std::thread;
use std::sync::Arc;

fn main() {
    let arc = Arc::new(String::from("value"));
    thread::spawn(|| {
        do_something(arc);
    });
}

fn do_something(_: Arc<String>) {
    // 
}

Arc允许在线程之间安全地读取内部数据，它类似于借用但用于异步代码块。

但是，Arc只允许读取。要使线程能够写入内部数据，数据需要用适当的锁定机制处理，即（Mutex<T>）。

Mutex<T>（读作：互斥）将阻止等待锁可用的线程。当在线程上调用lock()时，该线程将成为唯一可以访问数据的线程，Mutex<T>阻止其他线程访问数据，因此在锁未释放时改变数据是安全的。

为了安全地改变我们与状态共享的数据：

use std::thread;
use std::sync::{Arc, Mutex};

fn main() {
    let inner_data = String::from("Hello ");
    let mutex = Arc::new(Mutex::new(inner_data));
    let mutex_clone = mutex.clone();
    thread::spawn(move || {
        let mut inner_data = mutex.lock().unwrap();
        inner_data.push_str(" world (once)!")
    });
    thread::spawn(move || {
        let mut inner_data = mutex_clone.lock().unwrap();
        inner_data.push_str(" world (twice)!")
    });
}

我们将在"共享内存Arc<Mutex<>>"部分深入探讨如何使用Arc<Mutex<>>来共享可变状态。

要了解更多关于共享状态的信息：

• The book的共享状态并发章节。
• Tokio文档中有专门的页面介绍如何在异步任务之间共享状态。

在下面的部分中，我们将开始构建copilot服务器。

服务器发送事件（SSE）服务器

在这个PR中，我们为copilot客户端添加端点。

从《Code Llama作为Copilot代码补全的替代品》一文我们知道，copilot服务器本质上是一个HTTP服务器，它接受包含提示的请求，并以服务器发送事件（SSE）的形式返回JSON块。让我们尝试指定SSE端点并使用axum创建服务器发送事件（SSE）服务器。

端点的URL路径：/v1/engines/:engine/completions

端点应接受POST请求。

端点采用路径参数（:engine）和请求体。

端点返回文本块的SSE流（Content-Type: text/event-stream）。

由于此端点几乎与OpenAI的completions端点相同，我们可以使用curl来查看输入（请求体）和输出（SSE文本块）：

curl https://api.openai.com/v1/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "gpt-3.5-turbo-instruct",
    "prompt": "Say this is a test",
    "max_tokens": 7,
    "temperature": 0,
    "stream": true
  }'
# chuck 0
data: {"choices":[{"text":"This ","index":0,"logprobs":null,"finish_reason":null}],
"model":"gpt-3.5-turbo-instruct", "id":"...","object":"text_completion","created":1}
# chuck 1
data: {"choices":[{"text":"is ","index":0,"logprobs":null,"finish_reason":null}],
"model":"gpt-3.5-turbo-instruct", "id":"...","object":"text_completion","created":1}
# chuck 2
data: {"choices":[{"text":"a ","index":0,"logprobs":null,"finish_reason":null}],
"model":"gpt-3.5-turbo-instruct", "id":"...","object":"text_completion","created":1}
# chuck with `"finish_reason":"stop"`
data: {"choices":[{"text":"test.","index":0,"logprobs":null,"finish_reason":"stop"}],
"model":"gpt-3.5-turbo-instruct", "id":"...","object":"text_completion","created":1}
# SSE事件流结束
data: [DONE]

在测试中，我们将使用reqwest_eventsource及其相关库充当客户端，向我们的端点/v1/engines/:engine/completions发送请求，并断言响应符合预期。由于reqwest_eventsource及其相关库不在最终二进制文件中使用，让我们在Cargo.toml的dev-dependencies中添加它们。

[dev-dependencies]
reqwest = { version = "0.11.22", features = ["json", "strea