Improve error propagation to sender in SnapshotPartitionTask

Author: pcholakov

crates/partition-store/src/partition_store_manager.rs

@@ -201,7 +201,10 @@ impl PartitionStoreManager {
         snapshot_id: SnapshotId,
         snapshot_base_path: PathBuf,
     ) -> anyhow::Result<LocalPartitionSnapshot> {
-        let mut partition_store = self.lookup.lock().await
+        let mut partition_store = self
+            .lookup
+            .lock()
+            .await
             .live
             .get_mut(&partition_id)
             .ok_or(anyhow!("Unknown partition: {}", partition_id))?

crates/worker/src/partition_processor_manager/mod.rs

@@ -304,7 +304,8 @@ impl PartitionProcessorManager {
         }
     }
 
-    #[instrument(level = "debug", skip_all, fields(partition_id = %event.partition_id, event = %<&'static str as From<&EventKind>>::from(&event.inner)))]
+    #[instrument(level = "debug", skip_all, fields(partition_id = %event.partition_id, event = %<&'static str as From<&EventKind>>::from(&event.inner)
+    ))]
     fn on_asynchronous_event(&mut self, event: AsynchronousEvent) {
         let AsynchronousEvent {
             partition_id,
@@ -543,7 +544,7 @@ impl PartitionProcessorManager {
                     status.last_archived_log_lsn = self
                         .archived_lsn_channels
                         .get(partition_id)
-                        .and_then(|(_, rx)| rx.borrow().clone());
+                        .and_then(|(_, rx)| *rx.borrow());
 
                     Some((*partition_id, status))
                 } else {
@@ -786,15 +787,8 @@ impl PartitionProcessorManager {
                 TaskKind::PartitionSnapshotProducer,
                 "create-snapshot",
                 Some(partition_id),
-                async move {
-                    create_snapshot_task
-                        .create_snapshot()
-                        .await
-                        .inspect_err(|err| {
-                            warn!("Unhandled error in create_snapshot task: {}", err)
-                        })
-                }
-                .instrument(snapshot_span),
+                async move { create_snapshot_task.create_snapshot().await }
+                    .instrument(snapshot_span),
             );
 
             match spawn_task_result {

crates/worker/src/partition_processor_manager/snapshot_task.rs

@@ -10,7 +10,6 @@
 
 use std::path::PathBuf;
 use std::time::SystemTime;
-
 use tokio::sync::{oneshot, watch};
 use tracing::{debug, warn};
 
@@ -39,49 +38,80 @@ impl SnapshotPartitionTask {
             "Creating partition snapshot"
         );
 
-        let snapshot_id = SnapshotId::new();
-        let snapshot = self
-            .partition_store_manager
-            .export_partition_snapshot(self.partition_id, snapshot_id, self.snapshot_base_path)
-            .await?;
-
-        let metadata = write_snapshot_metadata_header(
-            snapshot_id,
+        let result = create_snapshot_inner(
+            self.partition_store_manager,
             self.cluster_name,
             self.node_name,
             self.partition_id,
-            snapshot,
+            self.snapshot_base_path,
+            self.archived_lsn_sender,
         )
-        .await?;
+        .await;
 
-        // todo(pavel): SnapshotRepository integration will go in here in a future PR
+        match result {
+            Ok(metadata) => {
+                let _ = self
+                    .result_sender
+                    .send(Ok(metadata.snapshot_id))
+                    .inspect_err(|err| {
+                        warn!(
+                        "Failed to send snapshot acknowledgement after snapshot {} was successfully created: {:?}",
+                        metadata.snapshot_id,
+                        err
+                    )
+                    });
 
-        self.archived_lsn_sender
-            .send(Some(metadata.min_applied_lsn))?;
-
-        let _ = self
-            .result_sender
-            .send(Ok(metadata.snapshot_id))
-            .inspect_err(|err| {
+                debug!(
+                    partition_id = %self.partition_id,
+                    snapshot_id = %metadata.snapshot_id,
+                    archived_lsn = %metadata.min_applied_lsn,
+                    "Partition snapshot created"
+                );
+                Ok(metadata.snapshot_id)
+            }
+            Err(err) => {
                 warn!(
-                    "Failed to send snapshot acknowledgement after snapshot {} was successfully created: {:?}",
-                    metadata.snapshot_id,
+                    partition_id = %self.partition_id,
+                    "Failed to create partition snapshot: {}",
                     err
-                )
-            });
-
-        debug!(
-            partition_id = %self.partition_id,
-            snapshot_id = %metadata.snapshot_id,
-            archived_lsn = %metadata.min_applied_lsn,
-            "Partition snapshot created"
-        );
+                );
 
-        Ok(metadata.snapshot_id)
+                Err(err)
+            }
+        }
     }
 }
 
-pub async fn write_snapshot_metadata_header(
+async fn create_snapshot_inner(
+    partition_store_manager: PartitionStoreManager,
+    cluster_name: String,
+    node_name: String,
+    partition_id: PartitionId,
+    snapshot_base_path: PathBuf,
+    archived_lsn_sender: watch::Sender<Option<Lsn>>,
+) -> anyhow::Result<PartitionSnapshotMetadata> {
+    let snapshot_id = SnapshotId::new();
+    let snapshot = partition_store_manager
+        .export_partition_snapshot(partition_id, snapshot_id, snapshot_base_path.clone())
+        .await?;
+
+    let metadata = write_snapshot_metadata_header(
+        snapshot_id,
+        cluster_name,
+        node_name,
+        partition_id,
+        snapshot,
+    )
+    .await?;
+
+    // todo(pavel): SnapshotRepository integration will go in here in a future PR
+
+    archived_lsn_sender.send(Some(metadata.min_applied_lsn))?;
+
+    Ok(metadata)
+}
+
+async fn write_snapshot_metadata_header(
     snapshot_id: SnapshotId,
     cluster_name: String,
     node_name: String,