port hrana for libsqlx server

we can now allocate primaries, and query them

Author: MarinPostma

Cargo.lock

@@ -7,15 +7,23 @@ edition = "2021"
 
 [dependencies]
 axum = "0.6.18"
+base64 = "0.21.2"
 bincode = "1.3.3"
+bytes = "1.4.0"
 clap = { version = "4.3.11", features = ["derive"] }
 color-eyre = "0.6.2"
 futures = "0.3.28"
+hmac = "0.12.1"
 hyper = { version = "0.14.27", features = ["h2", "server"] }
 libsqlx = { version = "0.1.0", path = "../libsqlx" }
 moka = { version = "0.11.2", features = ["future"] }
+parking_lot = "0.12.1"
+priority-queue = "1.3.2"
+rand = "0.8.5"
 regex = "1.9.1"
-serde = { version = "1.0.166", features = ["derive"] }
+serde = { version = "1.0.166", features = ["derive", "rc"] }
+serde_json = "1.0.100"
+sha2 = "0.10.7"
 sled = "0.34.7"
 thiserror = "1.0.43"
 tokio = { version = "1.29.1", features = ["full"] }

libsqlx-server/Cargo.toml

@@ -1,11 +1,15 @@
-use std::collections::HashMap;
 use std::path::PathBuf;
+use std::sync::Arc;
 
 use libsqlx::libsql::{LibsqlDatabase, LogCompactor, LogFile, PrimaryType};
 use libsqlx::Database as _;
 use tokio::sync::{mpsc, oneshot};
 use tokio::task::{block_in_place, JoinSet};
 
+use crate::hrana;
+use crate::hrana::http::handle_pipeline;
+use crate::hrana::http::proto::{PipelineRequestBody, PipelineResponseBody};
+
 use self::config::{AllocConfig, DbConfig};
 
 pub mod config;
