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, ViewBinding};
33use crate::table_data::{TableData, UpdateData};
34use crate::table_ref::TableRef;
35use crate::utils::*;
36use crate::view::{OnRemoveData, 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 pub fn get_split_rollup_modes(&self) -> Vec<crate::config::SplitRollupMode> {
112 if self.split_rollup_mode.is_empty() {
113 return vec![crate::config::SplitRollupMode::Flat];
114 }
115
116 self.split_rollup_mode
117 .iter()
118 .map(|x| {
119 crate::config::SplitRollupMode::from(
120 crate::proto::SplitRollupMode::try_from(*x).unwrap(),
121 )
122 })
123 .collect::<Vec<_>>()
124 }
125}
126
127impl Deref for Features {
128 type Target = GetFeaturesResp;
129
130 fn deref(&self) -> &Self::Target {
131 &self.0
132 }
133}
134
135impl GetFeaturesResp {
136 pub fn default_op(&self, col_type: ColumnType) -> Option<&str> {
137 self.filter_ops
138 .get(&(col_type as u32))?
139 .options
140 .first()
141 .map(|x| x.as_str())
142 }
143
144 pub fn get_window_aggregates(
147 &self,
148 col_type: ColumnType,
149 ) -> Vec<crate::proto::WindowAggregateArgs> {
150 self.window_aggregates
151 .get(&(col_type as u32))
152 .map(|x| x.options.clone())
153 .unwrap_or_default()
154 }
155
156 pub fn has_window_aggregates(&self) -> bool {
159 self.window_aggregates
160 .values()
161 .any(|x| !x.options.is_empty())
162 }
163}
164
165type Box2Fn<I, J, O> = Box<dyn Fn(I, J) -> O + Send + Sync + 'static>;
166
167type Subscriptions<C> = Arc<RwLock<HashMap<u32, C>>>;
168type UpdateCallback =
169 Arc<dyn Fn(Response) -> BoxFuture<'static, Result<(), ClientError>> + Send + Sync + 'static>;
170type OnErrorCallback =
171 Box2Fn<ClientError, Option<ReconnectCallback>, BoxFuture<'static, Result<(), ClientError>>>;
172
173type OnceCallback = Box<dyn FnOnce(Response) -> ClientResult<()> + Send + Sync + 'static>;
174type SendCallback = Arc<
175 dyn for<'a> Fn(&'a Request) -> BoxFuture<'a, Result<(), Box<dyn Error + Send + Sync>>>
176 + Send
177 + Sync
178 + 'static,
179>;
180
181pub trait ClientHandler: Clone + Send + Sync + 'static {
183 fn send_request(
184 &self,
185 msg: Vec<u8>,
186 ) -> impl Future<Output = Result<(), Box<dyn Error + Send + Sync>>> + Send;
187}
188
189mod name_registry {
190 use std::collections::HashSet;
191 use std::sync::{Arc, LazyLock, Mutex};
192
193 use crate::ClientError;
194 use crate::view::ClientResult;
195
196 static CLIENT_ID_GEN: LazyLock<Arc<Mutex<u32>>> = LazyLock::new(Arc::default);
197 static REGISTERED_CLIENTS: LazyLock<Arc<Mutex<HashSet<String>>>> = LazyLock::new(Arc::default);
198
199 pub(crate) fn generate_name(name: Option<&str>) -> ClientResult<String> {
200 if let Some(name) = name {
201 if let Some(name) = REGISTERED_CLIENTS
202 .lock()
203 .map_err(ClientError::from)?
204 .get(name)
205 {
206 Err(ClientError::DuplicateNameError(name.to_owned()))
207 } else {
208 Ok(name.to_owned())
209 }
210 } else {
211 let mut guard = CLIENT_ID_GEN.lock()?;
212 *guard += 1;
213 Ok(format!("client-{guard}"))
214 }
215 }
216}
217
218#[derive(Clone)]
225#[allow(clippy::type_complexity)]
226pub struct ReconnectCallback(
227 Arc<dyn Fn() -> LocalBoxFuture<'static, Result<(), Box<dyn Error>>> + Send + Sync>,
228);
229
230impl Deref for ReconnectCallback {
231 type Target = dyn Fn() -> LocalBoxFuture<'static, Result<(), Box<dyn Error>>> + Send + Sync;
232
233 fn deref(&self) -> &Self::Target {
234 &*self.0
235 }
236}
237
238impl ReconnectCallback {
239 pub fn new(
240 f: impl Fn() -> LocalBoxFuture<'static, Result<(), Box<dyn Error>>> + Send + Sync + 'static,
241 ) -> Self {
242 ReconnectCallback(Arc::new(f))
243 }
244}
245
246#[derive(Clone)]
250pub struct Client {
251 name: Arc<String>,
252 features: Arc<Mutex<Option<Features>>>,
253 send: SendCallback,
254 id_gen: IDGen,
255 subscriptions_errors: Subscriptions<OnErrorCallback>,
256 subscriptions_once: Subscriptions<OnceCallback>,
257 subscriptions: Subscriptions<UpdateCallback>,
258}
259
260impl PartialEq for Client {
261 fn eq(&self, other: &Self) -> bool {
262 self.name == other.name
263 }
264}
265
266impl std::fmt::Debug for Client {
267 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
268 f.debug_struct("Client").finish()
269 }
270}
271
272impl Client {
273 pub fn new_with_callback<T, U>(name: Option<&str>, send_request: T) -> ClientResult<Self>
276 where
277 T: Fn(Vec<u8>) -> U + 'static + Sync + Send,
278 U: Future<Output = Result<(), Box<dyn Error + Send + Sync>>> + Send + 'static,
279 {
280 let name = name_registry::generate_name(name)?;
281 let send_request = Arc::new(send_request);
282 let send: SendCallback = Arc::new(move |req| {
283 let mut bytes: Vec<u8> = Vec::new();
284 req.encode(&mut bytes).unwrap();
285 let send_request = send_request.clone();
286 Box::pin(async move { send_request(bytes).await })
287 });
288
289 Ok(Client {
290 name: Arc::new(name),
291 features: Arc::default(),
292 id_gen: IDGen::default(),
293 send,
294 subscriptions: Subscriptions::default(),
295 subscriptions_errors: Arc::default(),
296 subscriptions_once: Arc::default(),
297 })
298 }
299
300 pub fn new<T>(name: Option<&str>, client_handler: T) -> ClientResult<Self>
302 where
303 T: ClientHandler + 'static + Sync + Send,
304 {
305 Self::new_with_callback(
306 name,
307 asyncfn!(client_handler, async move |req| {
308 client_handler.send_request(req).await
309 }),
310 )
311 }
312
313 pub fn get_name(&self) -> &'_ str {
314 self.name.as_str()
315 }
316
317 pub async fn handle_response<'a>(&'a self, msg: &'a [u8]) -> ClientResult<bool> {
324 let msg = Response::decode(msg)?;
325 tracing::debug!("RECV {}", msg);
326 let mut wr = self.subscriptions_once.write().await;
327 if let Some(handler) = (*wr).remove(&msg.msg_id) {
328 drop(wr);
329 handler(msg)?;
330 return Ok(true);
331 }
332
333 let handler = self.subscriptions.read().await.get(&msg.msg_id).cloned();
334 drop(wr);
335 if let Some(handler) = handler {
336 handler(msg).await?;
337 return Ok(true);
338 }
339
340 if let Response {
341 client_resp: Some(ClientResp::ServerError(ServerError { message, .. })),
342 ..
343 } = &msg
344 {
345 tracing::error!("{}", message);
346 } else {
347 tracing::debug!("Received unsolicited server response: {}", msg);
348 }
349
350 Ok(false)
351 }
352
353 pub async fn handle_error<T, U>(
355 &self,
356 message: ClientError,
357 reconnect: Option<T>,
358 ) -> ClientResult<()>
359 where
360 T: Fn() -> U + Clone + Send + Sync + 'static,
361 U: Future<Output = ClientResult<()>>,
362 {
363 let subs = self.subscriptions_errors.read().await;
364 let tasks = join_all(subs.values().map(|callback| {
365 callback(
366 message.clone(),
367 reconnect.clone().map(move |f| {
368 ReconnectCallback(Arc::new(move || {
369 clone!(f);
370 Box::pin(async move { Ok(f().await?) }) as LocalBoxFuture<'static, _>
371 }))
372 }),
373 )
374 }));
375
376 tasks.await.into_iter().collect::<Result<(), _>>()?;
377 self.close_and_error_subscriptions(&message).await
378 }
379
380 async fn close_and_error_subscriptions(&self, message: &ClientError) -> ClientResult<()> {
385 let synthetic_error = |msg_id| Response {
386 msg_id,
387 entity_id: "".to_string(),
388 client_resp: Some(ClientResp::ServerError(ServerError {
389 message: format!("{message}"),
390 status_code: 2,
391 })),
392 };
393
394 self.subscriptions.write().await.clear();
395 let callbacks_once = self
396 .subscriptions_once
397 .write()
398 .await
399 .drain()
400 .collect::<Vec<_>>();
401
402 callbacks_once
403 .into_iter()
404 .try_for_each(|(msg_id, f)| f(synthetic_error(msg_id)))
405 }
406
407 pub async fn on_error<T, U, V>(&self, on_error: T) -> ClientResult<u32>
408 where
409 T: Fn(ClientError, Option<ReconnectCallback>) -> U + Clone + Send + Sync + 'static,
410 U: Future<Output = V> + Send + 'static,
411 V: Into<Result<(), ClientError>> + Sync + 'static,
412 {
413 let id = self.gen_id();
414 let callback = asyncfn!(on_error, async move |x, y| on_error(x, y).await.into());
415 self.subscriptions_errors
416 .write()
417 .await
418 .insert(id, Box::new(move |x, y| Box::pin(callback(x, y))));
419
420 Ok(id)
421 }
422
423 pub(crate) fn gen_id(&self) -> u32 {
425 self.id_gen.next()
426 }
427
428 pub(crate) async fn unsubscribe(&self, update_id: u32) -> ClientResult<()> {
429 let callback = self
430 .subscriptions
431 .write()
432 .await
433 .remove(&update_id)
434 .ok_or(ClientError::Unknown("remove_update".to_string()))?;
435
436 drop(callback);
437 Ok(())
438 }
439
440 pub(crate) async fn subscribe_once(
442 &self,
443 msg: &Request,
444 on_update: Box<dyn FnOnce(Response) -> ClientResult<()> + Send + Sync + 'static>,
445 ) -> ClientResult<()> {
446 self.subscriptions_once
447 .write()
448 .await
449 .insert(msg.msg_id, on_update);
450
451 tracing::debug!("SEND {}", msg);
452 if let Err(e) = (self.send)(msg).await {
453 self.subscriptions_once.write().await.remove(&msg.msg_id);
454 Err(ClientError::Unknown(e.to_string()))
455 } else {
456 Ok(())
457 }
458 }
459
460 pub(crate) async fn subscribe<T, U>(&self, msg: &Request, on_update: T) -> ClientResult<()>
461 where
462 T: Fn(Response) -> U + Send + Sync + 'static,
463 U: Future<Output = Result<(), ClientError>> + Send + 'static,
464 {
465 self.subscriptions
466 .write()
467 .await
468 .insert(msg.msg_id, Arc::new(move |x| Box::pin(on_update(x))));
469
470 tracing::debug!("SEND {}", msg);
471 if let Err(e) = (self.send)(msg).await {
472 self.subscriptions.write().await.remove(&msg.msg_id);
473 Err(ClientError::Unknown(e.to_string()))
474 } else {
475 Ok(())
476 }
477 }
478
479 pub(crate) async fn oneshot(&self, req: &Request) -> ClientResult<ClientResp> {
482 let (sender, receiver) = futures::channel::oneshot::channel::<ClientResp>();
483 let on_update = Box::new(move |res: Response| {
484 sender.send(res.client_resp.unwrap()).map_err(|x| x.into())
485 });
486
487 self.subscribe_once(req, on_update).await?;
488 receiver
489 .await
490 .map_err(|_| ClientError::Unknown(format!("Internal error for req {req}")))
491 }
492
493 pub(crate) async fn get_features(&self) -> ClientResult<Features> {
494 let mut guard = self.features.lock().await;
495 let features = if let Some(features) = &*guard {
496 features.clone()
497 } else {
498 let msg = Request {
499 msg_id: self.gen_id(),
500 entity_id: "".to_owned(),
501 client_req: Some(ClientReq::GetFeaturesReq(GetFeaturesReq {})),
502 };
503
504 let features = Features(Arc::new(match self.oneshot(&msg).await? {
505 ClientResp::GetFeaturesResp(features) => Ok(features),
506 resp => Err(resp),
507 }?));
508
509 *guard = Some(features.clone());
510 features
511 };
512
513 Ok(features)
514 }
515
516 pub async fn table(&self, input: TableData, options: TableInitOptions) -> ClientResult<Table> {
572 let entity_id = match options.name.clone() {
573 Some(x) => x.to_owned(),
574 None => randid(),
575 };
576
577 if let TableData::View(view) = &input {
578 let mut options = options;
579 let source_index = view.source.as_ref().and_then(|x| x.options.index.clone());
580 if let (None, Some(index)) = (&options.index, &source_index) {
581 let config = view.get_config().await?;
582 let is_flat = config.group_by.is_empty() && config.split_by.is_empty();
583 let has_index = config.columns.iter().flatten().any(|x| x == index);
584 if is_flat && has_index {
585 options.index = Some(index.clone());
586 }
587 }
588
589 if options.index.is_none() && options.limit.is_none() {
590 options.limit = view.source.as_ref().and_then(|x| x.options.limit);
591 }
592
593 let window = ViewWindow::default();
594 let arrow = view.to_arrow(window).await?;
595 let mut table = self
596 .crate_table_inner(UpdateData::Arrow(arrow).into(), options.into(), entity_id)
597 .await?;
598
599 let table_ = table.clone();
600 let callback = asyncfn!(table_, update, async move |update: OnUpdateData| {
601 let update = UpdateData::Arrow(update.delta.expect("Malformed message").into());
602 let options = crate::UpdateOptions::default();
603 table_.update(update, options).await.unwrap_or_log();
604 });
605
606 let options = OnUpdateOptions {
607 mode: Some(OnUpdateMode::Row),
608 };
609
610 let update_token = view.on_update(callback, options).await?;
611 let remove_token = if source_index.is_some() && source_index == table.get_index() {
612 let table_ = table.clone();
613 let callback = asyncfn!(table_, async move |removed: OnRemoveData| {
614 if let Some(indices) = removed.indices.as_ref().filter(|x| !x.is_empty()) {
615 let indices = UpdateData::Arrow(indices.clone().into());
616 table_.remove(indices).await.unwrap_or_log();
617 }
618 });
619
620 Some(view.on_remove(callback).await?)
621 } else {
622 None
623 };
624
625 table.view_binding = Some(ViewBinding {
626 view: view.clone(),
627 update_token,
628 remove_token,
629 });
630
631 Ok(table)
632 } else {
633 self.crate_table_inner(input, options.into(), entity_id)
634 .await
635 }
636 }
637
638 async fn crate_table_inner(
639 &self,
640 input: TableData,
641 options: TableOptions,
642 entity_id: String,
643 ) -> ClientResult<Table> {
644 let msg = Request {
645 msg_id: self.gen_id(),
646 entity_id: entity_id.clone(),
647 client_req: Some(ClientReq::MakeTableReq(MakeTableReq {
648 data: Some(input.into()),
649 options: Some(options.clone().try_into()?),
650 })),
651 };
652
653 let client = self.clone();
654 match self.oneshot(&msg).await? {
655 ClientResp::MakeTableResp(_) => Ok(Table::new(entity_id, client, options)),
656 resp => Err(resp.into()),
657 }
658 }
659
660 pub async fn join(
672 &self,
673 left: TableRef,
674 right: TableRef,
675 on: &str,
676 options: JoinOptions,
677 ) -> ClientResult<Table> {
678 let entity_id = options.name.unwrap_or_else(randid);
679 let join_type: JoinType = options.join_type.unwrap_or_default();
680 let right_on_column = options.right_on.unwrap_or_default();
681 let msg = Request {
682 msg_id: self.gen_id(),
683 entity_id: entity_id.clone(),
684 client_req: Some(ClientReq::MakeJoinTableReq(MakeJoinTableReq {
685 left_table_id: left.table_name().to_owned(),
686 right_table_id: right.table_name().to_owned(),
687 on_column: on.to_owned(),
688 join_type: join_type.into(),
689 right_on_column,
690 })),
691 };
692
693 let client = self.clone();
694 match self.oneshot(&msg).await? {
695 ClientResp::MakeJoinTableResp(_) => Ok(Table::new(entity_id, client, TableOptions {
696 index: Some(on.to_owned()),
697 limit: None,
698 page_to_disk: None,
699 list_flatten: None,
700 })),
701 resp => Err(resp.into()),
702 }
703 }
704
705 async fn get_table_infos(&self) -> ClientResult<Vec<HostedTable>> {
706 let msg = Request {
707 msg_id: self.gen_id(),
708 entity_id: "".to_owned(),
709 client_req: Some(ClientReq::GetHostedTablesReq(GetHostedTablesReq {
710 subscribe: false,
711 })),
712 };
713
714 match self.oneshot(&msg).await? {
715 ClientResp::GetHostedTablesResp(GetHostedTablesResp { table_infos }) => Ok(table_infos),
716 resp => Err(resp.into()),
717 }
718 }
719
720 pub async fn open_table(&self, entity_id: String) -> ClientResult<Table> {
737 let infos = self.get_table_infos().await?;
738
739 if let Some(info) = infos.into_iter().find(|i| i.entity_id == entity_id) {
741 let options = TableOptions {
742 index: info.index,
743 limit: info.limit,
744 page_to_disk: None,
745 list_flatten: None,
746 };
747
748 let client = self.clone();
749 Ok(Table::new(entity_id, client, options))
750 } else {
751 Err(ClientError::Unknown(format!(
752 "Unknown table \"{}\"",
753 entity_id
754 )))
755 }
756 }
757
758 pub async fn get_hosted_table_names(&self) -> ClientResult<Vec<String>> {
775 let msg = Request {
776 msg_id: self.gen_id(),
777 entity_id: "".to_owned(),
778 client_req: Some(ClientReq::GetHostedTablesReq(GetHostedTablesReq {
779 subscribe: false,
780 })),
781 };
782
783 match self.oneshot(&msg).await? {
784 ClientResp::GetHostedTablesResp(GetHostedTablesResp { table_infos }) => {
785 Ok(table_infos.into_iter().map(|i| i.entity_id).collect())
786 },
787 resp => Err(resp.into()),
788 }
789 }
790
791 pub async fn on_hosted_tables_update<T, U>(&self, on_update: T) -> ClientResult<u32>
795 where
796 T: Fn() -> U + Send + Sync + 'static,
797 U: Future<Output = ()> + Send + 'static,
798 {
799 let on_update = Arc::new(on_update);
800 let callback = asyncfn!(on_update, async move |resp: Response| {
801 match resp.client_resp {
802 Some(ClientResp::GetHostedTablesResp(_)) | None => {
803 on_update().await;
804 Ok(())
805 },
806 resp => Err(resp.into()),
807 }
808 });
809
810 let msg = Request {
811 msg_id: self.gen_id(),
812 entity_id: "".to_owned(),
813 client_req: Some(ClientReq::GetHostedTablesReq(GetHostedTablesReq {
814 subscribe: true,
815 })),
816 };
817
818 self.subscribe(&msg, callback).await?;
819 Ok(msg.msg_id)
820 }
821
822 pub async fn remove_hosted_tables_update(&self, update_id: u32) -> ClientResult<()> {
825 let msg = Request {
826 msg_id: self.gen_id(),
827 entity_id: "".to_owned(),
828 client_req: Some(ClientReq::RemoveHostedTablesUpdateReq(
829 RemoveHostedTablesUpdateReq { id: update_id },
830 )),
831 };
832
833 self.unsubscribe(update_id).await?;
834 match self.oneshot(&msg).await? {
835 ClientResp::RemoveHostedTablesUpdateResp(_) => Ok(()),
836 resp => Err(resp.into()),
837 }
838 }
839
840 pub async fn system_info(&self) -> ClientResult<SystemInfo> {
844 let msg = Request {
845 msg_id: self.gen_id(),
846 entity_id: "".to_string(),
847 client_req: Some(ClientReq::ServerSystemInfoReq(ServerSystemInfoReq {})),
848 };
849
850 match self.oneshot(&msg).await? {
851 ClientResp::ServerSystemInfoResp(resp) => {
852 #[cfg(not(target_family = "wasm"))]
853 let timestamp = Some(
854 std::time::SystemTime::now()
855 .duration_since(std::time::UNIX_EPOCH)?
856 .as_millis() as u64,
857 );
858
859 #[cfg(target_family = "wasm")]
860 let timestamp = None;
861
862 #[cfg(feature = "talc-allocator")]
863 let (client_used, client_heap) = {
864 let (client_used, client_heap) = crate::utils::get_used();
865 (Some(client_used as u64), Some(client_heap as u64))
866 };
867
868 #[cfg(not(feature = "talc-allocator"))]
869 let (client_used, client_heap) = (None, None);
870
871 let info = SystemInfo {
872 heap_size: resp.heap_size,
873 used_size: resp.used_size,
874 cpu_time: resp.cpu_time,
875 cpu_time_epoch: resp.cpu_time_epoch,
876 timestamp,
877 client_heap,
878 client_used,
879 };
880
881 Ok(info)
882 },
883 resp => Err(resp.into()),
884 }
885 }
886}