sokuto_runtime/dag/
pool.rs

1//! # ロックフリーセッション借用プールモジュール
2//!
3//! Tokio 非同期タスクとブロッキング推論スレッド間のセッション借用を
4//! `crossbeam_queue::ArrayQueue` と `tokio::sync::Semaphore` でロックフリーに制御する。
5
6use std::ops::{Deref, DerefMut};
7use std::sync::Arc;
8use std::time::Duration;
9
10use crossbeam_queue::ArrayQueue;
11use tokio::sync::{OwnedSemaphorePermit, Semaphore};
12
13/// プール借用時のエラー型。
14#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
15pub enum PoolError {
16    /// セッション取得のハードタイムアウト。
17    #[error("セッション借用がタイムアウトしました ({0:?})。プールが枯渇しています。")]
18    Timeout(Duration),
19
20    /// セマフォがクローズされた。
21    #[error("セッションプールのセマフォがクローズされました。")]
22    Closed,
23
24    /// プール内部キューの不整合。
25    #[error("セマフォ許可証の取得後にキューが空でした。")]
26    QueueEmpty,
27}
28
29/// ロックフリーなリソース借用プール。
30///
31/// Tokio の非同期スケジューラ上で許可証を獲得後、RAII ガード構造体として
32/// リソースを取り出し、`spawn_blocking` などのスレッド境界を越えて安全に転送可能にする。
33#[derive(Debug, Clone)]
34pub struct LowLatencyResourcePool<T: Send + 'static> {
35    resources: Arc<ArrayQueue<T>>,
36    permits: Arc<Semaphore>,
37    capacity: usize,
38}
39
40impl<T: Send + 'static> LowLatencyResourcePool<T> {
41    /// リソースリストからプールを新規作成する。
42    ///
43    /// # 引数
44    /// - `items`: プールに初期格納するリソース一覧。
45    pub fn new(items: Vec<T>) -> Self {
46        let capacity = items.len().max(1);
47        let queue = Arc::new(ArrayQueue::new(capacity));
48        for item in items {
49            let _ = queue.push(item);
50        }
51        Self {
52            resources: queue,
53            permits: Arc::new(Semaphore::new(capacity)),
54            capacity,
55        }
56    }
57
58    /// プールの最大キャパシティを取得する。
59    pub fn capacity(&self) -> usize {
60        self.capacity
61    }
62
63    /// 現在キュー内で待機中の利用可能リソース数を取得する。
64    pub fn available(&self) -> usize {
65        self.permits.available_permits()
66    }
67
68    /// ハードタイムアウト付きでリソースを非同期借用する。
69    ///
70    /// Tokio 非同期コンテキストでセマフォ許可証を取得後、キューから `pop` する。
71    /// 返却された `PooledResource` は `'static` であり、`spawn_blocking` への転送が可能である。
72    ///
73    /// # 引数
74    /// - `timeout`: 取得待機タイムアウト時間。
75    pub async fn acquire(&self, timeout: Duration) -> Result<PooledResource<T>, PoolError> {
76        let permit_result =
77            tokio::time::timeout(timeout, self.permits.clone().acquire_owned()).await;
78
79        let permit = match permit_result {
80            Ok(Ok(p)) => p,
81            Ok(Err(_)) => return Err(PoolError::Closed),
82            Err(_) => return Err(PoolError::Timeout(timeout)),
83        };
84
85        let resource = self.resources.pop().ok_or(PoolError::QueueEmpty)?;
86
87        Ok(PooledResource {
88            resources: self.resources.clone(),
89            resource: Some(resource),
90            _permit: permit,
91        })
92    }
93}
94
95/// 借用中のリソースを保持する RAII ガード。
96///
97/// スコープを抜けた際に自動で元のプールキューへリソースを返却し、
98/// セマフォ許可証を解放する。
99pub struct PooledResource<T: Send + 'static> {
100    resources: Arc<ArrayQueue<T>>,
101    resource: Option<T>,
102    _permit: OwnedSemaphorePermit,
103}
104
105impl<T: Send + 'static> Deref for PooledResource<T> {
106    type Target = T;
107
108    fn deref(&self) -> &Self::Target {
109        self.resource
110            .as_ref()
111            .expect("リソースは生存期間中常に存在します。")
112    }
113}
114
115impl<T: Send + 'static> DerefMut for PooledResource<T> {
116    fn deref_mut(&mut self) -> &mut Self::Target {
117        self.resource
118            .as_mut()
119            .expect("リソースは生存期間中常に存在します。")
120    }
121}
122
123impl<T: Send + 'static> Drop for PooledResource<T> {
124    fn drop(&mut self) {
125        if let Some(item) = self.resource.take() {
126            let _ = self.resources.push(item);
127        }
128    }
129}
130
131/// ONNX Runtime セッション専用の低遅延借用プール型エイリアス。
132pub type LowLatencySessionPool = LowLatencyResourcePool<ort::session::Session>;
133/// 借用済み ONNX Runtime セッション型エイリアス。
134pub type PooledSession = PooledResource<ort::session::Session>;
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139
140    #[tokio::test]
141    async fn test_resource_pool_acquire_and_release() {
142        let pool = LowLatencyResourcePool::new(vec![100, 200]);
143        assert_eq!(pool.capacity(), 2);
144        assert_eq!(pool.available(), 2);
145
146        {
147            let guard1 = pool
148                .acquire(Duration::from_millis(50))
149                .await
150                .expect("借用成功すること。");
151            assert_eq!(*guard1, 100);
152            assert_eq!(pool.available(), 1);
153
154            let guard2 = pool
155                .acquire(Duration::from_millis(50))
156                .await
157                .expect("借用成功すること。");
158            assert_eq!(*guard2, 200);
159            assert_eq!(pool.available(), 0);
160
161            // タイムアウトテスト (空きがない場合)
162            let timeout_res = pool.acquire(Duration::from_millis(10)).await;
163            assert!(matches!(timeout_res, Err(PoolError::Timeout(_))));
164        }
165
166        // drop 後に返却されていること
167        assert_eq!(pool.available(), 2);
168        let guard_reacquired = pool
169            .acquire(Duration::from_millis(50))
170            .await
171            .expect("再借用成功すること。");
172        assert!(*guard_reacquired == 100 || *guard_reacquired == 200);
173    }
174
175    #[tokio::test]
176    async fn test_transfer_to_spawn_blocking() {
177        let pool = LowLatencyResourcePool::new(vec!["session_a".to_string()]);
178        let guard = pool
179            .acquire(Duration::from_millis(50))
180            .await
181            .expect("借用成功すること。");
182
183        let handle = tokio::task::spawn_blocking(move || {
184            let val = guard.clone();
185            format!("{val}_processed")
186        });
187
188        let res = handle.await.expect("スレッド完了すること。");
189        assert_eq!(res, "session_a_processed");
190        assert_eq!(pool.available(), 1);
191    }
192}
sokuto_runtime/dag/pool.rs

sokuto_runtime/dag/
pool.rs
