1use crate::driver::Driver;
44use crate::registry::{BranchCtx, Registry};
45use crate::saga_rows;
46use crate::workflow::{WorkflowCtx, WorkflowRegistry, WorkflowResult};
47use dtmrs_core::{BranchResult, GlobalStatus, SagaStep, TransType};
48use dtmrs_store::Store;
49use std::future::Future;
50use std::sync::Arc;
51use std::time::Duration;
52
53pub struct EmbeddedBuilder {
54 db: String,
55 owner: String,
56 registry: Registry,
57 workflows: WorkflowRegistry,
58 tick: Duration,
59}
60
61impl EmbeddedBuilder {
62 pub fn handler<F, Fut>(mut self, name: &str, f: F) -> Self
64 where
65 F: Fn(BranchCtx) -> Fut + Send + Sync + 'static,
66 Fut: Future<Output = BranchResult> + Send + 'static,
67 {
68 self.registry.register(name, f);
69 self
70 }
71
72 pub fn owner(mut self, o: &str) -> Self {
73 self.owner = o.to_string();
74 self
75 }
76
77 pub fn tick(mut self, d: Duration) -> Self {
79 self.tick = d;
80 self
81 }
82
83 pub fn workflow<F, Fut>(mut self, name: &str, f: F) -> Self
90 where
91 F: Fn(WorkflowCtx) -> Fut + Send + Sync + 'static,
92 Fut: std::future::Future<Output = WorkflowResult<()>> + Send + 'static,
93 {
94 self.workflows.register(name, f);
95 self
96 }
97
98 pub async fn start(self) -> anyhow::Result<Embedded> {
99 let store = Store::open(&self.db).await?;
100 let registry = Arc::new(self.registry);
101 let workflows = Arc::new(self.workflows);
102 let driver = Driver::new(store.clone(), self.owner)
103 .with_registry(registry.clone())
104 .with_workflows(workflows.clone());
105 let task = tokio::spawn(driver.clone().run_forever(self.tick));
107 Ok(Embedded {
108 store,
109 registry,
110 workflows,
111 task: Some(task),
112 })
113 }
114}
115
116pub struct Embedded {
117 store: Store,
118 registry: Arc<Registry>,
119 workflows: Arc<WorkflowRegistry>,
120 task: Option<tokio::task::JoinHandle<()>>,
121}
122
123impl Embedded {
124 pub fn builder(db: &str) -> EmbeddedBuilder {
125 EmbeddedBuilder {
126 db: db.to_string(),
127 owner: format!("embedded-{}", std::process::id()),
128 registry: Registry::new(),
129 workflows: WorkflowRegistry::new(),
130 tick: Duration::from_millis(200),
131 }
132 }
133
134 pub fn saga(&self, gid: &str) -> SagaBuilder<'_> {
135 SagaBuilder {
136 tc: self,
137 gid: gid.to_string(),
138 steps: Vec::new(),
139 }
140 }
141
142 pub async fn submit_workflow(&self, gid: &str, name: &str, input: &str) -> anyhow::Result<()> {
150 if !self.workflows.contains(name) {
151 anyhow::bail!(
152 "workflow「{name}」没注册。已注册的: {:?}",
153 self.workflows.names()
154 );
155 }
156 let mut g = crate::tcc_rows(gid);
157 g.trans_type = TransType::Workflow;
158 g.status = GlobalStatus::Submitted;
159 g.payload = crate::workflow::encode_payload(name, input);
160 self.store.create_global(&g, &[]).await?;
162 Ok(())
163 }
164
165 pub async fn status(&self, gid: &str) -> anyhow::Result<Option<GlobalStatus>> {
166 Ok(self.store.get_global(gid).await?.map(|g| g.status))
167 }
168
169 pub async fn wait_final(&self, gid: &str, timeout: Duration) -> anyhow::Result<GlobalStatus> {
172 let deadline = std::time::Instant::now() + timeout;
173 loop {
174 if let Some(s) = self.status(gid).await? {
175 if s.is_final() {
176 return Ok(s);
177 }
178 }
179 if std::time::Instant::now() >= deadline {
180 anyhow::bail!("等 {gid} 落终态超时");
181 }
182 tokio::time::sleep(Duration::from_millis(20)).await;
183 }
184 }
185
186 pub fn store(&self) -> &Store {
187 &self.store
188 }
189}
190
191impl Drop for Embedded {
192 fn drop(&mut self) {
193 if let Some(t) = self.task.take() {
195 t.abort();
196 }
197 }
198}
199
200pub struct SagaBuilder<'a> {
201 tc: &'a Embedded,
202 gid: String,
203 steps: Vec<SagaStep>,
204}
205
206impl SagaBuilder<'_> {
207 pub fn step(mut self, action: &str, compensate: &str) -> Self {
209 self.steps.push(SagaStep::new(action, compensate));
210 self
211 }
212
213 pub fn step_with(mut self, action: &str, compensate: &str, payload: &str) -> Self {
217 self.steps
218 .push(SagaStep::with_payload(action, compensate, payload));
219 self
220 }
221
222 pub async fn submit(self) -> anyhow::Result<()> {
223 if self.steps.is_empty() {
224 anyhow::bail!("saga 至少要有一步");
225 }
226 let targets: Vec<String> = self
229 .steps
230 .iter()
231 .flat_map(|s| [s.action.clone(), s.compensate.clone()])
232 .collect();
233 if let Err(missing) = self.tc.registry.check_all(&targets) {
234 anyhow::bail!("这些本地分支没注册: {}", missing.join(", "));
235 }
236 let (g, branches) = saga_rows(&self.gid, &self.steps);
237 self.tc.store.create_global(&g, &branches).await?;
238 Ok(())
239 }
240}