Make complete_chunk and fail_chunk not error when doing so for previously completed, failed, or paused chunks

Author: SemMulder

libs/opsqueue_python/python/opsqueue/exceptions.py

@@ -92,15 +92,6 @@ class TryFromIntError(IncorrectUsageError):
     pass
 
 
-class ChunkNotFoundError(IncorrectUsageError):
-    """
-    Raised when a method is used to look up information about a chunk
-    but the chunk doesn't exist within the Opsqueue.
-    """
-
-    pass
-
-
 class SubmissionNotFoundError(IncorrectUsageError):
     """
     Raised when a method is used to look up information about a submission

libs/opsqueue_python/src/errors.rs

@@ -2,24 +2,22 @@
 /// so we have nice IDE support for docs-on-hover and for 'go to definition'.
 use std::error::Error;
 
-use opsqueue::common::chunk::ChunkId;
+
 use opsqueue::common::errors::{
-    ChunkNotFound, IncorrectUsage, SubmissionNotCancellable, SubmissionNotFound,
+    IncorrectUsage, SubmissionNotCancellable, SubmissionNotFound,
     UnexpectedOpsqueueConsumerServerResponse, E,
 };
 use pyo3::exceptions::PyBaseException;
 use pyo3::{import_exception, Bound, PyErr, Python};
 
 use crate::common;
-use crate::common::{ChunkIndex, SubmissionId};
 
 // Expected errors:
 import_exception!(opsqueue.exceptions, SubmissionFailedError);
 
 // Incorrect usage errors:
 import_exception!(opsqueue.exceptions, IncorrectUsageError);
 import_exception!(opsqueue.exceptions, TryFromIntError);
-import_exception!(opsqueue.exceptions, ChunkNotFoundError);
 import_exception!(opsqueue.exceptions, SubmissionNotFoundError);
 import_exception!(opsqueue.exceptions, SubmissionNotCancellableError);
 import_exception!(opsqueue.exceptions, NewObjectStoreClientError);
@@ -166,22 +164,6 @@ impl From<CError<crate::producer::SubmissionNotCompletedYetError>> for PyErr {
     }
 }
 
-impl From<CError<ChunkNotFound>> for PyErr {
-    fn from(value: CError<ChunkNotFound>) -> Self {
-        let ChunkId {
-            submission_id,
-            chunk_index,
-        } = value.0 .0;
-        ChunkNotFoundError::new_err((
-            value.0.to_string(),
-            (
-                SubmissionId::from(submission_id),
-                ChunkIndex::from(chunk_index),
-            ),
-        ))
-    }
-}
-
 impl From<CError<opsqueue::object_store::NewObjectStoreClientError>> for PyErr {
     fn from(value: CError<opsqueue::object_store::NewObjectStoreClientError>) -> Self {
         NewObjectStoreClientError::new_err(value.0.to_string())

opsqueue/src/common/chunk.rs

@@ -218,7 +218,7 @@ impl Chunk {
 #[cfg(feature = "server-logic")]
 pub mod db {
     use super::*;
-    use crate::common::errors::{ChunkNotFound, DatabaseError, SubmissionNotFound, E};
+    use crate::common::errors::{DatabaseError, SubmissionNotFound, E};
     use crate::db::{Connection, True, WriterConnection};
     use axum_prometheus::metrics::{counter, gauge};
     use sqlx::{query, query_as};
@@ -281,25 +281,18 @@ pub mod db {
         chunk_id: ChunkId,
         output_content: Option<Vec<u8>>,
         mut conn: impl WriterConnection,
-    ) -> Result<(), E<DatabaseError, E<SubmissionNotFound, ChunkNotFound>>> {
-        let _chunk_size: Result<ChunkSize, E<DatabaseError, E<SubmissionNotFound, ChunkNotFound>>> =
-            conn.transaction(move |mut tx| {
-                Box::pin(async move {
-                    let completed_work =
-                        complete_chunk_raw(chunk_id, output_content, &mut tx).await?;
-                    crate::common::submission::db::maybe_complete_submission(
-                        chunk_id.submission_id,
-                        &mut tx,
-                    )
-                    .await
-                    .map_err(|e| match e {
-                        E::L(e) => E::L(e),
-                        E::R(e) => E::R(E::L(e)),
-                    })?;
-                    Ok(completed_work.unwrap_or_default())
-                })
+    ) -> Result<(), E<DatabaseError, SubmissionNotFound>> {
+        conn.transaction(move |mut tx| {
+            Box::pin(async move {
+                complete_chunk_raw(chunk_id, output_content, &mut tx).await?;
+                crate::common::submission::db::maybe_complete_submission(
+                    chunk_id.submission_id,
+                    &mut tx,
+                )
+                .await
             })
-            .await;
+        })
+        .await?;
 
         counter!(crate::prometheus::CHUNKS_COMPLETED_COUNTER).increment(1);
         Ok(())
@@ -311,17 +304,16 @@ pub mod db {
         chunk_id: ChunkId,
         output_content: Option<Vec<u8>>,
         mut tx: impl WriterConnection<Transaction = True>,
-    ) -> sqlx::Result<Option<ChunkSize>> {
+    ) -> sqlx::Result<()> {
         let now = chrono::prelude::Utc::now();
-        query!(
+        let rows_affected = query!(
             "
         INSERT INTO chunks_completed
         (submission_id, chunk_index, output_content, completed_at)
         SELECT submission_id, chunk_index, $1, julianday($2) FROM chunks
         WHERE chunks.submission_id = $3 AND chunks.chunk_index = $4;
 
-        DELETE FROM chunks WHERE chunks.submission_id = $5 AND chunks.chunk_index = $6
-        RETURNING submission_id, chunk_index;
+        DELETE FROM chunks WHERE chunks.submission_id = $5 AND chunks.chunk_index = $6;
         ",
             output_content,
             now,
@@ -330,26 +322,40 @@ pub mod db {
             chunk_id.submission_id,
             chunk_id.chunk_index,
         )
-        .fetch_one(tx.get_inner())
-        .await?;
-        // Defense in depth: Above query should never be called twice on the same chunk.
-        // If it _does_ happen, it means that either a consumer is attempting a chunk they didn't reserve,
-        // or we gave out the same reservation twice.
+        .execute(tx.get_inner())
+        .await?
+        .rows_affected() as i64;
+        // Defense in depth: Above query could be called twice on the same chunk. For instance,
+        // when the server was restarted and the reservations are forgotten, and the same chunk
+        // was reserved again.
+        //
+        // By using the rows affected to update `chunks_done` we ensure that we never mess up the
+        // submission's `chunks_done` counter.
         //
-        // By returning early if the chunk was not found,
-        // we ensure that even in these situations
-        // we never mess up the submission's `chunks_done` counter.
+        // In addition, cancelling or pausing a submission while a chunk is reserved also results
+        // in `rows_affected` being 0. Which is fine, because cancelled submissions count as failed,
+        // and for paused submissions we will retry the chunk when it becomes unpaused.
         //
         // (Not doing that resulted in a hard-to-track-down bug in the past.
         // https://github.com/channable/opsqueue/issues/76
         // )
+
+        if rows_affected == 0 {
+            tracing::warn!(
+                "Could not complete chunk because it was either: completed before,\
+             or cancelled, or paused. Ignoring."
+            );
+        }
+
         sqlx::query_scalar!(
-            "UPDATE submissions SET chunks_done = chunks_done + 1 WHERE submissions.id = $1 RETURNING submissions.chunk_size;",
+            "UPDATE submissions SET chunks_done = chunks_done + $1 WHERE submissions.id = $2;",
+            rows_affected,
             chunk_id.submission_id,
         )
-        .fetch_one(tx.get_inner())
-        .await
-        .map(|opt| opt.map(ChunkSize))
+        .execute(tx.get_inner())
+        .await?;
+
+        Ok(())
     }
 
     #[tracing::instrument(skip(conn))]

opsqueue/src/common/errors.rs

@@ -12,7 +12,7 @@ use thiserror::Error;
 use crate::consumer::common::SyncServerToClientResponse;
 
 use super::{
-    chunk::{ChunkFailed, ChunkId},
+    chunk::ChunkFailed,
     submission::{SubmissionCancelled, SubmissionCompleted, SubmissionFailed, SubmissionId},
 };
 
@@ -31,10 +31,6 @@ impl<T> From<DatabaseError> for E<DatabaseError, T> {
     }
 }
 
-#[derive(Error, Debug)]
-#[error("Chunk not found for ID {0:?}")]
-pub struct ChunkNotFound(pub ChunkId);
-
 #[derive(Error, Debug, Deserialize, Serialize)]
 #[error("Submission not found for ID {0:?}")]
 pub struct SubmissionNotFound(pub SubmissionId);