1.前言
选择RisingWave是因为我本身对Flink有些了解,再加上RW的代码本身写得还不错,比较易读,因此以它来做为生产级Rust代码的学习对象。
本文基于RW v2.2.0。面向的读者主要是Rust初学者,希望通过这个系列文章可以让你快速了解生产级的Rust代码是什么样的。
2.读代码
为了方便大家理清思路,这边直接给出一个简单的代码地图。
rust
-- utils/pgwite/pg_server.rs
-- pg_serve
-- handle_connection
-- utils/pgwite/pg_protocol.rs
-- process
-- do_process #分析见2.1
-- do_process_inner
-- process_query_msg #选择query分支
-- inner_process_query_msg #SqlParser在这里做
-- inner_process_query_msg_one_stmt
-- frontend/session.rs
-- run_one_query
-- handle
-- handler/mod.rs
-- handle #选择Query Insert Delete Update分支
-- handle_query #在别的系列文章中会分析
-- frontend/handler/query.rs
-- execute
-- create_stream
-- distribute_execute
-- fronted/scheduler/distributed/query_manager.rs
-- schedule
-- fronted/scheduler/distributed/query.rs
-- start #分析见2.2
-- run
-- fronted/scheduler/distributed/stage.rs
-- start
-- StageRunner.run #分析见2.3
-- schedule_tasks_for_all
-- schedule_tasks
-- schedule_task #分析见2.4
-- rpc_client/compute_client.rs
-- create_task
-- prost/task_service.rs
-- create_task
-- fire_task
-- prost/task_execution.rs
-- async_execute
-- run
2.1 do_pcocess:weak与pin
这个函数主要用于处理来自客户端的消息。在函数的一开始我们可以看到一些奇怪的代码。
rust
async fn do_process(&mut self, msg: FeMessage) -> Option<()> {
let span = self.root_span_for_msg(&msg);
let weak_session = self
.session
.as_ref()
.map(|s| Arc::downgrade(s) as Weak<dyn Any + Send + Sync>);
// Processing the message itself.
//
// Note: pin the future to avoid stack overflow as we'll wrap it multiple times
// in the following code.
let fut = Box::pin(self.do_process_inner(msg));
这个代码有两个点可以拎出来说一说:
- Weak是干啥的
- Box::pin是干啥的
那为什么要从强引用变成weak_session呢?
因为上层的PgProtocol调用的这个代码,然后会有多个session进来,PgProtocol里维护了一个session: Option<Arc<SM::Session>>
, 每个session进来相当于会赋到这个成员上。而使用weak并不会添加Arc的引用,不过我并不觉得在这里用Arc会有什么问题,因为这个栈调用完以后,引用肯定会被释放掉。
那么Arc
是啥?(Atomic Reference Counting,原子引用计数)是标准库提供的一个智能指针类型,用于在多线程环境中安全地共享所有权。
Box::pin是干啥的?Box::pin 用于将一个异步操作(future)分配到堆上并将其固定(pin),确保该 future 在其生命周期内不会被移动。这是为了满足某些异步操作的要求,特别是当这些操作依赖于它们的内存位置不变时。
通过 Box::pin,代码可以安全地对 future 进行多次转换和包装,而不会导致不可预期的行为或性能问题。
2.2 start
这段代码主要负责查询的启动逻辑,检查并更新查询状态,创建阶段执行计划,设置超时任务以取消查询,启动查询执行器并返回结果获取器。
rust
/// Start execution of this query.
/// Note the two shutdown channel sender and receivers are not dual.
/// One is used for propagate error to `QueryResultFetcher`, one is used for listening on
/// cancel request (from ctrl-c, cli, ui etc).
#[allow(clippy::too_many_arguments)]
pub async fn start(
self: Arc<Self>,
context: ExecutionContextRef,
worker_node_manager: WorkerNodeSelector,
batch_query_epoch: BatchQueryEpoch,
compute_client_pool: ComputeClientPoolRef,
catalog_reader: CatalogReader,
query_execution_info: QueryExecutionInfoRef,
query_metrics: Arc<DistributedQueryMetrics>,
) -> SchedulerResult<QueryResultFetcher> {
//忽略一些代码
match cur_state {
QueryState::Pending { msg_receiver } => {
*state = QueryState::Running;
// Start a timer to cancel the query
let mut timeout_abort_task_handle: Option<JoinHandle<()>> = None;
if let Some(timeout) = context.timeout() {
let this = self.clone();
timeout_abort_task_handle = Some(tokio::spawn(async move {
tokio::time::sleep(timeout).await;
warn!(
"Query {:?} timeout after {} seconds, sending cancel message.",
this.query.query_id,
timeout.as_secs(),
);
this.abort(format!("timeout after {} seconds", timeout.as_secs()))
.await;
}));
}
- QueryState::Pending { msg_receiver } => { ... } 表示当 cur_state 是 QueryState::Pending 变体时,会进入这个分支。同时,它将 msg_receiver 字段从 Pending 变体中解构出来,并绑定到同名变量 msg_receiver,以便在后续代码中使用。
// Start a timer to cancel the query
是用来做超时检测,sleep几秒后开始检测,而tokio::spawn是起一个以协程为单位的异步任务。
2.3 StageRunner.run
实现了查询执行的异步运行逻辑。启动查询的叶子阶段,处理消息队列中的事件,如阶段调度、根阶段完成、阶段失败和查询取消,根据事件更新查询状态并处理相应的清理工作。
rust
impl StageRunner {
async fn run(mut self, shutdown_rx: ShutdownToken) {
if let Err(e) = self.schedule_tasks_for_all(shutdown_rx).await {
error!(
error = %e.as_report(),
query_id = ?self.stage.query_id,
stage_id = ?self.stage.id,
"Failed to schedule tasks"
);
self.send_event(QueryMessage::Stage(Failed {
id: self.stage.id,
reason: e,
}))
.await;
}
}
这里的代码是rust的if let形式,熟悉go语言的同学应该会觉得有点像。
rust
if str, ok := i.(string); ok {
fmt.Printf("The string is: %s\n", str)
} else {
fmt.Println("Not a string.")
}
- 调用 self.schedule_tasks_for_all(shutdown_rx) 并等待其完成。
- 如果该调用返回的结果是 Err(e)(即发生了错误),则将错误值绑定到变量 e,并执行后续的大括号内的代码块。
- 如果返回的是 Ok,则跳过这个 if let 块。
2.4 schedule_task
用于调度任务到计算节点。获取或选择一个工作节点,创建计算客户端并获取任务流状态,更新任务状态,返回任务流状态。
rust
async fn schedule_task(
&self,
task_id: PbTaskId,
plan_fragment: PlanFragment,
worker: Option<WorkerNode>,
expr_context: ExprContext,
) -> SchedulerResult<Fuse<Streaming<TaskInfoResponse>>> {
let mut worker = worker.unwrap_or(self.worker_node_manager.next_random_worker()?);
let worker_node_addr = worker.host.take().unwrap();
let compute_client = self
.compute_client_pool
.get_by_addr((&worker_node_addr).into())
.await
.inspect_err(|_| self.mask_failed_serving_worker(&worker))
.map_err(|e| anyhow!(e))?;
let t_id = task_id.task_id;
let stream_status: Fuse<Streaming<TaskInfoResponse>> = compute_client
.create_task(task_id, plan_fragment, self.epoch, expr_context)
.await
.inspect_err(|_| self.mask_failed_serving_worker(&worker))
.map_err(|e| anyhow!(e))?
.fuse();
self.tasks[&t_id].inner.store(Arc::new(TaskStatus {
_task_id: t_id,
location: Some(worker_node_addr),
}));
Ok(stream_status)
}
这里面出现一个有趣的结构体,叫做fuse。
fuse 是 Rust 中流(Stream)处理的一个常用方法,通常用于确保流在完成或出错后不再产生额外的项。具体来说:
-
当流调用 fuse 方法后,会返回一个新的 Fuse 包装器。
-
如果流已经完成(即调用了 poll_next 并返回了 None),那么后续的所有轮询都会立即返回 None,而不会再次尝试从原始流中获取数据。
-
这样可以避免对已完成流的重复轮询,防止不必要的计算或副作用。
简单来说,fuse 方法确保流一旦结束就不会再被重新激活,从而提供更好的行为控制和性能优化。
rust
pin_project! {
/// Stream for the [`fuse`](super::StreamExt::fuse) method.
#[derive(Debug)]
#[must_use = "streams do nothing unless polled"]
pub struct Fuse<St> {
#[pin]
stream: St,
done: bool,
}
}
2.5 run
异步任务的执行和状态管理。初始化任务并增加任务计数,处理数据流。发送数据块并处理发送错误。监听关闭信号并根据信号类型设置任务状态。关闭通道并通知状态变化。减少任务计数。
ini
async fn run(
&self,
root: BoxedExecutor,
mut sender: ChanSenderImpl,
state_tx: Option<&mut StateReporter>,
) {
self.context
.batch_metrics()
.as_ref()
.inspect(|m| m.batch_manager_metrics().task_num.inc());
let mut data_chunk_stream = root.execute();
let mut state;
let mut error = None;
let mut shutdown_rx = self.shutdown_rx.clone();
loop {
select! {
biased;
// `shutdown_rx` can't be removed here to avoid `sender.send(data_chunk)` blocked whole execution.
_ = shutdown_rx.cancelled() => {
match self.shutdown_rx.message() {
ShutdownMsg::Abort(e) => {
error = Some(BatchError::Aborted(e));
state = TaskStatus::Aborted;
break;
}
ShutdownMsg::Cancel => {
state = TaskStatus::Cancelled;
break;
}
ShutdownMsg::Init => {
unreachable!("Init message should not be received here!")
}
}
}
data_chunk = data_chunk_stream.next()=> {
match data_chunk {
Some(Ok(data_chunk)) => {
if let Err(e) = sender.send(data_chunk).await {
match e {
BatchError::SenderError => {
// This is possible since when we have limit executor in parent
// stage, it may early stop receiving data from downstream, which
// leads to close of channel.
warn!("Task receiver closed!");
state = TaskStatus::Finished;
break;
}
x => {
error!("Failed to send data!");
error = Some(x);
state = TaskStatus::Failed;
break;
}
}
}
}
Some(Err(e)) => match self.shutdown_rx.message() {
ShutdownMsg::Init => {
// There is no message received from shutdown channel, which means it caused
// task failed.
error!(error = %e.as_report(), "Batch task failed");
error = Some(e);
state = TaskStatus::Failed;
break;
}
ShutdownMsg::Abort(_) => {
error = Some(e);
state = TaskStatus::Aborted;
break;
}
ShutdownMsg::Cancel => {
state = TaskStatus::Cancelled;
break;
}
},
None => {
debug!("Batch task {:?} finished successfully.", self.task_id);
state = TaskStatus::Finished;
break;
}
}
}
}
}
//忽略一些代码
}
相信很多人看到select! {
会很懵逼。select! 是个宏,具体含义如下:
- loop { ... } 创建一个无限循环,持续执行其中的代码块,直到遇到 break 语句。
- select! { biased; ... } 是 Tokio 提供的一个宏,用于在多个异步操作之间进行选择。biased 关键字确保优先处理第一个分支,避免其他分支抢占先机。
宏就是在编译期时展开的代码逻辑增强。写过Java的同学可以理解为AOP的AspectJ实现。