Use tokio::task::coop::consume_budget instead of manually yielding

Author: sandhose

crates/syn2mas/src/migration.rs

@@ -273,14 +273,9 @@ async fn migrate_users(
             let mut user_buffer = MasWriteBuffer::new(&mas, MasWriter::write_users);
             let mut password_buffer = MasWriteBuffer::new(&mas, MasWriter::write_passwords);
 
-            let mut i = 0;
             while let Ok(user) = rx.recv_async().await {
-                // Force yielding to the scheduler every once in a while, else we'll have long
-                // tick times
-                i += 1;
-                if i % 128 == 0 {
-                    tokio::task::yield_now().await;
-                }
+                // Consume some coop budget to give it a chance to yield
+                tokio::task::coop::consume_budget().await;
 
                 // Handling an edge case: some AS users may have invalid localparts containing
                 // extra `:` characters. These users are ignored and a warning is logged.
@@ -622,14 +617,9 @@ async fn migrate_devices(
         async move {
             let mut write_buffer = MasWriteBuffer::new(&mas, MasWriter::write_compat_sessions);
 
-            let mut i = 0;
             while let Ok(device) = rx.recv_async().await {
-                // Force yielding to the scheduler every once in a while, else we'll have long
-                // tick times
-                i += 1;
-                if i % 128 == 0 {
-                    tokio::task::yield_now().await;
-                }
+                // Consume some coop budget to give it a chance to yield
+                tokio::task::coop::consume_budget().await;
 
                 let SynapseDevice {
                     user_id: synapse_user_id,
@@ -783,14 +773,9 @@ async fn migrate_unrefreshable_access_tokens(
             let mut deviceless_session_write_buffer =
                 MasWriteBuffer::new(&mas, MasWriter::write_compat_sessions);
 
-            let mut i = 0;
             while let Ok(token) = rx.recv_async().await {
-                // Force yielding to the scheduler every once in a while, else we'll have long
-                // tick times
-                i += 1;
-                if i % 128 == 0 {
-                    tokio::task::yield_now().await;
-                }
+                // Consume some coop budget to give it a chance to yield
+                tokio::task::coop::consume_budget().await;
 
                 let SynapseAccessToken {
                     user_id: synapse_user_id,