Merge 'introduce program execution state in order to run stmt to completion in case of finalize or reset' from Nikita Sivukhin

This PR introduces program execution state in order for statement to be
aware of its state - is it terminal (Done, Failed, Interrupted) or not.
The particular problem right now is that statements like `INSERT INTO t
VALUES (1), (2), (3) RETURNING x` will execute inserts one by one and
interleave them with rows generation. This means that if statement
consumer will just read one row and then finalize the statement -
nothing will be actually committed (because transaction will be
aborted).
In order to quickly mitigate this issue - program state is introduced
which can help to decide what to do in the finalize.

Reviewed-by: Jussi Saurio <[email protected]>

Closes #4038

Author: jussisaurio

core/lib.rs

@@ -31,6 +31,7 @@ mod util;
 #[cfg(feature = "uuid")]
 mod uuid;
 mod vdbe;
+pub type ProgramExecutionState = vdbe::ProgramExecutionState;
 pub mod vector;
 mod vtab;
 
@@ -2650,6 +2651,10 @@ impl Statement {
         self.state.interrupt();
     }
 
+    pub fn execution_state(&self) -> ProgramExecutionState {
+        self.state.execution_state
+    }
+
     fn _step(&mut self, waker: Option<&Waker>) -> Result<StepResult> {
         if let Some(busy_timeout) = self.busy_timeout.as_ref() {
             if self.pager.io.now() < busy_timeout.timeout {

core/vdbe/mod.rs

@@ -273,6 +273,39 @@ pub enum TxnCleanup {
     RollbackTxn,
 }
 
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub enum ProgramExecutionState {
+    /// No steps of the program was executed
+    Init,
+    /// Program started execution but didn't reach any terminal state
+    Running,
+    /// Interrupt requested for the program
+    Interrupting,
+    /// Terminal state: program interrupted
+    Interrupted,
+    /// Terminal state: program finished successfully
+    Done,
+    /// Terminal state: program failed with error
+    Failed,
+}
+
+impl ProgramExecutionState {
+    pub fn is_running(&self) -> bool {
+        matches!(
+            self,
+            ProgramExecutionState::Interrupting | ProgramExecutionState::Running
+        )
+    }
+    pub fn is_terminal(&self) -> bool {
+        matches!(
+            self,
+            ProgramExecutionState::Interrupted
+                | ProgramExecutionState::Failed
+                | ProgramExecutionState::Done
+        )
+    }
+}
+
 /// The program state describes the environment in which the program executes.
 pub struct ProgramState {
     pub io_completions: Option<IOCompletions>,
@@ -287,7 +320,7 @@ pub struct ProgramState {
     /// Indicate whether an [Insn::Once] instruction at a given program counter position has already been executed, well, once.
     once: SmallVec<u32, 4>,
     regex_cache: RegexCache,
-    interrupted: bool,
+    pub execution_state: ProgramExecutionState,
     pub parameters: HashMap<NonZero<usize>, Value>,
     commit_state: CommitState,
     #[cfg(feature = "json")]
@@ -356,7 +389,7 @@ impl ProgramState {
             ended_coroutine: Bitfield::new(),
             once: SmallVec::<u32, 4>::new(),
             regex_cache: RegexCache::new(),
-            interrupted: false,
+            execution_state: ProgramExecutionState::Init,
             parameters: HashMap::new(),
             commit_state: CommitState::Ready,
             #[cfg(feature = "json")]
@@ -409,11 +442,7 @@ impl ProgramState {
     }
 
     pub fn interrupt(&mut self) {
-        self.interrupted = true;
-    }
-
-    pub fn is_interrupted(&self) -> bool {
-        self.interrupted
+        self.execution_state = ProgramExecutionState::Interrupting;
     }
 
     pub fn bind_at(&mut self, index: NonZero<usize>, value: Value) {
@@ -450,7 +479,7 @@ impl ProgramState {
         self.deferred_seeks.iter_mut().for_each(|s| *s = None);
         self.ended_coroutine.0 = [0; 4];
         self.regex_cache.like.clear();
-        self.interrupted = false;
+        self.execution_state = ProgramExecutionState::Init;
         self.current_collation = None;
         #[cfg(feature = "json")]
         self.json_cache.clear();
@@ -650,11 +679,25 @@ impl Program {
         query_mode: QueryMode,
         waker: Option<&Waker>,
     ) -> Result<StepResult> {
-        match query_mode {
+        state.execution_state = ProgramExecutionState::Running;
+        let result = match query_mode {
             QueryMode::Normal => self.normal_step(state, mv_store, pager, waker),
             QueryMode::Explain => self.explain_step(state, mv_store, pager),
             QueryMode::ExplainQueryPlan => self.explain_query_plan_step(state, mv_store, pager),
+        };
+        match &result {
+            Ok(StepResult::Done) => {
+                state.execution_state = ProgramExecutionState::Done;
+            }
+            Ok(StepResult::Interrupt) => {
+                state.execution_state = ProgramExecutionState::Interrupted;
+            }
+            Err(_) => {
+                state.execution_state = ProgramExecutionState::Failed;
+            }
+            _ => {}
         }
+        result
     }
 
     fn explain_step(
@@ -673,7 +716,7 @@ impl Program {
             return Err(LimboError::InternalError("Connection closed".to_string()));
         }
 
-        if state.is_interrupted() {
+        if matches!(state.execution_state, ProgramExecutionState::Interrupting) {
             return Ok(StepResult::Interrupt);
         }
 
@@ -823,7 +866,7 @@ impl Program {
                 return Err(LimboError::InternalError("Connection closed".to_string()));
             }
 
-            if state.is_interrupted() {
+            if matches!(state.execution_state, ProgramExecutionState::Interrupting) {
                 return Ok(StepResult::Interrupt);
             }
 
@@ -871,10 +914,11 @@ impl Program {
                 }
                 return Err(LimboError::InternalError("Connection closed".to_string()));
             }
-            if state.is_interrupted() {
+            if matches!(state.execution_state, ProgramExecutionState::Interrupting) {
                 self.abort(mv_store, &pager, None, state);
                 return Ok(StepResult::Interrupt);
             }
+
             if let Some(io) = &state.io_completions {
                 if !io.finished() {
                     io.set_waker(waker);

sqlite3/src/lib.rs

@@ -325,13 +325,32 @@ pub unsafe extern "C" fn sqlite3_prepare_v2(
     SQLITE_OK
 }
 
+unsafe fn stmt_run_to_completion(stmt: *mut sqlite3_stmt) -> ffi::c_int {
+    let stmt_ref = &mut *stmt;
+    while stmt_ref.stmt.execution_state().is_running() {
+        let result = sqlite3_step(stmt);
+        if result != SQLITE_DONE && result != SQLITE_ROW {
+            return result;
+        }
+    }
+    SQLITE_OK
+}
+
 #[no_mangle]
 pub unsafe extern "C" fn sqlite3_finalize(stmt: *mut sqlite3_stmt) -> ffi::c_int {
     if stmt.is_null() {
         return SQLITE_MISUSE;
     }
     let stmt_ref = &mut *stmt;
 
+    // first, finalize any execution if it was unfinished
+    // (for example, many drivers can consume just one row and finalize statement after that, while there still can be work to do)
+    // (this is necessary because queries like INSERT INTO t VALUES (1), (2), (3) RETURNING id return values within a transaction)
+    let result = stmt_run_to_completion(stmt);
+    if result != SQLITE_OK {
+        return result;
+    }
+
     if !stmt_ref.db.is_null() {
         let db = &mut *stmt_ref.db;
         let mut db_inner = db.inner.lock().unwrap();
@@ -669,6 +688,13 @@ fn split_sql_statements(sql: &str) -> Vec<&str> {
 #[no_mangle]
 pub unsafe extern "C" fn sqlite3_reset(stmt: *mut sqlite3_stmt) -> ffi::c_int {
     let stmt = &mut *stmt;
+    // first, finalize any execution if it was unfinished
+    // (for example, many drivers can consume just one row and finalize statement after that, while there still can be work to do)
+    // (this is necessary because queries like INSERT INTO t VALUES (1), (2), (3) RETURNING id return values within a transaction)
+    let result = stmt_run_to_completion(stmt);
+    if result != SQLITE_OK {
+        return result;
+    }
     stmt.stmt.reset();
     stmt.clear_text_cache();
     SQLITE_OK
@@ -1420,10 +1446,10 @@ unsafe extern "C" fn sqlite_get_table_cb(
     for i in 0..n_column {
         let value = *argv.add(i as usize);
         let value_cstring = if !value.is_null() {
-            let len = libc::strlen(value);
-            let mut buf = Vec::with_capacity(len + 1);
-            libc::strncpy(buf.as_mut_ptr() as *mut ffi::c_char, value, len);
-            buf.set_len(len + 1);
+            let value_cstr = CStr::from_ptr(value).to_bytes();
+            let len = value_cstr.len();
+            let mut buf = vec![0u8; len + 1];
+            buf[0..len].copy_from_slice(value_cstr);
             CString::from_vec_with_nul(buf).unwrap()
         } else {
             CString::new("NULL").unwrap()

sqlite3/tests/sqlite3_tests.c

@@ -20,6 +20,7 @@ void test_sqlite3_column_type();
 void test_sqlite3_column_decltype();
 void test_sqlite3_next_stmt();
 void test_sqlite3_table_column_metadata();
+void test_sqlite3_insert_returning();
 
 int allocated = 0;
 
@@ -39,6 +40,7 @@ int main(void)
     test_sqlite3_column_decltype();
     test_sqlite3_next_stmt();
     test_sqlite3_table_column_metadata();
+    test_sqlite3_insert_returning();
     return 0;
 }
 
@@ -746,3 +748,47 @@ void test_sqlite3_table_column_metadata()
 	printf("sqlite3_table_column_metadata test passed\n");
 	sqlite3_close(db);
 }
+
+void test_sqlite3_insert_returning()
+{
+    sqlite3 *db;
+    sqlite3_stmt *stmt;
+    char *err_msg = NULL;
+    int rc;
+
+    rc = sqlite3_open(":memory:", &db);
+    assert(rc == SQLITE_OK);
+
+    rc = sqlite3_exec(db,
+                      "CREATE TABLE t(x)",
+                      NULL, NULL, &err_msg);
+    assert(rc == SQLITE_OK);
+
+    if (err_msg)
+    {
+        sqlite3_free(err_msg);
+        err_msg = NULL;
+    }
+    rc = sqlite3_prepare_v2(db, "INSERT INTO t (x) VALUES (1), (2), (3) RETURNING x;", -1, &stmt, NULL);
+    assert(rc == SQLITE_OK);
+
+    rc = sqlite3_step(stmt);
+    assert(rc == SQLITE_ROW);
+
+    rc = sqlite3_finalize(stmt);
+    assert(rc == SQLITE_OK);
+
+    rc = sqlite3_prepare_v2(db, "SELECT COUNT(*) FROM t;", -1, &stmt, NULL);
+    assert(rc == SQLITE_OK);
+
+    rc = sqlite3_step(stmt);
+    assert(rc == SQLITE_ROW);
+
+    sqlite_int64 fetched = sqlite3_column_int64(stmt, 0);
+    assert(fetched == 3);
+
+    sqlite3_finalize(stmt);
+    sqlite3_close(db);
+
+    printf("test_sqlite3_insert_retuning test passed\n");
+}