Rust数据加载与批处理库dataloader的使用，高效实现GraphQL批量查询与缓存优化

Dataloader简介

Rust实现的DataLoader库，基于Facebook的DataLoader概念，使用async-await特性实现高效的批量数据加载和缓存。

主要特性

• 支持带缓存的批量数据加载
• 支持不带缓存的批量数据加载
• 异步处理机制
• 可配置的批量大小

安装与配置

运行时选择

默认使用async-std运行时：

[dependencies]
dataloader = "0.18"

使用Tokio运行时：

[dependencies]
dataloader = { version = "0.18", default-features = false, features = ["runtime-tokio"]}
futures = "0.3"

基础使用示例

use dataloader::cached::Loader;
use dataloader::BatchFn;
use futures::executor::block_on;
use futures::future::ready;
use std::collections::HashMap;
use std::thread;

// 定义批量加载函数
struct MyLoadFn;

impl BatchFn<usize, usize> for MyLoadFn {
    async fn load(&mut self, keys: &[usize]) -> HashMap<usize, usize> {
        println!("批量加载键: {:?}", keys);
        keys.iter()
            .map(|v| (v.clone(), v.clone()))
            .collect::<HashMap<_, _>>()
    }
}

fn main() {
    let mut i = 0;
    while i < 2 {
        let a = MyLoadFn;
        // 创建Loader实例，设置最大批量大小为4
        let loader = Loader::new(a).with_max_batch_size(4);

        // 线程1中加载数据
        let l1 = loader.clone();
        let h1 = thread::spawn(move || {
            let r1 = l1.load(1);
            let r2 = l1.load(2);
            let r3 = l1.load(3);
            let r4 = l1.load_many(vec![2, 3, 4, 5, 6, 7, 8]);
            
            // 等待所有加载完成
            let f = futures::future::join4(r1, r2, r3, r4);
            println!("线程1结果: {:?}", block_on(f));
        });

        // 线程2中加载数据
        let l2 = loader.clone();
        let h2 = thread::spawn(move || {
            let r1 = l2.load(1);
            let r2 = l2.load(2);
            let r3 = l2.load(3);
            let r4 = l2.load(4);
            
            // 等待所有加载完成
            let f = futures::future::join4(r1, r2, r3, r4);
            println!("线程2结果: {:?}", block_on(f));
        });

        h1.join().unwrap();
        h2.join().unwrap();
        i += 1;
    }
}

GraphQL集成完整示例

use async_graphql::{Context, FieldResult, Object};
use dataloader::{BatchFn, Loader};
use futures::executor::block_on;
use std::collections::HashMap;
use std::sync::Arc;

// 用户数据模型
#[derive(Clone)]
struct User {
    id: i32,
    name: String,
}

// 用户数据加载器
struct UserLoader;

#[async_trait::async_trait]
impl BatchFn<i32, User> for UserLoader {
    async fn load(&mut self, keys: &[i32]) -> HashMap<i32, User> {
        println!("批量加载用户ID: {:?}", keys);
        keys.iter()
            .map(|id| {
                (
                    *id,
                    User {
                        id: *id,
                        name: format!("用户 {}", id),
                    },
                )
            })
            .collect()
    }
}

// GraphQL查询定义
struct Query;

#[Object]
impl Query {
    // 查询单个用户
    async fn user(&self, ctx: &Context<'_>, id: i32) -> FieldResult<User> {
        let loader = ctx.data::<Loader<i32, User>>()?;
        Ok(loader.load(id).await)
    }

    // 批量查询多个用户
    async fn users(&self, ctx: &Context<'_>, ids: Vec<i32>) -> FieldResult<Vec<User>> {
        let loader = ctx.data::<Loader<i32, User>>()?;
        Ok(loader.load_many(ids).await)
    }
}

fn main() {
    // 创建数据加载器实例，设置最大批量大小为10
    let loader = Loader::new(UserLoader).with_max_batch_size(10);

    // 构建GraphQL Schema
    let schema = async_graphql::Schema::build(Query, EmptyMutation, EmptySubscription)
        .data(loader)
        .finish();

    // 执行GraphQL查询
    block_on(async {
        // 查询单个用户
        let res = schema.execute("{ user(id: 1) { id name } }").await;
        println!("单用户查询结果: {:?}", res);
        
        // 批量查询多个用户
        let res = schema.execute("{ users(ids: [1,2,3]) { id name } }").await;
        println!("多用户查询结果: {:?}", res);
    });
}

许可证

本项目采用以下任一许可证:

• Apache License, Version 2.0
• MIT license