feat: pool spacetime client connections
This commit is contained in:
@@ -31,6 +31,7 @@ GENARRATIVE_API_TARGET="http://127.0.0.1:3100"
|
|||||||
GENARRATIVE_INTERNAL_API_SECRET="CHANGE_ME_FOR_PRODUCTION"
|
GENARRATIVE_INTERNAL_API_SECRET="CHANGE_ME_FOR_PRODUCTION"
|
||||||
GENARRATIVE_SPACETIME_SERVER_URL="http://127.0.0.1:3001"
|
GENARRATIVE_SPACETIME_SERVER_URL="http://127.0.0.1:3001"
|
||||||
GENARRATIVE_SPACETIME_DATABASE="genarrative-dev"
|
GENARRATIVE_SPACETIME_DATABASE="genarrative-dev"
|
||||||
|
GENARRATIVE_SPACETIME_POOL_SIZE="4"
|
||||||
|
|
||||||
# Local Caddy upstream target used for dist-based testing.
|
# Local Caddy upstream target used for dist-based testing.
|
||||||
CADDY_API_UPSTREAM="http://127.0.0.1:8081"
|
CADDY_API_UPSTREAM="http://127.0.0.1:8081"
|
||||||
|
|||||||
@@ -0,0 +1,141 @@
|
|||||||
|
# api-server SpacetimeClient 连接池化设计 2026-04-23
|
||||||
|
|
||||||
|
更新时间:`2026-04-23`
|
||||||
|
|
||||||
|
## 1. 背景
|
||||||
|
|
||||||
|
当前 `api-server` 虽然在 `AppState` 中只持有一个 `SpacetimeClient` 实例,但 `spacetime-client` 内部仍然是:
|
||||||
|
|
||||||
|
1. 每次 procedure / reducer 调用都执行一次 `DbConnection::builder().build()`
|
||||||
|
2. 建连后立即 `run_threaded()`
|
||||||
|
3. 拿到结果后立刻 `disconnect()`
|
||||||
|
|
||||||
|
也就是说,当前问题不是 `api-server` 每次请求都 new 一个 client,而是:
|
||||||
|
|
||||||
|
**每次 client 调用都新建并销毁一条 SpacetimeDB 连接。**
|
||||||
|
|
||||||
|
## 2. 本轮目标
|
||||||
|
|
||||||
|
本轮不继续维持“每次 HTTP 请求一条短连接”的阶段性策略。
|
||||||
|
|
||||||
|
本轮目标改为:
|
||||||
|
|
||||||
|
1. `api-server` 进程内预热并持有一组可复用的 SpacetimeDB 连接
|
||||||
|
2. 每次 HTTP 请求只从池里借一个可用连接执行 procedure / reducer
|
||||||
|
3. 请求完成后归还连接,不主动断开
|
||||||
|
4. 连接失效时自动剔除并按需重建
|
||||||
|
|
||||||
|
## 3. 为什么不直接引第三方池库
|
||||||
|
|
||||||
|
当前仓库使用的是 `spacetimedb-sdk` 生成的 `DbConnection`,不是传统 SQL client。
|
||||||
|
|
||||||
|
它的连接模型包含:
|
||||||
|
|
||||||
|
1. `on_connect`
|
||||||
|
2. `on_disconnect`
|
||||||
|
3. `run_threaded`
|
||||||
|
4. reducer / procedure callback
|
||||||
|
|
||||||
|
这类对象不是标准的 `bb8` / `deadpool` 资源接口。
|
||||||
|
|
||||||
|
当前仓库也没有已经接入的通用资源池库,因此本轮优先在 `spacetime-client` 内实现最小可控池化层,而不是强行套第三方 SQL 风格池库。
|
||||||
|
|
||||||
|
## 4. 池化设计
|
||||||
|
|
||||||
|
## 4.1 结构
|
||||||
|
|
||||||
|
`SpacetimeClient` 内新增一个共享池状态:
|
||||||
|
|
||||||
|
1. `pool_size`
|
||||||
|
2. `Semaphore`
|
||||||
|
3. `Vec<Mutex<Option<PooledConnection>>>`
|
||||||
|
|
||||||
|
其中 `PooledConnection` 持有:
|
||||||
|
|
||||||
|
1. `DbConnection`
|
||||||
|
2. `run_threaded` 返回的后台线程句柄
|
||||||
|
3. 连接唯一 id
|
||||||
|
|
||||||
|
## 4.2 借还模型
|
||||||
|
|
||||||
|
每次调用 procedure / reducer 时:
|
||||||
|
|
||||||
|
1. 先获取 `Semaphore permit`
|
||||||
|
2. 选取一个空闲槽位
|
||||||
|
3. 若槽位已有健康连接,则直接复用
|
||||||
|
4. 若槽位为空或连接已坏,则现场重建
|
||||||
|
5. 调用完成后归还槽位,但不主动断开连接
|
||||||
|
|
||||||
|
## 4.3 健康判断
|
||||||
|
|
||||||
|
当前阶段不做复杂心跳表。
|
||||||
|
|
||||||
|
最小健康策略如下:
|
||||||
|
|
||||||
|
1. procedure / reducer callback 正常完成:连接保持在池中
|
||||||
|
2. 连接在调用期间触发 `on_disconnect`:标记该槽位失效
|
||||||
|
3. 下次借用该槽位时重建连接
|
||||||
|
|
||||||
|
## 4.4 并发策略
|
||||||
|
|
||||||
|
不共享同一个 `DbConnection` 给多个并发请求同时发 procedure。
|
||||||
|
|
||||||
|
原因:
|
||||||
|
|
||||||
|
1. SDK callback 是异步回调模型
|
||||||
|
2. 当前仓库调用层没有 request id 级别的统一 dispatcher
|
||||||
|
3. 多请求共用一条连接容易把回调和调用方绑定关系搞乱
|
||||||
|
|
||||||
|
所以本轮采取:
|
||||||
|
|
||||||
|
**一个池槽位同一时刻只服务一个请求。**
|
||||||
|
|
||||||
|
这本质上是“连接池”,不是“多路复用单连接”。
|
||||||
|
|
||||||
|
## 5. 默认规模
|
||||||
|
|
||||||
|
默认池大小取小值,避免本地开发和轻量部署浪费连接:
|
||||||
|
|
||||||
|
1. 默认 `4`
|
||||||
|
2. 允许通过环境变量覆盖,例如 `GENARRATIVE_SPACETIME_POOL_SIZE`
|
||||||
|
|
||||||
|
## 6. 错误与超时策略
|
||||||
|
|
||||||
|
沿用现有 `SpacetimeClientError` 口径:
|
||||||
|
|
||||||
|
1. 建连失败:`Build` / `Runtime`
|
||||||
|
2. 连接在返回前断开:`ConnectDropped` 或 `Procedure`
|
||||||
|
3. 调用超时:`Timeout`
|
||||||
|
|
||||||
|
新增规则:
|
||||||
|
|
||||||
|
1. 借用池槽位超时,也映射为 `Timeout`
|
||||||
|
2. 某槽位一旦确认断线,必须在池中清空,不能继续复用脏连接
|
||||||
|
3. procedure / reducer 等待结果无论成功、失败还是超时,都必须先归还租约再向上层返回,避免槽位泄漏把池卡死
|
||||||
|
4. 调用期间若连接先收到 `on_disconnect`,当前阶段只标记坏连接;若业务回调未及时返回,则最终由调用超时路径统一清槽并回传错误
|
||||||
|
|
||||||
|
## 7. 与现有文档的关系
|
||||||
|
|
||||||
|
之前 [`AXUM_TO_SPACETIMEDB_ASSET_OBJECT_CONFIRM_CALL_DESIGN_2026-04-21.md`](D:/Genarrative/docs/technical/AXUM_TO_SPACETIMEDB_ASSET_OBJECT_CONFIRM_CALL_DESIGN_2026-04-21.md)
|
||||||
|
中写明“当前阶段每次 HTTP 请求可以建立一条短连接,待真实链路验证稳定后再评估连接池或长连接复用”。
|
||||||
|
|
||||||
|
本轮就是进入这个“下一阶段”:
|
||||||
|
|
||||||
|
1. 保留 `on_connect` 后再发请求的约束
|
||||||
|
2. 去掉“请求完成立即断开”的短连接策略
|
||||||
|
3. 改成 `spacetime-client` 进程内连接池
|
||||||
|
|
||||||
|
## 8. 验收标准
|
||||||
|
|
||||||
|
落地后至少满足:
|
||||||
|
|
||||||
|
1. `api-server` 启动后,`SpacetimeClient` 不再为每次调用单独建连
|
||||||
|
2. 同一进程内连续多个 API 请求可以复用池中连接
|
||||||
|
3. 单个连接断开后不会污染后续请求
|
||||||
|
4. `api-server` 调用侧无需修改业务 handler
|
||||||
|
|
||||||
|
## 9. 一句话结论
|
||||||
|
|
||||||
|
本轮不引第三方 SQL 风格池库,而是在 `spacetime-client` 内实现一层:
|
||||||
|
|
||||||
|
**面向 `DbConnection` 的最小连接池,让 `api-server` 复用长活连接,而不是每次调用都单独建连。**
|
||||||
@@ -67,6 +67,7 @@ pub struct AppConfig {
|
|||||||
pub spacetime_server_url: String,
|
pub spacetime_server_url: String,
|
||||||
pub spacetime_database: String,
|
pub spacetime_database: String,
|
||||||
pub spacetime_token: Option<String>,
|
pub spacetime_token: Option<String>,
|
||||||
|
pub spacetime_pool_size: u32,
|
||||||
pub llm_provider: LlmProvider,
|
pub llm_provider: LlmProvider,
|
||||||
pub llm_base_url: String,
|
pub llm_base_url: String,
|
||||||
pub llm_api_key: Option<String>,
|
pub llm_api_key: Option<String>,
|
||||||
@@ -139,6 +140,7 @@ impl Default for AppConfig {
|
|||||||
spacetime_server_url: "http://127.0.0.1:3000".to_string(),
|
spacetime_server_url: "http://127.0.0.1:3000".to_string(),
|
||||||
spacetime_database: "genarrative-dev".to_string(),
|
spacetime_database: "genarrative-dev".to_string(),
|
||||||
spacetime_token: None,
|
spacetime_token: None,
|
||||||
|
spacetime_pool_size: 4,
|
||||||
llm_provider: LlmProvider::Ark,
|
llm_provider: LlmProvider::Ark,
|
||||||
llm_base_url: DEFAULT_ARK_BASE_URL.to_string(),
|
llm_base_url: DEFAULT_ARK_BASE_URL.to_string(),
|
||||||
llm_api_key: None,
|
llm_api_key: None,
|
||||||
@@ -370,6 +372,11 @@ impl AppConfig {
|
|||||||
}
|
}
|
||||||
|
|
||||||
config.spacetime_token = read_first_non_empty_env(&["GENARRATIVE_SPACETIME_TOKEN"]);
|
config.spacetime_token = read_first_non_empty_env(&["GENARRATIVE_SPACETIME_TOKEN"]);
|
||||||
|
if let Some(spacetime_pool_size) =
|
||||||
|
read_first_positive_u32_env(&["GENARRATIVE_SPACETIME_POOL_SIZE"])
|
||||||
|
{
|
||||||
|
config.spacetime_pool_size = spacetime_pool_size;
|
||||||
|
}
|
||||||
|
|
||||||
if let Some(llm_provider) =
|
if let Some(llm_provider) =
|
||||||
read_first_llm_provider_env(&["GENARRATIVE_LLM_PROVIDER", "LLM_PROVIDER"])
|
read_first_llm_provider_env(&["GENARRATIVE_LLM_PROVIDER", "LLM_PROVIDER"])
|
||||||
@@ -620,3 +627,51 @@ fn parse_positive_u16(raw: &str) -> Option<u16> {
|
|||||||
|
|
||||||
Some(value)
|
Some(value)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::AppConfig;
|
||||||
|
use std::sync::{Mutex, OnceLock};
|
||||||
|
|
||||||
|
static ENV_LOCK: OnceLock<Mutex<()>> = OnceLock::new();
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn from_env_reads_spacetime_pool_size() {
|
||||||
|
let _guard = ENV_LOCK
|
||||||
|
.get_or_init(|| Mutex::new(()))
|
||||||
|
.lock()
|
||||||
|
.expect("env lock should not poison");
|
||||||
|
|
||||||
|
unsafe {
|
||||||
|
std::env::remove_var("GENARRATIVE_SPACETIME_POOL_SIZE");
|
||||||
|
std::env::set_var("GENARRATIVE_SPACETIME_POOL_SIZE", "7");
|
||||||
|
}
|
||||||
|
|
||||||
|
let config = AppConfig::from_env();
|
||||||
|
assert_eq!(config.spacetime_pool_size, 7);
|
||||||
|
|
||||||
|
unsafe {
|
||||||
|
std::env::remove_var("GENARRATIVE_SPACETIME_POOL_SIZE");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn from_env_ignores_zero_spacetime_pool_size() {
|
||||||
|
let _guard = ENV_LOCK
|
||||||
|
.get_or_init(|| Mutex::new(()))
|
||||||
|
.lock()
|
||||||
|
.expect("env lock should not poison");
|
||||||
|
|
||||||
|
unsafe {
|
||||||
|
std::env::remove_var("GENARRATIVE_SPACETIME_POOL_SIZE");
|
||||||
|
std::env::set_var("GENARRATIVE_SPACETIME_POOL_SIZE", "0");
|
||||||
|
}
|
||||||
|
|
||||||
|
let config = AppConfig::from_env();
|
||||||
|
assert_eq!(config.spacetime_pool_size, 4);
|
||||||
|
|
||||||
|
unsafe {
|
||||||
|
std::env::remove_var("GENARRATIVE_SPACETIME_POOL_SIZE");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -116,6 +116,7 @@ impl AppState {
|
|||||||
server_url: config.spacetime_server_url.clone(),
|
server_url: config.spacetime_server_url.clone(),
|
||||||
database: config.spacetime_database.clone(),
|
database: config.spacetime_database.clone(),
|
||||||
token: config.spacetime_token.clone(),
|
token: config.spacetime_token.clone(),
|
||||||
|
pool_size: config.spacetime_pool_size,
|
||||||
});
|
});
|
||||||
let llm_client = build_llm_client(&config)?;
|
let llm_client = build_llm_client(&config)?;
|
||||||
|
|
||||||
|
|||||||
@@ -3,7 +3,9 @@ pub mod module_bindings;
|
|||||||
use std::{
|
use std::{
|
||||||
error::Error,
|
error::Error,
|
||||||
fmt,
|
fmt,
|
||||||
|
sync::atomic::{AtomicBool, Ordering},
|
||||||
sync::{Arc, Mutex},
|
sync::{Arc, Mutex},
|
||||||
|
thread::JoinHandle,
|
||||||
time::Duration,
|
time::Duration,
|
||||||
};
|
};
|
||||||
|
|
||||||
@@ -95,7 +97,10 @@ use module_story::{
|
|||||||
};
|
};
|
||||||
use shared_kernel::format_timestamp_micros;
|
use shared_kernel::format_timestamp_micros;
|
||||||
use spacetimedb_sdk::DbContext;
|
use spacetimedb_sdk::DbContext;
|
||||||
use tokio::{sync::oneshot, time::timeout};
|
use tokio::{
|
||||||
|
sync::{OwnedSemaphorePermit, Semaphore, oneshot},
|
||||||
|
time::timeout,
|
||||||
|
};
|
||||||
|
|
||||||
use crate::module_bindings::{
|
use crate::module_bindings::{
|
||||||
AiResultReferenceInput as BindingAiResultReferenceInput,
|
AiResultReferenceInput as BindingAiResultReferenceInput,
|
||||||
@@ -350,11 +355,13 @@ pub struct SpacetimeClientConfig {
|
|||||||
pub server_url: String,
|
pub server_url: String,
|
||||||
pub database: String,
|
pub database: String,
|
||||||
pub token: Option<String>,
|
pub token: Option<String>,
|
||||||
|
pub pool_size: u32,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Clone, Debug)]
|
#[derive(Clone)]
|
||||||
pub struct SpacetimeClient {
|
pub struct SpacetimeClient {
|
||||||
config: SpacetimeClientConfig,
|
config: SpacetimeClientConfig,
|
||||||
|
pool: Arc<SpacetimeConnectionPool>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug)]
|
#[derive(Debug)]
|
||||||
@@ -372,9 +379,45 @@ type ProcedureResultSender<T> =
|
|||||||
Arc<Mutex<Option<oneshot::Sender<Result<T, SpacetimeClientError>>>>>;
|
Arc<Mutex<Option<oneshot::Sender<Result<T, SpacetimeClientError>>>>>;
|
||||||
type ReducerResultSender = Arc<Mutex<Option<oneshot::Sender<Result<(), SpacetimeClientError>>>>>;
|
type ReducerResultSender = Arc<Mutex<Option<oneshot::Sender<Result<(), SpacetimeClientError>>>>>;
|
||||||
|
|
||||||
|
struct SpacetimeConnectionPool {
|
||||||
|
slots: Vec<tokio::sync::Mutex<PooledConnectionSlot>>,
|
||||||
|
permits: Arc<Semaphore>,
|
||||||
|
}
|
||||||
|
|
||||||
|
struct PooledConnectionSlot {
|
||||||
|
connection: Option<PooledConnection>,
|
||||||
|
in_use: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
struct PooledConnection {
|
||||||
|
connection: DbConnection,
|
||||||
|
runner: Option<JoinHandle<()>>,
|
||||||
|
broken: Arc<AtomicBool>,
|
||||||
|
}
|
||||||
|
|
||||||
|
struct PooledConnectionLease {
|
||||||
|
slot_index: usize,
|
||||||
|
connection: Option<PooledConnection>,
|
||||||
|
_permit: OwnedSemaphorePermit,
|
||||||
|
}
|
||||||
|
|
||||||
impl SpacetimeClient {
|
impl SpacetimeClient {
|
||||||
pub fn new(config: SpacetimeClientConfig) -> Self {
|
pub fn new(config: SpacetimeClientConfig) -> Self {
|
||||||
Self { config }
|
let pool_size = config.pool_size.max(1) as usize;
|
||||||
|
let slots = (0..pool_size)
|
||||||
|
.map(|_| {
|
||||||
|
tokio::sync::Mutex::new(PooledConnectionSlot {
|
||||||
|
connection: None,
|
||||||
|
in_use: false,
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
let pool = Arc::new(SpacetimeConnectionPool {
|
||||||
|
slots,
|
||||||
|
permits: Arc::new(Semaphore::new(pool_size)),
|
||||||
|
});
|
||||||
|
|
||||||
|
Self { config, pool }
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn create_ai_task(
|
pub async fn create_ai_task(
|
||||||
@@ -2213,70 +2256,122 @@ impl SpacetimeClient {
|
|||||||
where
|
where
|
||||||
T: Send + 'static,
|
T: Send + 'static,
|
||||||
{
|
{
|
||||||
let config = self.config.clone();
|
|
||||||
let (sender, receiver) = oneshot::channel();
|
let (sender, receiver) = oneshot::channel();
|
||||||
let result_sender = Arc::new(Mutex::new(Some(sender)));
|
let result_sender = Arc::new(Mutex::new(Some(sender)));
|
||||||
let connect_sender = result_sender.clone();
|
let lease = self.acquire_connection().await?;
|
||||||
let disconnect_sender = result_sender.clone();
|
let final_result = if let Some(connection) = lease.connection.as_ref() {
|
||||||
|
call(&connection.connection, result_sender.clone());
|
||||||
|
match timeout(CONFIRM_ASSET_OBJECT_TIMEOUT, receiver).await {
|
||||||
|
Ok(inner) => match inner {
|
||||||
|
Ok(value) => value,
|
||||||
|
Err(_) => Err(SpacetimeClientError::ConnectDropped),
|
||||||
|
},
|
||||||
|
Err(_) => Err(Self::resolve_timeout_error(Some(connection))),
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
Err(SpacetimeClientError::Runtime(
|
||||||
|
"SpacetimeDB 连接租约缺少连接".to_string(),
|
||||||
|
))
|
||||||
|
};
|
||||||
|
self.release_connection(lease).await;
|
||||||
|
|
||||||
let connection = tokio::task::spawn_blocking(move || {
|
final_result
|
||||||
DbConnection::builder()
|
|
||||||
.with_uri(config.server_url)
|
|
||||||
.with_database_name(config.database)
|
|
||||||
.with_token(config.token)
|
|
||||||
.on_connect(move |connection, _, _| {
|
|
||||||
// SDK 收到 IdentityToken 后才调用 procedure,避免 WebSocket 已建好但身份握手未完成时丢请求。
|
|
||||||
call(connection, connect_sender);
|
|
||||||
})
|
|
||||||
.on_disconnect(move |_, error| {
|
|
||||||
let message = error
|
|
||||||
.map(|error| error.to_string())
|
|
||||||
.unwrap_or_else(|| "SpacetimeDB 连接在 procedure 返回前断开".to_string());
|
|
||||||
send_once(
|
|
||||||
&disconnect_sender,
|
|
||||||
Err(SpacetimeClientError::Procedure(message)),
|
|
||||||
);
|
|
||||||
})
|
|
||||||
.build()
|
|
||||||
.map_err(|error| SpacetimeClientError::Build(error.to_string()))
|
|
||||||
})
|
|
||||||
.await
|
|
||||||
.map_err(|error| SpacetimeClientError::Runtime(error.to_string()))??;
|
|
||||||
|
|
||||||
let runner = connection.run_threaded();
|
|
||||||
let result = timeout(CONFIRM_ASSET_OBJECT_TIMEOUT, receiver).await;
|
|
||||||
let _ = connection.disconnect();
|
|
||||||
// SDK 线程会在断开消息被处理后自行退出;HTTP 请求不能同步等待该线程,否则 Windows 本地联调可能卡在收尾阶段。
|
|
||||||
drop(runner);
|
|
||||||
|
|
||||||
result
|
|
||||||
.map_err(|_| SpacetimeClientError::Timeout)?
|
|
||||||
.map_err(|_| SpacetimeClientError::ConnectDropped)?
|
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn call_reducer_after_connect(
|
async fn call_reducer_after_connect(
|
||||||
&self,
|
&self,
|
||||||
call: impl FnOnce(&DbConnection, ReducerResultSender) + Send + 'static,
|
call: impl FnOnce(&DbConnection, ReducerResultSender) + Send + 'static,
|
||||||
) -> Result<(), SpacetimeClientError> {
|
) -> Result<(), SpacetimeClientError> {
|
||||||
let config = self.config.clone();
|
|
||||||
let (sender, receiver) = oneshot::channel();
|
let (sender, receiver) = oneshot::channel();
|
||||||
let result_sender = Arc::new(Mutex::new(Some(sender)));
|
let result_sender = Arc::new(Mutex::new(Some(sender)));
|
||||||
let connect_sender = result_sender.clone();
|
let lease = self.acquire_connection().await?;
|
||||||
let disconnect_sender = result_sender.clone();
|
let final_result = if let Some(connection) = lease.connection.as_ref() {
|
||||||
|
call(&connection.connection, result_sender.clone());
|
||||||
|
match timeout(CONFIRM_ASSET_OBJECT_TIMEOUT, receiver).await {
|
||||||
|
Ok(inner) => match inner {
|
||||||
|
Ok(value) => value,
|
||||||
|
Err(_) => Err(SpacetimeClientError::ConnectDropped),
|
||||||
|
},
|
||||||
|
Err(_) => Err(Self::resolve_timeout_error(Some(connection))),
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
Err(SpacetimeClientError::Runtime(
|
||||||
|
"SpacetimeDB 连接租约缺少连接".to_string(),
|
||||||
|
))
|
||||||
|
};
|
||||||
|
self.release_connection(lease).await;
|
||||||
|
|
||||||
|
final_result
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn acquire_connection(&self) -> Result<PooledConnectionLease, SpacetimeClientError> {
|
||||||
|
let permit = timeout(
|
||||||
|
CONFIRM_ASSET_OBJECT_TIMEOUT,
|
||||||
|
self.pool.permits.clone().acquire_owned(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.map_err(|_| SpacetimeClientError::Timeout)?
|
||||||
|
.map_err(|error| SpacetimeClientError::Runtime(error.to_string()))?;
|
||||||
|
|
||||||
|
loop {
|
||||||
|
for (slot_index, slot) in self.pool.slots.iter().enumerate() {
|
||||||
|
if let Ok(mut slot_guard) = slot.try_lock() {
|
||||||
|
if slot_guard.in_use {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let reusable_connection = slot_guard
|
||||||
|
.connection
|
||||||
|
.take()
|
||||||
|
.filter(|connection| !connection.is_broken());
|
||||||
|
slot_guard.in_use = true;
|
||||||
|
drop(slot_guard);
|
||||||
|
|
||||||
|
let connection = if let Some(connection) = reusable_connection {
|
||||||
|
connection
|
||||||
|
} else {
|
||||||
|
match self.build_pooled_connection().await {
|
||||||
|
Ok(connection) => connection,
|
||||||
|
Err(error) => {
|
||||||
|
let mut slot_guard = self.pool.slots[slot_index].lock().await;
|
||||||
|
slot_guard.in_use = false;
|
||||||
|
return Err(error);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
return Ok(PooledConnectionLease {
|
||||||
|
slot_index,
|
||||||
|
connection: Some(connection),
|
||||||
|
_permit: permit,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
tokio::task::yield_now().await;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn build_pooled_connection(&self) -> Result<PooledConnection, SpacetimeClientError> {
|
||||||
|
let config = self.config.clone();
|
||||||
|
let broken = Arc::new(AtomicBool::new(false));
|
||||||
|
let (sender, receiver) = oneshot::channel::<Result<(), SpacetimeClientError>>();
|
||||||
|
let connect_sender = Arc::new(Mutex::new(Some(sender)));
|
||||||
|
let broken_flag = broken.clone();
|
||||||
|
let disconnect_sender = connect_sender.clone();
|
||||||
let connection = tokio::task::spawn_blocking(move || {
|
let connection = tokio::task::spawn_blocking(move || {
|
||||||
DbConnection::builder()
|
DbConnection::builder()
|
||||||
.with_uri(config.server_url)
|
.with_uri(config.server_url)
|
||||||
.with_database_name(config.database)
|
.with_database_name(config.database)
|
||||||
.with_token(config.token)
|
.with_token(config.token)
|
||||||
.on_connect(move |connection, _, _| {
|
.on_connect(move |_, _, _| {
|
||||||
call(connection, connect_sender);
|
send_connect_once(&connect_sender, Ok(()));
|
||||||
})
|
})
|
||||||
.on_disconnect(move |_, error| {
|
.on_disconnect(move |_, error| {
|
||||||
|
broken_flag.store(true, Ordering::SeqCst);
|
||||||
let message = error
|
let message = error
|
||||||
.map(|error| error.to_string())
|
.map(|error| error.to_string())
|
||||||
.unwrap_or_else(|| "SpacetimeDB 连接在 reducer 返回前断开".to_string());
|
.unwrap_or_else(|| "SpacetimeDB 连接已断开".to_string());
|
||||||
send_reducer_once(
|
send_connect_once(
|
||||||
&disconnect_sender,
|
&disconnect_sender,
|
||||||
Err(SpacetimeClientError::Procedure(message)),
|
Err(SpacetimeClientError::Procedure(message)),
|
||||||
);
|
);
|
||||||
@@ -2288,13 +2383,70 @@ impl SpacetimeClient {
|
|||||||
.map_err(|error| SpacetimeClientError::Runtime(error.to_string()))??;
|
.map_err(|error| SpacetimeClientError::Runtime(error.to_string()))??;
|
||||||
|
|
||||||
let runner = connection.run_threaded();
|
let runner = connection.run_threaded();
|
||||||
let result = timeout(CONFIRM_ASSET_OBJECT_TIMEOUT, receiver).await;
|
timeout(CONFIRM_ASSET_OBJECT_TIMEOUT, receiver)
|
||||||
let _ = connection.disconnect();
|
.await
|
||||||
drop(runner);
|
|
||||||
|
|
||||||
result
|
|
||||||
.map_err(|_| SpacetimeClientError::Timeout)?
|
.map_err(|_| SpacetimeClientError::Timeout)?
|
||||||
.map_err(|_| SpacetimeClientError::ConnectDropped)?
|
.map_err(|_| SpacetimeClientError::ConnectDropped)??;
|
||||||
|
|
||||||
|
Ok(PooledConnection {
|
||||||
|
connection,
|
||||||
|
runner: Some(runner),
|
||||||
|
broken,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn release_connection(&self, mut lease: PooledConnectionLease) {
|
||||||
|
let mut slot_guard = self.pool.slots[lease.slot_index].lock().await;
|
||||||
|
slot_guard.in_use = false;
|
||||||
|
let Some(connection) = lease.connection.take() else {
|
||||||
|
slot_guard.connection = None;
|
||||||
|
return;
|
||||||
|
};
|
||||||
|
if connection.is_broken() {
|
||||||
|
slot_guard.connection = None;
|
||||||
|
} else {
|
||||||
|
slot_guard.connection = Some(connection);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 超时后必须统一归还租约;若连接已先一步断开则回传断线,否则标记坏连接并回传超时。
|
||||||
|
fn resolve_timeout_error(connection: Option<&PooledConnection>) -> SpacetimeClientError {
|
||||||
|
if let Some(connection) = connection {
|
||||||
|
if connection.is_broken() {
|
||||||
|
return SpacetimeClientError::ConnectDropped;
|
||||||
|
}
|
||||||
|
connection.mark_broken();
|
||||||
|
}
|
||||||
|
|
||||||
|
SpacetimeClientError::Timeout
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl PooledConnection {
|
||||||
|
fn is_broken(&self) -> bool {
|
||||||
|
self.broken.load(Ordering::SeqCst)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn mark_broken(&self) {
|
||||||
|
self.broken.store(true, Ordering::SeqCst);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Drop for PooledConnection {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
let _ = self.connection.disconnect();
|
||||||
|
if let Some(runner) = self.runner.take() {
|
||||||
|
drop(runner);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl fmt::Debug for SpacetimeClient {
|
||||||
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||||
|
f.debug_struct("SpacetimeClient")
|
||||||
|
.field("config", &self.config)
|
||||||
|
.field("pool_size", &self.pool.slots.len())
|
||||||
|
.finish()
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -2318,6 +2470,19 @@ fn send_reducer_once(sender: &ReducerResultSender, result: Result<(), SpacetimeC
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn send_connect_once(
|
||||||
|
sender: &Arc<Mutex<Option<oneshot::Sender<Result<(), SpacetimeClientError>>>>>,
|
||||||
|
result: Result<(), SpacetimeClientError>,
|
||||||
|
) {
|
||||||
|
if let Some(sender) = sender
|
||||||
|
.lock()
|
||||||
|
.expect("spacetime connect sender should not poison")
|
||||||
|
.take()
|
||||||
|
{
|
||||||
|
let _ = sender.send(result);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn map_entity_binding_input(
|
fn map_entity_binding_input(
|
||||||
input: module_assets::AssetEntityBindingInput,
|
input: module_assets::AssetEntityBindingInput,
|
||||||
) -> BindingAssetEntityBindingInput {
|
) -> BindingAssetEntityBindingInput {
|
||||||
|
|||||||
Reference in New Issue
Block a user