use crate::driver::Driver;
use crate::registry::{BranchCtx, Registry};
use crate::saga_rows;
use crate::workflow::{WorkflowCtx, WorkflowRegistry, WorkflowResult};
use dtmrs_core::{BranchResult, GlobalStatus, SagaStep, TransType};
use dtmrs_store::Store;
use std::future::Future;
use std::sync::Arc;
use std::time::Duration;
pub struct EmbeddedBuilder {
db: String,
owner: String,
registry: Registry,
workflows: WorkflowRegistry,
tick: Duration,
}
impl EmbeddedBuilder {
pub fn handler<F, Fut>(mut self, name: &str, f: F) -> Self
where
F: Fn(BranchCtx) -> Fut + Send + Sync + 'static,
Fut: Future<Output = BranchResult> + Send + 'static,
{
self.registry.register(name, f);
self
}
pub fn owner(mut self, o: &str) -> Self {
self.owner = o.to_string();
self
}
pub fn tick(mut self, d: Duration) -> Self {
self.tick = d;
self
}
pub fn workflow<F, Fut>(mut self, name: &str, f: F) -> Self
where
F: Fn(WorkflowCtx) -> Fut + Send + Sync + 'static,
Fut: std::future::Future<Output = WorkflowResult<()>> + Send + 'static,
{
self.workflows.register(name, f);
self
}
pub async fn start(self) -> anyhow::Result<Embedded> {
let store = Store::open(&self.db).await?;
let registry = Arc::new(self.registry);
let workflows = Arc::new(self.workflows);
let driver = Driver::new(store.clone(), self.owner)
.with_registry(registry.clone())
.with_workflows(workflows.clone());
let task = tokio::spawn(driver.clone().run_forever(self.tick));
Ok(Embedded {
store,
registry,
workflows,
task: Some(task),
})
}
}
pub struct Embedded {
store: Store,
registry: Arc<Registry>,
workflows: Arc<WorkflowRegistry>,
task: Option<tokio::task::JoinHandle<()>>,
}
impl Embedded {
pub fn builder(db: &str) -> EmbeddedBuilder {
EmbeddedBuilder {
db: db.to_string(),
owner: format!("embedded-{}", std::process::id()),
registry: Registry::new(),
workflows: WorkflowRegistry::new(),
tick: Duration::from_millis(200),
}
}
pub fn saga(&self, gid: &str) -> SagaBuilder<'_> {
SagaBuilder {
tc: self,
gid: gid.to_string(),
steps: Vec::new(),
}
}
pub async fn submit_workflow(&self, gid: &str, name: &str, input: &str) -> anyhow::Result<()> {
if !self.workflows.contains(name) {
anyhow::bail!(
"workflow「{name}」没注册。已注册的: {:?}",
self.workflows.names()
);
}
let mut g = crate::tcc_rows(gid);
g.trans_type = TransType::Workflow;
g.status = GlobalStatus::Submitted;
g.payload = crate::workflow::encode_payload(name, input);
self.store.create_global(&g, &[]).await?;
Ok(())
}
pub async fn status(&self, gid: &str) -> anyhow::Result<Option<GlobalStatus>> {
Ok(self.store.get_global(gid).await?.map(|g| g.status))
}
pub async fn wait_final(&self, gid: &str, timeout: Duration) -> anyhow::Result<GlobalStatus> {
let deadline = std::time::Instant::now() + timeout;
loop {
if let Some(s) = self.status(gid).await? {
if s.is_final() {
return Ok(s);
}
}
if std::time::Instant::now() >= deadline {
anyhow::bail!("等 {gid} 落终态超时");
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
}
pub fn store(&self) -> &Store {
&self.store
}
}
impl Drop for Embedded {
fn drop(&mut self) {
if let Some(t) = self.task.take() {
t.abort();
}
}
}
pub struct SagaBuilder<'a> {
tc: &'a Embedded,
gid: String,
steps: Vec<SagaStep>,
}
impl SagaBuilder<'_> {
pub fn step(mut self, action: &str, compensate: &str) -> Self {
self.steps.push(SagaStep::new(action, compensate));
self
}
pub fn step_with(mut self, action: &str, compensate: &str, payload: &str) -> Self {
self.steps
.push(SagaStep::with_payload(action, compensate, payload));
self
}
pub async fn submit(self) -> anyhow::Result<()> {
if self.steps.is_empty() {
anyhow::bail!("saga 至少要有一步");
}
let targets: Vec<String> = self
.steps
.iter()
.flat_map(|s| [s.action.clone(), s.compensate.clone()])
.collect();
if let Err(missing) = self.tc.registry.check_all(&targets) {
anyhow::bail!("这些本地分支没注册: {}", missing.join(", "));
}
let (g, branches) = saga_rows(&self.gid, &self.steps);
self.tc.store.create_global(&g, &branches).await?;
Ok(())
}
}