fix remote execution

Author: MarinPostma

sqld/src/database/libsql.rs

@@ -207,8 +207,8 @@ impl Connection {
     fn run(&mut self, pgm: Program) -> Vec<Option<QueryResult>> {
         let mut results = Vec::with_capacity(pgm.steps.len());
 
-        for step in pgm.steps {
-            let res = self.execute_step(&step, &results);
+        for step in pgm.steps() {
+            let res = self.execute_step(step, &results);
             results.push(res);
         }
 

sqld/src/database/mod.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use crate::query::{Params, Query, QueryResult};
 use crate::query_analysis::{State, Statement};
 use crate::Result;
@@ -9,24 +11,34 @@ pub mod write_proxy;
 
 const TXN_TIMEOUT_SECS: u64 = 5;
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct Program {
-    pub steps: Vec<Step>,
+    pub steps: Arc<Vec<Step>>,
 }
 
 impl Program {
+    pub fn new(steps: Vec<Step>) -> Self {
+        Self {
+            steps: Arc::new(steps),
+        }
+    }
+
     pub fn is_read_only(&self) -> bool {
         self.steps.iter().all(|s| s.query.stmt.is_read_only())
     }
+
+    pub fn steps(&self) -> &[Step] {
+        self.steps.as_slice()
+    }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct Step {
     pub cond: Option<Cond>,
     pub query: Query,
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub enum Cond {
     Ok { step: usize },
     Err { step: usize },
@@ -42,9 +54,7 @@ pub trait Database: Send + Sync {
 
     /// Unconditionnaly execute a query as part of a program
     async fn execute_one(&self, query: Query) -> Result<(QueryResult, State)> {
-        let pgm = Program {
-            steps: vec![Step { cond: None, query }],
-        };
+        let pgm = Program::new(vec![Step { cond: None, query }]);
 
         let (results, state) = self.execute_program(pgm).await?;
         Ok((results.into_iter().next().unwrap().unwrap(), state))
@@ -85,12 +95,25 @@ pub trait Database: Send + Sync {
             })
         }
 
-        let pgm = Program { steps };
+        let pgm = Program::new(steps);
 
         let (mut results, state) = self.execute_program(pgm).await?;
         // remove the rollback result
         results.pop();
 
         Ok((results, state))
     }
+
+    async fn rollback(&self) -> Result<()> {
+        let (results, _) = self
+            .execute_one(Query {
+                stmt: Statement::parse("ROLLBACK").next().unwrap().unwrap(),
+                params: Params::empty(),
+            })
+            .await?;
+
+        results?;
+
+        Ok(())
+    }
 }

sqld/src/database/write_proxy.rs

@@ -10,7 +10,7 @@ use uuid::Uuid;
 
 use crate::error::Error;
 use crate::query::{QueryResponse, QueryResult};
-use crate::query_analysis::{final_state, State};
+use crate::query_analysis::State;
 use crate::replication::client::PeriodicDbUpdater;
 use crate::rpc::proxy::rpc::proxy_client::ProxyClient;
 use crate::rpc::proxy::rpc::query_result::RowResult;
@@ -115,50 +115,62 @@ impl WriteProxyDatabase {
             client_id: Uuid::new_v4(),
         })
     }
+
+    async fn execute_remote(
+        &self,
+        pgm: Program,
+        state: &mut State,
+    ) -> Result<(Vec<Option<QueryResult>>, State)> {
+        let mut client = self.write_proxy.clone();
+        let req = crate::rpc::proxy::rpc::ProgramReq {
+            client_id: self.client_id.to_string(),
+            pgm: Some(pgm.into()),
+        };
+        match client.execute(req).await {
+            Ok(r) => {
+                let execute_result = r.into_inner();
+                *state = execute_result.state().into();
+                let results = execute_result
+                    .results
+                    .into_iter()
+                    .map(|r| -> Option<QueryResult> {
+                        let result = r.row_result?;
+                        match result {
+                            RowResult::Row(res) => Some(Ok(QueryResponse::ResultSet(res.into()))),
+                            RowResult::Error(e) => Some(Err(Error::RpcQueryError(e))),
+                        }
+                    })
+                    .collect();
+
+                Ok((results, *state))
+            }
+            Err(e) => {
+                // Set state to invalid, so next call is sent to remote, and we have a chance
+                // to recover state.
+                *state = State::Invalid;
+                Err(Error::RpcQueryExecutionError(e))
+            }
+        }
+    }
 }
 
 #[async_trait::async_trait]
 impl Database for WriteProxyDatabase {
     async fn execute_program(&self, pgm: Program) -> Result<(Vec<Option<QueryResult>>, State)> {
         let mut state = self.state.lock().await;
-        if *state == State::Init
-            && pgm.is_read_only()
-            && final_state(*state, pgm.steps.iter().map(|s| &s.query.stmt)) == State::Init
-        {
-            self.read_db.execute_program(pgm).await
-        } else {
-            let mut client = self.write_proxy.clone();
-            let req = crate::rpc::proxy::rpc::ProgramReq {
-                client_id: self.client_id.to_string(),
-                pgm: Some(pgm.into()),
-            };
-            match client.execute(req).await {
-                Ok(r) => {
-                    let execute_result = r.into_inner();
-                    *state = execute_result.state().into();
-                    let results = execute_result
-                        .results
-                        .into_iter()
-                        .map(|r| -> Option<QueryResult> {
-                            let result = r.row_result?;
-                            match result {
-                                RowResult::Row(res) => {
-                                    Some(Ok(QueryResponse::ResultSet(res.into())))
-                                }
-                                RowResult::Error(e) => Some(Err(Error::RpcQueryError(e))),
-                            }
-                        })
-                        .collect();
-
-                    Ok((results, *state))
-                }
-                Err(e) => {
-                    // Set state to invalid, so next call is sent to remote, and we have a chance
-                    // to recover state.
-                    *state = State::Invalid;
-                    Err(Error::RpcQueryExecutionError(e))
-                }
+        if *state == State::Init && pgm.is_read_only() {
+            // We know that this program won't perform any writes. We attempt to run it on the
+            // replica. If it leaves an open transaction, then this program is an interactive
+            // transaction, so we rollback the replica, and execute again on the primary.
+            let (results, new_state) = self.read_db.execute_program(pgm.clone()).await?;
+            if new_state != State::Init {
+                self.read_db.rollback().await?;
+                self.execute_remote(pgm, &mut state).await
+            } else {
+                Ok((results, new_state))
             }
+        } else {
+            self.execute_remote(pgm, &mut state).await
         }
     }
 }

sqld/src/hrana/batch.rs

@@ -68,7 +68,7 @@ fn batch_to_program(batch: &proto::Batch) -> Result<Program> {
         steps.push(step);
     }
 
-    Ok(Program { steps })
+    Ok(Program::new(steps))
 }
 
 pub async fn execute_batch(db: &dyn Database, batch: &proto::Batch) -> Result<proto::BatchResult> {

sqld/src/http/mod.rs

@@ -28,7 +28,7 @@ use crate::error::Error;
 use crate::hrana;
 use crate::http::types::HttpQuery;
 use crate::query::{self, Query, QueryResult, ResultSet};
-use crate::query_analysis::{final_state, State, Statement};
+use crate::query_analysis::{predict_final_state, State, Statement};
 use crate::utils::services::idle_shutdown::IdleShutdownLayer;
 
 use self::types::QueryObject;
@@ -140,7 +140,7 @@ fn parse_queries(queries: Vec<QueryObject>) -> anyhow::Result<Vec<Query>> {
         out.push(query);
     }
 
-    match final_state(State::Init, out.iter().map(|q| &q.stmt)) {
+    match predict_final_state(State::Init, out.iter().map(|q| &q.stmt)) {
         State::Txn => anyhow::bail!("interactive transaction not allowed in HTTP queries"),
         State::Init => (),
         // maybe we should err here, but let's sqlite deal with that.

sqld/src/query.rs

@@ -271,7 +271,7 @@ pub enum QueryResponse {
     ResultSet(ResultSet),
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct Query {
     pub stmt: Statement,
     pub params: Params,
@@ -291,7 +291,7 @@ impl ToSql for Value {
     }
 }
 
-#[derive(Debug, Serialize)]
+#[derive(Debug, Serialize, Clone)]
 pub enum Params {
     Named(HashMap<String, Value>),
     Positional(Vec<Value>),

sqld/src/query_analysis.rs

@@ -6,7 +6,7 @@ use sqlite3_parser::{
 };
 
 /// A group of statements to be executed together.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct Statement {
     pub stmt: String,
     pub kind: StmtKind,
@@ -70,11 +70,7 @@ pub enum State {
 impl State {
     pub fn step(&mut self, kind: StmtKind) {
         *self = match (*self, kind) {
-            // those two transition will cause an error, but since we are interested in what the
-            // pesimistic final state is, and we will adjust when we get the actual state back.
-            (State::Txn, StmtKind::TxnBegin) => State::Txn,
-            (State::Init, StmtKind::TxnEnd) => State::Init,
-
+            (State::Txn, StmtKind::TxnBegin) | (State::Init, StmtKind::TxnEnd) => State::Invalid,
             (State::Txn, StmtKind::TxnEnd) => State::Init,
             (state, StmtKind::Other | StmtKind::Write | StmtKind::Read) => state,
             (State::Invalid, _) => State::Invalid,
@@ -145,9 +141,12 @@ impl Statement {
     }
 }
 
-/// Given a an initial state and an array of queries, return the final state obtained if all the
-/// queries succeeded
-pub fn final_state<'a>(mut state: State, stmts: impl Iterator<Item = &'a Statement>) -> State {
+/// Given a an initial state and an array of queries, attempts to predict what the final state will
+/// be
+pub fn predict_final_state<'a>(
+    mut state: State,
+    stmts: impl Iterator<Item = &'a Statement>,
+) -> State {
     for stmt in stmts {
         state.step(stmt.kind);
     }

sqld/src/rpc/proxy.rs

@@ -15,6 +15,8 @@ use self::rpc::{Ack, DisconnectMessage, ExecuteResults};
 pub mod rpc {
     #![allow(clippy::all)]
 
+    use std::sync::Arc;
+
     use anyhow::Context;
 
     use crate::query::QueryResponse;
@@ -145,13 +147,13 @@ pub mod rpc {
         type Error = anyhow::Error;
 
         fn try_from(pgm: Program) -> Result<Self, Self::Error> {
-            Ok(Self {
-                steps: pgm
-                    .steps
-                    .into_iter()
-                    .map(TryInto::try_into)
-                    .collect::<anyhow::Result<_>>()?,
-            })
+            let steps = pgm
+                .steps
+                .into_iter()
+                .map(TryInto::try_into)
+                .collect::<anyhow::Result<_>>()?;
+
+            Ok(Self::new(steps))
         }
     }
 
@@ -215,8 +217,14 @@ pub mod rpc {
 
     impl From<database::Program> for Program {
         fn from(pgm: database::Program) -> Self {
+            // TODO: use unwrap_or_clone when stable
+            let steps = match Arc::try_unwrap(pgm.steps) {
+                Ok(steps) => steps,
+                Err(arc) => (*arc).clone(),
+            };
+
             Self {
-                steps: pgm.steps.into_iter().map(|s| s.into()).collect(),
+                steps: steps.into_iter().map(|s| s.into()).collect(),
             }
         }
     }