From c74ecff3f0a684dc044acb594247d1ecb3192e5c Mon Sep 17 00:00:00 2001
From: Andy Balaam <andy.balaam@matrix.org>
Date: Mon, 24 Feb 2025 16:08:46 +0000
Subject: [PATCH] refactor(timeline): Move finding retry indices into
 DecryptionRetryTask

---
 .../controller/decryption_retry_task.rs       | 51 ++++++++++++++++---
 .../src/timeline/controller/mod.rs            | 38 +-------------
 2 files changed, 46 insertions(+), 43 deletions(-)

diff --git a/crates/matrix-sdk-ui/src/timeline/controller/decryption_retry_task.rs b/crates/matrix-sdk-ui/src/timeline/controller/decryption_retry_task.rs
index 5407ce072..b50242da7 100644
--- a/crates/matrix-sdk-ui/src/timeline/controller/decryption_retry_task.rs
+++ b/crates/matrix-sdk-ui/src/timeline/controller/decryption_retry_task.rs
@@ -16,11 +16,7 @@ use std::{collections::BTreeSet, sync::Arc};
 
 use matrix_sdk::deserialized_responses::TimelineEventKind as SdkTimelineEventKind;
 use tokio::sync::RwLock;
-use tracing::{
-    error,
-    field::{self, debug},
-    info, info_span, Instrument as _,
-};
+use tracing::{debug, error, field, info, info_span, Instrument as _};
 
 use crate::timeline::{
     controller::{TimelineSettings, TimelineState},
@@ -50,6 +46,48 @@ impl<P: RoomDataProvider> DecryptionRetryTask<P> {
         &self,
         decryptor: impl Decryptor,
         session_ids: Option<BTreeSet<String>>,
+    ) {
+        let state = self.state.clone().read_owned().await;
+
+        let should_retry = |session_id: &str| {
+            if let Some(session_ids) = &session_ids {
+                session_ids.contains(session_id)
+            } else {
+                true
+            }
+        };
+
+        let retry_indices: Vec<_> = state
+            .items
+            .iter()
+            .enumerate()
+            .filter_map(|(idx, item)| match item.as_event()?.content().as_unable_to_decrypt()? {
+                EncryptedMessage::MegolmV1AesSha2 { session_id, .. }
+                    if should_retry(session_id) =>
+                {
+                    Some(idx)
+                }
+                EncryptedMessage::MegolmV1AesSha2 { .. }
+                | EncryptedMessage::OlmV1Curve25519AesSha2 { .. }
+                | EncryptedMessage::Unknown => None,
+            })
+            .collect();
+
+        if retry_indices.is_empty() {
+            return;
+        }
+
+        drop(state);
+
+        debug!("Retrying decryption");
+
+        self.decrypt_by_index(decryptor, session_ids, retry_indices).await;
+    }
+
+    async fn decrypt_by_index(
+        &self,
+        decryptor: impl Decryptor,
+        session_ids: Option<BTreeSet<String>>,
         retry_indices: Vec<usize>,
     ) {
         let should_retry = move |session_id: &str| {
@@ -93,7 +131,8 @@ impl<P: RoomDataProvider> DecryptionRetryTask<P> {
                         return None;
                     };
 
-                    tracing::Span::current().record("event_id", debug(&remote_event.event_id));
+                    tracing::Span::current()
+                        .record("event_id", field::debug(&remote_event.event_id));
 
                     let Some(original_json) = &remote_event.original_json else {
                         error!("UTD item must contain original JSON");
diff --git a/crates/matrix-sdk-ui/src/timeline/controller/mod.rs b/crates/matrix-sdk-ui/src/timeline/controller/mod.rs
index cbe78d0c9..d0e143cc4 100644
--- a/crates/matrix-sdk-ui/src/timeline/controller/mod.rs
+++ b/crates/matrix-sdk-ui/src/timeline/controller/mod.rs
@@ -1062,49 +1062,13 @@ impl<P: RoomDataProvider> TimelineController<P> {
         decryptor: impl Decryptor,
         session_ids: Option<BTreeSet<String>>,
     ) {
-        use super::EncryptedMessage;
-
-        let state = self.state.clone().read_owned().await;
-
-        let should_retry = |session_id: &str| {
-            if let Some(session_ids) = &session_ids {
-                session_ids.contains(session_id)
-            } else {
-                true
-            }
-        };
-
-        let retry_indices: Vec<_> = state
-            .items
-            .iter()
-            .enumerate()
-            .filter_map(|(idx, item)| match item.as_event()?.content().as_unable_to_decrypt()? {
-                EncryptedMessage::MegolmV1AesSha2 { session_id, .. }
-                    if should_retry(session_id) =>
-                {
-                    Some(idx)
-                }
-                EncryptedMessage::MegolmV1AesSha2 { .. }
-                | EncryptedMessage::OlmV1Curve25519AesSha2 { .. }
-                | EncryptedMessage::Unknown => None,
-            })
-            .collect();
-
-        if retry_indices.is_empty() {
-            return;
-        }
-
-        drop(state);
-
-        debug!("Retrying decryption");
-
         let decryption_retry_task = DecryptionRetryTask::new(
             self.state.clone(),
             self.settings.clone(),
             self.room_data_provider.clone(),
         );
 
-        decryption_retry_task.decrypt(decryptor, session_ids, retry_indices).await;
+        decryption_retry_task.decrypt(decryptor, session_ids).await;
     }
 
     pub(super) async fn set_sender_profiles_pending(&self) {