生产级Rust代码品鉴(一)RisingWave一条SQL到运行的流程

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));

这个代码有两个点可以拎出来说一说:

  1. Weak是干啥的
  2. 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实现。

相关推荐
该用户已不存在4 分钟前
不知道这些工具,难怪的你的Python开发那么慢丨Python 开发必备的6大工具
前端·后端·python
Xy91010 分钟前
开发者视角:App Trace 一键拉起(Deep Linking)技术详解
java·前端·后端
嘻嘻哈哈开森15 分钟前
技术分享:深入了解 PlantUML
后端·面试·架构
vvw&20 分钟前
Linux 中的 .bashrc 是什么?配置详解
linux·运维·服务器·chrome·后端·ubuntu·centos
厚道26 分钟前
Elasticsearch 的存储原理
后端·elasticsearch
不甘打工的程序猿27 分钟前
nacos-client模块学习《心跳维持》
后端·架构
方块海绵27 分钟前
mysql 中使用 json 类型的字段
后端
今夜星辉灿烂29 分钟前
nestjs微服务-系列4
javascript·后端
敏叔V58734 分钟前
SpringBoot实现MCP
java·spring boot·后端