wd_pool 0.1.0

A pool crate for rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112

use std::f32::consts::E;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicI32, Ordering};
use std::time::Duration;
use crate::{TaskPool};
use async_channel::{Sender,Receiver};


pub struct TaskPoolChannel {
    status: Arc<AtomicBool>,
    cache_len:Option<usize>,
    idle:i32,
    max:i32,
    idle_sleep_time:Duration,
    active:Arc<AtomicI32>,
    sender:Sender<Pin<Box<dyn Future<Output=()> + Send > >>,
    receiver:Receiver<Pin<Box<dyn Future<Output=()> + Send >>>
}

impl Default for TaskPoolChannel {
    fn default() -> Self {
        let (sender,receiver) = async_channel::unbounded();
        Self{
            status: Arc::new(AtomicBool::new(true)),
            cache_len:None,
            idle: 8, //计算cpu的个数
            idle_sleep_time: Duration::from_secs(1),
            max: i32::MAX,
            active: Arc::new(Default::default()),
            sender,receiver
        }
    }
}

impl TaskPoolChannel{
    pub fn set_cache_len(mut self,len:usize)->Self{
        self.cache_len = Some(len);self
    }
    pub fn set_idle(mut self,idle:i32)->Self{
        self.idle = idle;self
    }
    pub fn set_max_parallel(mut self,max:i32)->Self{
        self.max = max;self
    }
    pub fn set_idle_sleep_time(mut self,idle_sleep_time:Duration)->Self{
        self.idle_sleep_time = idle_sleep_time;self
    }
    pub fn build(mut self) -> impl TaskPool<Out=()> {
        self.init();
        self
    }
}

impl TaskPoolChannel {
    fn init(&mut self){
        if let Some(len) = self.cache_len{
            let (sender,receiver) = async_channel::bounded(len);
            self.sender = sender;
            self.receiver = receiver;
        }
        for _ in 0..self.idle {
            Self::generate_worker(self.idle,self.active.clone(),self.receiver.clone());
        }
    }
    fn generate_worker(idle:i32,active:Arc<AtomicI32>,receiver:Receiver<Pin<Box<dyn Future<Output=()> + Send >>>){
        active.fetch_add(1,Ordering::Relaxed);
        tokio::spawn(async move{
            while !receiver.is_closed() {
                let task = match receiver.try_recv() {
                    Ok(o)=>o,
                    Err(_)=>{
                        if active.load(Ordering::Relaxed) > idle {
                            break
                        }
                        tokio::time::sleep(Duration::from_millis(100)).await;
                        continue
                    }
                };
                task.await;
            }
            active.fetch_sub(1,Ordering::Relaxed);
        });
    }
}


#[async_trait::async_trait]
impl TaskPool for TaskPoolChannel {
    type Out = ();

    async fn push<F: Future<Output=()> + Send + 'static>(&self, task: F)  ->anyhow::Result<Self::Out> {
        if !self.status.load(Ordering::Relaxed) {
            return Err(anyhow::anyhow!("TaskPoolChannel is not normal"))
        }
        if !self.sender.is_closed() {
            let _ = self.sender.send(Box::pin(task)).await;
        }
        Ok(())
    }

    async fn close(&self, timeout: Duration) -> anyhow::Result<()> {
        self.sender.close();
        let active = self.active.clone();
        tokio::time::timeout(timeout,async move{
            while active.load(Ordering::Relaxed) > 0 {
                tokio::time::sleep(Duration::from_secs(1)).await;
            }
        }).await?;Ok(())
    }
}