@@ -19,16 +23,11 @@ pub struct ConnectionId {
 }
 
 pub enum AllocationMessage {
-    /// Execute callback against connection
-    Exec {
-        connection_id: ConnectionId,
-        exec: ExecFn,
-    },
-    /// Create a new connection, execute the callback and return the connection id.
-    NewConnExec {
-        exec: ExecFn,
-        ret: oneshot::Sender<ConnectionId>,
-    },
+    NewConnection(oneshot::Sender<ConnectionHandle>),
+    HranaPipelineReq {
+        req: PipelineRequestBody,
+        ret: oneshot::Sender<crate::Result<PipelineResponseBody>>,
+    }
 }
 
 pub enum Database {
@@ -73,12 +72,34 @@ impl Database {
 pub struct Allocation {
     pub inbox: mpsc::Receiver<AllocationMessage>,
     pub database: Database,
-    /// senders to the spawned connections
-    pub connections: HashMap<u32, mpsc::Sender<ExecFn>>,
     /// spawned connection futures, returning their connection id on completion.
     pub connections_futs: JoinSet<u32>,
     pub next_conn_id: u32,
     pub max_concurrent_connections: u32,
+
+    pub hrana_server: Arc<hrana::http::Server>,
+}
+
+pub struct ConnectionHandle {
+    exec: mpsc::Sender<ExecFn>,
+    exit: oneshot::Sender<()>,
+}
+
+impl ConnectionHandle {
+    pub async fn exec<F, R>(&self, f: F) -> crate::Result<R>
+        where F: for<'a> FnOnce(&'a mut (dyn libsqlx::Connection + 'a)) -> R + Send + 'static,
+              R: Send + 'static,
+    {
+        let (sender, ret) = oneshot::channel();
+        let cb = move |conn: &mut dyn libsqlx::Connection| {
+            let res = f(conn);
+            let _ = sender.send(res);
+        };
+
+        self.exec.send(Box::new(cb)).await.unwrap();
+
+        Ok(ret.await?)
+    }
 }
 
 impl Allocation {
@@ -87,34 +108,36 @@ impl Allocation {
             tokio::select! {
                 Some(msg) = self.inbox.recv() => {
                     match msg {
-                        AllocationMessage::Exec { connection_id, exec } => {
-                            if let Some(sender) = self.connections.get(&connection_id.id) {
-                                if let Err(_) = sender.send(exec).await {
-                                    tracing::debug!("connection {} closed.", connection_id.id);
-                                    self.connections.remove_entry(&connection_id.id);
-                                }
-                            }
-                        },
-                        AllocationMessage::NewConnExec { exec, ret } => {
-                            let id = self.new_conn_exec(exec).await;
-                            let _ = ret.send(id);
+                        AllocationMessage::NewConnection(ret) => {
+                            let _ =ret.send(self.new_conn().await);
                         },
+                        AllocationMessage::HranaPipelineReq { req, ret} => {
+                            let res = handle_pipeline(&self.hrana_server.clone(), req, || async {
+                                let conn= self.new_conn().await;
+                                dbg!();
+                                Ok(conn)
+                            }).await;
+                            let _ = ret.send(res);
+                        }
                     }
                 },
                 maybe_id = self.connections_futs.join_next() => {
-                    if let Some(Ok(id)) = maybe_id {
-                        self.connections.remove_entry(&id);
+                    if let Some(Ok(_id)) = maybe_id {
+                        // self.connections.remove_entry(&id);
                     }
                 },
                 else => break,
             }
         }
     }
 
-    async fn new_conn_exec(&mut self, exec: ExecFn) -> ConnectionId {
+    async fn new_conn(&mut self) -> ConnectionHandle {
+        dbg!();
         let id = self.next_conn_id();
+        dbg!();
         let conn = block_in_place(|| self.database.connect());
-        let (close_sender, exit) = mpsc::channel(1);
+        dbg!();
+        let (close_sender, exit) = oneshot::channel();
         let (exec_sender, exec_receiver) = mpsc::channel(1);
         let conn = Connection {
             id,
@@ -123,36 +146,40 @@ impl Allocation {
             exec: exec_receiver,
         };
 
+        dbg!();
         self.connections_futs.spawn(conn.run());
-        // This should never block!
-        assert!(exec_sender.try_send(exec).is_ok());
-        assert!(self.connections.insert(id, exec_sender).is_none());
+        dbg!();
+
+        ConnectionHandle {
+            exec: exec_sender,
+            exit: close_sender,
+        }
 
-        ConnectionId { id, close_sender }
     }
 
     fn next_conn_id(&mut self) -> u32 {
         loop {
             self.next_conn_id = self.next_conn_id.wrapping_add(1);
-            if !self.connections.contains_key(&self.next_conn_id) {
-                return self.next_conn_id;
-            }
+            return self.next_conn_id;
+            // if !self.connections.contains_key(&self.next_conn_id) {
+            //     return self.next_conn_id;
+            // }
         }
     }
 }
 
 struct Connection {
     id: u32,
     conn: Box<dyn libsqlx::Connection + Send>,
-    exit: mpsc::Receiver<()>,
+    exit: oneshot::Receiver<()>,
     exec: mpsc::Receiver<ExecFn>,
 }
 
 impl Connection {
     async fn run(mut self) -> u32 {
         loop {
             tokio::select! {
-                _ = self.exit.recv() => break,
+                _ = &mut self.exit => break,
                 Some(exec) = self.exec.recv() => {
                     tokio::task::block_in_place(|| exec(&mut *self.conn));
                 }

libsqlx-server/src/allocation/mod.rs

@@ -0,0 +1,19 @@
+use tokio::sync::{mpsc, oneshot};
+
+use crate::hrana::http::proto::{PipelineRequestBody, PipelineResponseBody};
+use crate::allocation::{AllocationMessage, ConnectionHandle};
+
+pub struct Database {
+    pub sender: mpsc::Sender<AllocationMessage>,
+}
+
+impl Database {
+    pub async fn hrana_pipeline(&self, req: PipelineRequestBody) -> crate::Result<PipelineResponseBody> {
+        dbg!();
+        let (sender, ret) = oneshot::channel();
+        dbg!();
+        self.sender.send(AllocationMessage::HranaPipelineReq { req, ret: sender }).await.unwrap();
+        dbg!();
+        ret.await.unwrap()
+    }
+}