Rust中如何为Future/async fn添加必执行的终止清理逻辑？

可以为Future或async fn添加在终止时（无论正常完成还是被丢弃）始终运行的清理代码，核心思路是借助Rust的Drop特性结合Future封装来实现。

实现原理 #

Rust的Drop trait会在值被销毁时自动执行指定逻辑，我们可以把原Future和清理代码封装到自定义结构体中：

• 让结构体实现Future，代理原Future的执行流程
• 实现Drop trait，确保结构体被丢弃时触发清理
• 在原Future正常完成时，手动触发清理（避免Drop重复执行，同时保证完成后立即执行清理）

示例代码 #

基础封装实现 #

use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

struct WithCleanup<F, C> {
    future: Option<F>,
    cleanup: Option<C>,
}

impl<F, C> WithCleanup<F, C> {
    fn new(future: F, cleanup: C) -> Self {
        Self {
            future: Some(future),
            cleanup: Some(cleanup),
        }
    }
}

// 实现Drop，确保丢弃时执行清理
impl<F, C: FnOnce()> Drop for WithCleanup<F, C> {
    fn drop(&mut self) {
        // 取出清理逻辑并执行，确保只执行一次
        if let Some(cleanup) = self.cleanup.take() {
            cleanup();
        }
    }
}

// 实现Future，代理原Future的poll逻辑
impl<F: Future + Unpin, C: FnOnce()> Future for WithCleanup<F, C> {
    type Output = F::Output;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let future = self.future.as_mut().expect("future already completed");
        match Pin::new(future).poll(cx) {
            Poll::Ready(output) => {
                // 正常完成时先执行清理
                if let Some(cleanup) = self.cleanup.take() {
                    cleanup();
                }
                self.future.take(); // 标记Future已完成
                Poll::Ready(output)
            }
            Poll::Pending => Poll::Pending,
        }
    }
}

// 辅助函数，简化调用
fn with_cleanup<F, C>(future: F, cleanup: C) -> WithCleanup<F, C> {
    WithCleanup::new(future, cleanup)
}

// 使用示例
async fn short_task() {
    println!("短任务执行中");
    std::thread::sleep(std::time::Duration::from_secs(1));
    println!("短任务完成");
}

async fn long_task() {
    println!("长任务执行中");
    std::thread::sleep(std::time::Duration::from_secs(5));
    println!("长任务完成");
}

#[tokio::main]
async fn main() {
    // 场景1：Future正常完成
    println!("=== 正常完成场景 ===");
    with_cleanup(short_task(), || {
        println!("清理逻辑执行（正常完成）");
    }).await;

    // 场景2：Future中途被丢弃
    println!("=== 中途丢弃场景 ===");
    let task_handle = tokio::spawn(with_cleanup(long_task(), || {
        println!("清理逻辑执行（中途丢弃）");
    }));

    std::thread::sleep(std::time::Duration::from_secs(1));
    task_handle.abort(); // 主动终止任务，触发Future丢弃
}

关键注意事项 #

• 避免重复执行：通过take()转移清理逻辑的所有权，确保清理代码只会执行一次（无论正常完成还是被丢弃）。
• 安全约束：Drop中的逻辑不能panic，否则会直接导致程序终止；若清理可能出错，需内部处理错误。
• Unpin适配：示例要求原Future实现Unpin，若需处理非Unpin的Future，可将结构体中的F改为Pin<Box<F>>。

现成方案 #

如果不想手动实现，可借助第三方库（如async-cleanup），其核心逻辑同样基于Drop特性；使用Tokio时，也可通过JoinHandle的abort结合自定义封装实现类似效果。

内容的提问来源于stack exchange，提问作者laptou