1use std::collections::HashMap;
14use std::error::Error;
15use std::ops::Deref;
16use std::sync::Arc;
17
18use async_lock::{Mutex, RwLock};
19use futures::Future;
20use futures::future::{BoxFuture, LocalBoxFuture, join_all};
21use prost::Message;
22use serde::{Deserialize, Serialize};
23use ts_rs::TS;
24
25use crate::proto::request::ClientReq;
26use crate::proto::response::ClientResp;
27use crate::proto::{
28 ColumnType, GetFeaturesReq, GetFeaturesResp, GetHostedTablesReq, GetHostedTablesResp,
29 HostedTable, JoinType, MakeJoinTableReq, MakeTableReq, RemoveHostedTablesUpdateReq, Request,
30 Response, ServerError, ServerSystemInfoReq,
31};
32use crate::table::{JoinOptions, Table, TableInitOptions, TableOptions};
33use crate::table_data::{TableData, UpdateData};
34use crate::table_ref::TableRef;
35use crate::utils::*;
36use crate::view::{OnUpdateData, ViewWindow};
37use crate::{OnUpdateMode, OnUpdateOptions, asyncfn, clone};
38
39#[derive(Clone, Debug, Serialize, Deserialize, TS)]
41pub struct SystemInfo<T = u64> {
42 pub heap_size: T,
44
45 pub used_size: T,
47
48 pub cpu_time: u32,
53
54 pub cpu_time_epoch: u32,
56
57 pub timestamp: Option<T>,
61
62 pub client_heap: Option<T>,
65
66 pub client_used: Option<T>,
69}
70
71impl<U: Copy + 'static> SystemInfo<U> {
72 pub fn cast<T: Copy + 'static>(&self) -> SystemInfo<T>
75 where
76 U: num_traits::AsPrimitive<T>,
77 {
78 SystemInfo {
79 heap_size: self.heap_size.as_(),
80 used_size: self.used_size.as_(),
81 cpu_time: self.cpu_time,
82 cpu_time_epoch: self.cpu_time_epoch,
83 timestamp: self.timestamp.map(|x| x.as_()),
84 client_heap: self.client_heap.map(|x| x.as_()),
85 client_used: self.client_used.map(|x| x.as_()),
86 }
87 }
88}
89
90#[derive(Clone, Debug, Default, PartialEq)]
93pub struct Features(Arc<GetFeaturesResp>);
94
95impl Features {
96 pub fn get_group_rollup_modes(&self) -> Vec<crate::config::GroupRollupMode> {
97 self.group_rollup_mode
98 .iter()
99 .map(|x| {
100 crate::config::GroupRollupMode::from(
101 crate::proto::GroupRollupMode::try_from(*x).unwrap(),
102 )
103 })
104 .collect::<Vec<_>>()
105 }
106}
107
108impl Deref for Features {
109 type Target = GetFeaturesResp;
110
111 fn deref(&self) -> &Self::Target {
112 &self.0
113 }
114}
115
116impl GetFeaturesResp {
117 pub fn default_op(&self, col_type: ColumnType) -> Option<&str> {
118 self.filter_ops
119 .get(&(col_type as u32))?
120 .options
121 .first()
122 .map(|x| x.as_str())
123 }
124
125 pub fn get_window_aggregates(
128 &self,
129 col_type: ColumnType,
130 ) -> Vec<crate::config::WindowAggregate> {
131 self.window_aggregates
132 .get(&(col_type as u32))
133 .map(|x| {
134 x.options
135 .iter()
136 .filter_map(|x| crate::proto::WindowAggregate::try_from(*x).ok())
137 .map(|x| x.into())
138 .collect()
139 })
140 .unwrap_or_default()
141 }
142
143 pub fn has_window_aggregates(&self) -> bool {
146 self.window_aggregates
147 .values()
148 .any(|x| !x.options.is_empty())
149 }
150}
151
152type BoxFn<I, O> = Box<dyn Fn(I) -> O + Send + Sync + 'static>;
153type Box2Fn<I, J, O> = Box<dyn Fn(I, J) -> O + Send + Sync + 'static>;
154
155type Subscriptions<C> = Arc<RwLock<HashMap<u32, C>>>;
156type OnErrorCallback =
157 Box2Fn<ClientError, Option<ReconnectCallback>, BoxFuture<'static, Result<(), ClientError>>>;
158
159type OnceCallback = Box<dyn FnOnce(Response) -> ClientResult<()> + Send + Sync + 'static>;
160type SendCallback = Arc<
161 dyn for<'a> Fn(&'a Request) -> BoxFuture<'a, Result<(), Box<dyn Error + Send + Sync>>>
162 + Send
163 + Sync
164 + 'static,
165>;
166
167pub trait ClientHandler: Clone + Send + Sync + 'static {
169 fn send_request(
170 &self,
171 msg: Vec<u8>,
172 ) -> impl Future<Output = Result<(), Box<dyn Error + Send + Sync>>> + Send;
173}
174
175mod name_registry {
176 use std::collections::HashSet;
177 use std::sync::{Arc, LazyLock, Mutex};
178
179 use crate::ClientError;
180 use crate::view::ClientResult;
181
182 static CLIENT_ID_GEN: LazyLock<Arc<Mutex<u32>>> = LazyLock::new(Arc::default);
183 static REGISTERED_CLIENTS: LazyLock<Arc<Mutex<HashSet<String>>>> = LazyLock::new(Arc::default);
184
185 pub(crate) fn generate_name(name: Option<&str>) -> ClientResult<String> {
186 if let Some(name) = name {
187 if let Some(name) = REGISTERED_CLIENTS
188 .lock()
189 .map_err(ClientError::from)?
190 .get(name)
191 {
192 Err(ClientError::DuplicateNameError(name.to_owned()))
193 } else {
194 Ok(name.to_owned())
195 }
196 } else {
197 let mut guard = CLIENT_ID_GEN.lock()?;
198 *guard += 1;
199 Ok(format!("client-{guard}"))
200 }
201 }
202}
203
204#[derive(Clone)]
211#[allow(clippy::type_complexity)]
212pub struct ReconnectCallback(
213 Arc<dyn Fn() -> LocalBoxFuture<'static, Result<(), Box<dyn Error>>> + Send + Sync>,
214);
215
216impl Deref for ReconnectCallback {
217 type Target = dyn Fn() -> LocalBoxFuture<'static, Result<(), Box<dyn Error>>> + Send + Sync;
218
219 fn deref(&self) -> &Self::Target {
220 &*self.0
221 }
222}
223
224impl ReconnectCallback {
225 pub fn new(
226 f: impl Fn() -> LocalBoxFuture<'static, Result<(), Box<dyn Error>>> + Send + Sync + 'static,
227 ) -> Self {
228 ReconnectCallback(Arc::new(f))
229 }
230}
231
232#[derive(Clone)]
236pub struct Client {
237 name: Arc<String>,
238 features: Arc<Mutex<Option<Features>>>,
239 send: SendCallback,
240 id_gen: IDGen,
241 subscriptions_errors: Subscriptions<OnErrorCallback>,
242 subscriptions_once: Subscriptions<OnceCallback>,
243 subscriptions: Subscriptions<BoxFn<Response, BoxFuture<'static, Result<(), ClientError>>>>,
244}
245
246impl PartialEq for Client {
247 fn eq(&self, other: &Self) -> bool {
248 self.name == other.name
249 }
250}
251
252impl std::fmt::Debug for Client {
253 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
254 f.debug_struct("Client").finish()
255 }
256}
257
258impl Client {
259 pub fn new_with_callback<T, U>(name: Option<&str>, send_request: T) -> ClientResult<Self>
262 where
263 T: Fn(Vec<u8>) -> U + 'static + Sync + Send,
264 U: Future<Output = Result<(), Box<dyn Error + Send + Sync>>> + Send + 'static,
265 {
266 let name = name_registry::generate_name(name)?;
267 let send_request = Arc::new(send_request);
268 let send: SendCallback = Arc::new(move |req| {
269 let mut bytes: Vec<u8> = Vec::new();
270 req.encode(&mut bytes).unwrap();
271 let send_request = send_request.clone();
272 Box::pin(async move { send_request(bytes).await })
273 });
274
275 Ok(Client {
276 name: Arc::new(name),
277 features: Arc::default(),
278 id_gen: IDGen::default(),
279 send,
280 subscriptions: Subscriptions::default(),
281 subscriptions_errors: Arc::default(),
282 subscriptions_once: Arc::default(),
283 })
284 }
285
286 pub fn new<T>(name: Option<&str>, client_handler: T) -> ClientResult<Self>
288 where
289 T: ClientHandler + 'static + Sync + Send,
290 {
291 Self::new_with_callback(
292 name,
293 asyncfn!(client_handler, async move |req| {
294 client_handler.send_request(req).await
295 }),
296 )
297 }
298
299 pub fn get_name(&self) -> &'_ str {
300 self.name.as_str()
301 }
302
303 pub async fn handle_response<'a>(&'a self, msg: &'a [u8]) -> ClientResult<bool> {
310 let msg = Response::decode(msg)?;
311 tracing::debug!("RECV {}", msg);
312 let mut wr = self.subscriptions_once.write().await;
313 if let Some(handler) = (*wr).remove(&msg.msg_id) {
314 drop(wr);
315 handler(msg)?;
316 return Ok(true);
317 } else if let Some(handler) = self.subscriptions.try_read().unwrap().get(&msg.msg_id) {
318 drop(wr);
319 handler(msg).await?;
320 return Ok(true);
321 }
322
323 if let Response {
324 client_resp: Some(ClientResp::ServerError(ServerError { message, .. })),
325 ..
326 } = &msg
327 {
328 tracing::error!("{}", message);
329 } else {
330 tracing::debug!("Received unsolicited server response: {}", msg);
331 }
332
333 Ok(false)
334 }
335
336 pub async fn handle_error<T, U>(
338 &self,
339 message: ClientError,
340 reconnect: Option<T>,
341 ) -> ClientResult<()>
342 where
343 T: Fn() -> U + Clone + Send + Sync + 'static,
344 U: Future<Output = ClientResult<()>>,
345 {
346 let subs = self.subscriptions_errors.read().await;
347 let tasks = join_all(subs.values().map(|callback| {
348 callback(
349 message.clone(),
350 reconnect.clone().map(move |f| {
351 ReconnectCallback(Arc::new(move || {
352 clone!(f);
353 Box::pin(async move { Ok(f().await?) }) as LocalBoxFuture<'static, _>
354 }))
355 }),
356 )
357 }));
358
359 tasks.await.into_iter().collect::<Result<(), _>>()?;
360 self.close_and_error_subscriptions(&message).await
361 }
362
363 async fn close_and_error_subscriptions(&self, message: &ClientError) -> ClientResult<()> {
368 let synthetic_error = |msg_id| Response {
369 msg_id,
370 entity_id: "".to_string(),
371 client_resp: Some(ClientResp::ServerError(ServerError {
372 message: format!("{message}"),
373 status_code: 2,
374 })),
375 };
376
377 self.subscriptions.write().await.clear();
378 let callbacks_once = self
379 .subscriptions_once
380 .write()
381 .await
382 .drain()
383 .collect::<Vec<_>>();
384
385 callbacks_once
386 .into_iter()
387 .try_for_each(|(msg_id, f)| f(synthetic_error(msg_id)))
388 }
389
390 pub async fn on_error<T, U, V>(&self, on_error: T) -> ClientResult<u32>
391 where
392 T: Fn(ClientError, Option<ReconnectCallback>) -> U + Clone + Send + Sync + 'static,
393 U: Future<Output = V> + Send + 'static,
394 V: Into<Result<(), ClientError>> + Sync + 'static,
395 {
396 let id = self.gen_id();
397 let callback = asyncfn!(on_error, async move |x, y| on_error(x, y).await.into());
398 self.subscriptions_errors
399 .write()
400 .await
401 .insert(id, Box::new(move |x, y| Box::pin(callback(x, y))));
402
403 Ok(id)
404 }
405
406 pub(crate) fn gen_id(&self) -> u32 {
408 self.id_gen.next()
409 }
410
411 pub(crate) async fn unsubscribe(&self, update_id: u32) -> ClientResult<()> {
412 let callback = self
413 .subscriptions
414 .write()
415 .await
416 .remove(&update_id)
417 .ok_or(ClientError::Unknown("remove_update".to_string()))?;
418
419 drop(callback);
420 Ok(())
421 }
422
423 pub(crate) async fn subscribe_once(
425 &self,
426 msg: &Request,
427 on_update: Box<dyn FnOnce(Response) -> ClientResult<()> + Send + Sync + 'static>,
428 ) -> ClientResult<()> {
429 self.subscriptions_once
430 .write()
431 .await
432 .insert(msg.msg_id, on_update);
433
434 tracing::debug!("SEND {}", msg);
435 if let Err(e) = (self.send)(msg).await {
436 self.subscriptions_once.write().await.remove(&msg.msg_id);
437 Err(ClientError::Unknown(e.to_string()))
438 } else {
439 Ok(())
440 }
441 }
442
443 pub(crate) async fn subscribe<T, U>(&self, msg: &Request, on_update: T) -> ClientResult<()>
444 where
445 T: Fn(Response) -> U + Send + Sync + 'static,
446 U: Future<Output = Result<(), ClientError>> + Send + 'static,
447 {
448 self.subscriptions
449 .write()
450 .await
451 .insert(msg.msg_id, Box::new(move |x| Box::pin(on_update(x))));
452
453 tracing::debug!("SEND {}", msg);
454 if let Err(e) = (self.send)(msg).await {
455 self.subscriptions.write().await.remove(&msg.msg_id);
456 Err(ClientError::Unknown(e.to_string()))
457 } else {
458 Ok(())
459 }
460 }
461
462 pub(crate) async fn oneshot(&self, req: &Request) -> ClientResult<ClientResp> {
465 let (sender, receiver) = futures::channel::oneshot::channel::<ClientResp>();
466 let on_update = Box::new(move |res: Response| {
467 sender.send(res.client_resp.unwrap()).map_err(|x| x.into())
468 });
469
470 self.subscribe_once(req, on_update).await?;
471 receiver
472 .await
473 .map_err(|_| ClientError::Unknown(format!("Internal error for req {req}")))
474 }
475
476 pub(crate) async fn get_features(&self) -> ClientResult<Features> {
477 let mut guard = self.features.lock().await;
478 let features = if let Some(features) = &*guard {
479 features.clone()
480 } else {
481 let msg = Request {
482 msg_id: self.gen_id(),
483 entity_id: "".to_owned(),
484 client_req: Some(ClientReq::GetFeaturesReq(GetFeaturesReq {})),
485 };
486
487 let features = Features(Arc::new(match self.oneshot(&msg).await? {
488 ClientResp::GetFeaturesResp(features) => Ok(features),
489 resp => Err(resp),
490 }?));
491
492 *guard = Some(features.clone());
493 features
494 };
495
496 Ok(features)
497 }
498
499 pub async fn table(&self, input: TableData, options: TableInitOptions) -> ClientResult<Table> {
555 let entity_id = match options.name.clone() {
556 Some(x) => x.to_owned(),
557 None => randid(),
558 };
559
560 if let TableData::View(view) = &input {
561 let window = ViewWindow::default();
562 let arrow = view.to_arrow(window).await?;
563 let mut table = self
564 .crate_table_inner(UpdateData::Arrow(arrow).into(), options.into(), entity_id)
565 .await?;
566
567 let table_ = table.clone();
568 let callback = asyncfn!(table_, update, async move |update: OnUpdateData| {
569 let update = UpdateData::Arrow(update.delta.expect("Malformed message").into());
570 let options = crate::UpdateOptions::default();
571 table_.update(update, options).await.unwrap_or_log();
572 });
573
574 let options = OnUpdateOptions {
575 mode: Some(OnUpdateMode::Row),
576 };
577
578 let on_update_token = view.on_update(callback, options).await?;
579 table.view_update_token = Some(on_update_token);
580 Ok(table)
581 } else {
582 self.crate_table_inner(input, options.into(), entity_id)
583 .await
584 }
585 }
586
587 async fn crate_table_inner(
588 &self,
589 input: TableData,
590 options: TableOptions,
591 entity_id: String,
592 ) -> ClientResult<Table> {
593 let msg = Request {
594 msg_id: self.gen_id(),
595 entity_id: entity_id.clone(),
596 client_req: Some(ClientReq::MakeTableReq(MakeTableReq {
597 data: Some(input.into()),
598 options: Some(options.clone().try_into()?),
599 })),
600 };
601
602 let client = self.clone();
603 match self.oneshot(&msg).await? {
604 ClientResp::MakeTableResp(_) => Ok(Table::new(entity_id, client, options)),
605 resp => Err(resp.into()),
606 }
607 }
608
609 pub async fn join(
621 &self,
622 left: TableRef,
623 right: TableRef,
624 on: &str,
625 options: JoinOptions,
626 ) -> ClientResult<Table> {
627 let entity_id = options.name.unwrap_or_else(randid);
628 let join_type: JoinType = options.join_type.unwrap_or_default();
629 let right_on_column = options.right_on.unwrap_or_default();
630 let msg = Request {
631 msg_id: self.gen_id(),
632 entity_id: entity_id.clone(),
633 client_req: Some(ClientReq::MakeJoinTableReq(MakeJoinTableReq {
634 left_table_id: left.table_name().to_owned(),
635 right_table_id: right.table_name().to_owned(),
636 on_column: on.to_owned(),
637 join_type: join_type.into(),
638 right_on_column,
639 })),
640 };
641
642 let client = self.clone();
643 match self.oneshot(&msg).await? {
644 ClientResp::MakeJoinTableResp(_) => Ok(Table::new(entity_id, client, TableOptions {
645 index: Some(on.to_owned()),
646 limit: None,
647 page_to_disk: None,
648 })),
649 resp => Err(resp.into()),
650 }
651 }
652
653 async fn get_table_infos(&self) -> ClientResult<Vec<HostedTable>> {
654 let msg = Request {
655 msg_id: self.gen_id(),
656 entity_id: "".to_owned(),
657 client_req: Some(ClientReq::GetHostedTablesReq(GetHostedTablesReq {
658 subscribe: false,
659 })),
660 };
661
662 match self.oneshot(&msg).await? {
663 ClientResp::GetHostedTablesResp(GetHostedTablesResp { table_infos }) => Ok(table_infos),
664 resp => Err(resp.into()),
665 }
666 }
667
668 pub async fn open_table(&self, entity_id: String) -> ClientResult<Table> {
685 let infos = self.get_table_infos().await?;
686
687 if let Some(info) = infos.into_iter().find(|i| i.entity_id == entity_id) {
689 let options = TableOptions {
690 index: info.index,
691 limit: info.limit,
692 page_to_disk: None,
695 };
696
697 let client = self.clone();
698 Ok(Table::new(entity_id, client, options))
699 } else {
700 Err(ClientError::Unknown(format!(
701 "Unknown table \"{}\"",
702 entity_id
703 )))
704 }
705 }
706
707 pub async fn get_hosted_table_names(&self) -> ClientResult<Vec<String>> {
724 let msg = Request {
725 msg_id: self.gen_id(),
726 entity_id: "".to_owned(),
727 client_req: Some(ClientReq::GetHostedTablesReq(GetHostedTablesReq {
728 subscribe: false,
729 })),
730 };
731
732 match self.oneshot(&msg).await? {
733 ClientResp::GetHostedTablesResp(GetHostedTablesResp { table_infos }) => {
734 Ok(table_infos.into_iter().map(|i| i.entity_id).collect())
735 },
736 resp => Err(resp.into()),
737 }
738 }
739
740 pub async fn on_hosted_tables_update<T, U>(&self, on_update: T) -> ClientResult<u32>
744 where
745 T: Fn() -> U + Send + Sync + 'static,
746 U: Future<Output = ()> + Send + 'static,
747 {
748 let on_update = Arc::new(on_update);
749 let callback = asyncfn!(on_update, async move |resp: Response| {
750 match resp.client_resp {
751 Some(ClientResp::GetHostedTablesResp(_)) | None => {
752 on_update().await;
753 Ok(())
754 },
755 resp => Err(resp.into()),
756 }
757 });
758
759 let msg = Request {
760 msg_id: self.gen_id(),
761 entity_id: "".to_owned(),
762 client_req: Some(ClientReq::GetHostedTablesReq(GetHostedTablesReq {
763 subscribe: true,
764 })),
765 };
766
767 self.subscribe(&msg, callback).await?;
768 Ok(msg.msg_id)
769 }
770
771 pub async fn remove_hosted_tables_update(&self, update_id: u32) -> ClientResult<()> {
774 let msg = Request {
775 msg_id: self.gen_id(),
776 entity_id: "".to_owned(),
777 client_req: Some(ClientReq::RemoveHostedTablesUpdateReq(
778 RemoveHostedTablesUpdateReq { id: update_id },
779 )),
780 };
781
782 self.unsubscribe(update_id).await?;
783 match self.oneshot(&msg).await? {
784 ClientResp::RemoveHostedTablesUpdateResp(_) => Ok(()),
785 resp => Err(resp.into()),
786 }
787 }
788
789 pub async fn system_info(&self) -> ClientResult<SystemInfo> {
793 let msg = Request {
794 msg_id: self.gen_id(),
795 entity_id: "".to_string(),
796 client_req: Some(ClientReq::ServerSystemInfoReq(ServerSystemInfoReq {})),
797 };
798
799 match self.oneshot(&msg).await? {
800 ClientResp::ServerSystemInfoResp(resp) => {
801 #[cfg(not(target_family = "wasm"))]
802 let timestamp = Some(
803 std::time::SystemTime::now()
804 .duration_since(std::time::UNIX_EPOCH)?
805 .as_millis() as u64,
806 );
807
808 #[cfg(target_family = "wasm")]
809 let timestamp = None;
810
811 #[cfg(feature = "talc-allocator")]
812 let (client_used, client_heap) = {
813 let (client_used, client_heap) = crate::utils::get_used();
814 (Some(client_used as u64), Some(client_heap as u64))
815 };
816
817 #[cfg(not(feature = "talc-allocator"))]
818 let (client_used, client_heap) = (None, None);
819
820 let info = SystemInfo {
821 heap_size: resp.heap_size,
822 used_size: resp.used_size,
823 cpu_time: resp.cpu_time,
824 cpu_time_epoch: resp.cpu_time_epoch,
825 timestamp,
826 client_heap,
827 client_used,
828 };
829
830 Ok(info)
831 },
832 resp => Err(resp.into()),
833 }
834 }
835}