1use dtmrs_core::{BranchOp, BranchResult};
19use std::collections::HashMap;
20use std::future::Future;
21use std::pin::Pin;
22use std::sync::Arc;
23
24#[derive(Debug, Clone)]
26pub struct BranchCtx {
27 pub gid: String,
28 pub branch_id: String,
29 pub op: BranchOp,
30 pub trans_type: String,
31}
32
33type BoxFut = Pin<Box<dyn Future<Output = BranchResult> + Send>>;
34type Handler = Arc<dyn Fn(BranchCtx) -> BoxFut + Send + Sync>;
35
36#[derive(Debug, Clone, PartialEq, Eq)]
41pub enum Target {
42 Local(String),
44 Http(String),
46 Grpc(GrpcTarget),
48}
49
50#[derive(Debug, Clone, PartialEq, Eq)]
55pub struct GrpcTarget {
56 pub endpoint: String,
58 pub path: String,
60}
61
62pub fn parse_target(s: &str) -> Target {
63 if let Some(name) = s.strip_prefix("local://") {
64 return Target::Local(name.to_string());
65 }
66 if let Some(rest) = s.strip_prefix("grpc://") {
69 if let Some(t) = parse_grpc(rest) {
70 return Target::Grpc(t);
71 }
72 }
73 Target::Http(s.to_string())
74}
75
76fn parse_grpc(rest: &str) -> Option<GrpcTarget> {
81 let (authority, path) = rest.split_once('/')?;
82 if authority.is_empty() {
83 return None;
84 }
85 let segs: Vec<&str> = path.split('/').filter(|s| !s.is_empty()).collect();
86 if segs.len() != 2 {
87 return None;
88 }
89 Some(GrpcTarget {
90 endpoint: format!("http://{authority}"),
91 path: format!("/{}/{}", segs[0], segs[1]),
92 })
93}
94
95#[derive(Default)]
96pub struct Registry {
97 handlers: HashMap<String, Handler>,
98}
99
100impl Registry {
101 pub fn new() -> Self {
102 Self::default()
103 }
104
105 pub fn register<F, Fut>(&mut self, name: &str, f: F) -> &mut Self
107 where
108 F: Fn(BranchCtx) -> Fut + Send + Sync + 'static,
109 Fut: Future<Output = BranchResult> + Send + 'static,
110 {
111 let h: Handler = Arc::new(move |ctx| Box::pin(f(ctx)));
112 self.handlers.insert(name.to_string(), h);
113 self
114 }
115
116 pub fn get(&self, name: &str) -> Option<Handler> {
117 self.handlers.get(name).cloned()
118 }
119
120 pub fn names(&self) -> Vec<&str> {
121 self.handlers.keys().map(String::as_str).collect()
122 }
123
124 pub fn check_all(&self, targets: &[String]) -> Result<(), Vec<String>> {
129 let missing: Vec<String> = targets
130 .iter()
131 .filter_map(|t| match parse_target(t) {
132 Target::Local(n) if !self.handlers.contains_key(&n) => Some(n),
133 _ => None,
134 })
135 .collect();
136 if missing.is_empty() {
137 Ok(())
138 } else {
139 Err(missing)
140 }
141 }
142}
143
144impl std::fmt::Debug for Registry {
145 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
146 f.debug_struct("Registry")
147 .field("handlers", &self.names())
148 .finish()
149 }
150}
151
152#[cfg(test)]
153mod tests {
154 use super::*;
155
156 #[test]
157 fn 前缀区分本地与远端() {
158 assert_eq!(
159 parse_target("local://deduct"),
160 Target::Local("deduct".into())
161 );
162 assert_eq!(
163 parse_target("http://busi/deduct"),
164 Target::Http("http://busi/deduct".into())
165 );
166 assert_eq!(
168 parse_target("https://a/b"),
169 Target::Http("https://a/b".into())
170 );
171 }
172
173 #[test]
174 fn grpc地址拆成端点与方法路径() {
175 let Target::Grpc(t) = parse_target("grpc://127.0.0.1:9000/busi.Busi/Deduct") else {
176 panic!("应该认成 grpc");
177 };
178 assert_eq!(t.endpoint, "http://127.0.0.1:9000");
180 assert_eq!(t.path, "/busi.Busi/Deduct");
181 }
182
183 #[test]
184 fn 畸形grpc地址不猜而是落回http() {
185 for bad in [
188 "grpc://127.0.0.1:9000/onlyservice",
189 "grpc://127.0.0.1:9000/",
190 "grpc://127.0.0.1:9000/a/b/c",
191 "grpc:///a/b",
192 "grpc://noslash",
193 ] {
194 assert!(
195 matches!(parse_target(bad), Target::Http(_)),
196 "{bad} 不该被当成合法 grpc 地址"
197 );
198 }
199 }
200
201 #[tokio::test]
202 async fn 注册与调用() {
203 let mut r = Registry::new();
204 r.register("ok", |_ctx| async { BranchResult::Success });
205 let h = r.get("ok").expect("应该能查到");
206 let ctx = BranchCtx {
207 gid: "g".into(),
208 branch_id: "01".into(),
209 op: BranchOp::Action,
210 trans_type: "saga".into(),
211 };
212 assert_eq!(h(ctx).await, BranchResult::Success);
213 assert!(r.get("nope").is_none());
214 }
215
216 #[test]
217 fn 提交前能查出漏注册的分支() {
218 let mut r = Registry::new();
219 r.register("a", |_| async { BranchResult::Success });
220 let targets = vec![
221 "local://a".to_string(),
222 "local://missing".to_string(),
223 "http://x/y".to_string(),
224 ];
225 let err = r.check_all(&targets).unwrap_err();
226 assert_eq!(err, vec!["missing"], "只报本地漏的,http 不管");
227 assert!(r.check_all(&["local://a".to_string()]).is_ok());
228 }
229}