Make complete_chunk and fail_chunk not error when processing 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,21 @@
 /// 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 +163,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,10 +218,10 @@ 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};
+    use sqlx::{query, query_as, query_scalar};
     use sqlx::{QueryBuilder, Sqlite};
 
     impl<'q> sqlx::Encode<'q, Sqlite> for super::ChunkIndex {
@@ -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,9 +304,9 @@ 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 chunk_moved = query!(
             "
         INSERT INTO chunks_completed
         (submission_id, chunk_index, output_content, completed_at)
@@ -330,26 +323,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.
+        .fetch_optional(tx.get_inner())
+        .await?
+        .is_some();
+        // 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.
+        //
+        // In addition, cancelling or pausing a submission while a chunk is reserved also results
+        // in the chunk not being in the `chunks` table. Which is fine, because cancelled
+        // submissions count as failed, and for paused submissions we will retry the chunk when it
+        // becomes unpaused.
         //
-        // 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.
+        // By only updating `chunks_done` when we actually moved a chunk, we ensure that we never
+        // mess up the submission's `chunks_done` counter.
         //
         // (Not doing that resulted in a hard-to-track-down bug in the past.
         // https://github.com/channable/opsqueue/issues/76
         // )
-        sqlx::query_scalar!(
-            "UPDATE submissions SET chunks_done = chunks_done + 1 WHERE submissions.id = $1 RETURNING submissions.chunk_size;",
-            chunk_id.submission_id,
-        )
-        .fetch_one(tx.get_inner())
-        .await
-        .map(|opt| opt.map(ChunkSize))
+        if chunk_moved {
+            sqlx::query_scalar!(
+                "UPDATE submissions SET chunks_done = chunks_done + 1 WHERE submissions.id = $1 RETURNING submissions.chunk_size;",
+                chunk_id.submission_id,
+            )
+                .fetch_one(tx.get_inner())
+                .await?;
+        } else {
+            tracing::warn!(
+                "Could not complete chunk {:?} because it was either: \
+                completed, failed, cancelled, or paused before. Ignoring.",
+                chunk_id
+            );
+        }
+
+        Ok(())
     }
 
     #[tracing::instrument(skip(conn))]
@@ -366,7 +373,7 @@ pub mod db {
                         submission_id,
                         chunk_index,
                     } = chunk_id;
-                    let fields = query!(
+                    let retries = query_scalar!(
                         "
         UPDATE chunks SET retries = retries + 1
         WHERE submission_id = $1 AND chunk_index = $2
@@ -375,24 +382,36 @@ pub mod db {
                         submission_id,
                         chunk_index
                     )
-                    .fetch_one(tx.get_inner())
+                    .fetch_optional(tx.get_inner())
                     .await?;
-                    tracing::trace!("Retries: {}", fields.retries);
-                    if fields.retries >= max_retries.into() {
-                        crate::common::submission::db::fail_submission_notx(
-                            submission_id,
-                            chunk_index,
-                            failure,
-                            &mut tx,
-                        )
-                        .await?;
-
-                        Ok::<_, sqlx::Error>(true)
-                    } else {
-                        counter!(crate::prometheus::CHUNKS_RETRIED_COUNTER).increment(1);
-                        // When retrying, the chunk re-enters ('stays') in the backlog,
-                        // so we *don't* decrement the backlog gauge here.
-                        Ok::<_, sqlx::Error>(false)
+                    tracing::trace!("Retries: {:?}", retries);
+                    match retries {
+                        Some(retries) if retries >= max_retries.into() => {
+                            crate::common::submission::db::fail_submission_notx(
+                                submission_id,
+                                chunk_index,
+                                failure,
+                                &mut tx,
+                            )
+                            .await?;
+
+                            Ok::<_, sqlx::Error>(true)
+                        }
+                        Some(_) => {
+                            counter!(crate::prometheus::CHUNKS_RETRIED_COUNTER).increment(1);
+                            // When retrying, the chunk re-enters ('stays') in the backlog,
+                            // so we *don't* decrement the backlog gauge here.
+                            Ok::<_, sqlx::Error>(false)
+                        }
+                        None => {
+                            tracing::warn!(
+                                "Could not fail chunk {:?} because it was either: \
+                                completed, failed, cancelled, or paused before. Ignoring.",
+                                chunk_id
+                            );
+
+                            Ok::<_, sqlx::Error>(false)
+                        }
                     }
                 })
             })

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);