fix: use oneshot channel for cancellation

Author: levydsa

libsql/src/database/builder.rs

@@ -644,7 +644,10 @@ cfg_sync! {
 
             let mut bg_abort: Option<std::sync::Arc<crate::sync::DropAbort>> = None;
 
+
             if let Some(sync_interval) = sync_interval {
+                let (cancel_tx, mut cancel_rx) = tokio::sync::oneshot::channel::<()>();
+
                 let sync_span = tracing::debug_span!("sync_interval");
                 let _enter = sync_span.enter();
 
@@ -659,7 +662,8 @@ cfg_sync! {
                 // `bootstrap_db` (for synced dbs) before calling connect. Otherwise, the sync
                 // protocol skips calling `export` endpoint causing slowdown in initial bootstrap.
                 let conn = db.connect()?;
-                let jh = tokio::spawn(
+
+                tokio::spawn(
                     async move {
                         let mut interval = tokio::time::interval(sync_interval);
 
@@ -669,21 +673,32 @@ cfg_sync! {
                             interval.tick().await;
 
                             let mut ctx = sync_ctx.lock().await;
+
                             if remote_writes {
-                                if let Err(e) = crate::sync::try_pull(&mut ctx, &conn).await {
-                                    tracing::error!("sync error: {}", e);
+                                tokio::select! {
+                                    _ = &mut cancel_rx => break,
+                                    result = crate::sync::try_pull(&mut ctx, &conn) => {
+                                        if let Err(e) = result {
+                                            tracing::error!("sync error: {}", e);
+                                        }
+                                    }
                                 }
                             } else {
-                                if let Err(e) = crate::sync::sync_offline(&mut ctx, &conn).await {
-                                    tracing::error!("sync error: {}", e);
+                                tokio::select! {
+                                    _ = &mut cancel_rx => break,
+                                    result = crate::sync::sync_offline(&mut ctx, &conn) => {
+                                        if let Err(e) = result {
+                                            tracing::error!("sync error: {}", e);
+                                        }
+                                    }
                                 }
                             }
                         }
                     }
                     .instrument(tracing::debug_span!("sync interval thread")),
                 );
 
-                bg_abort.replace(std::sync::Arc::new(crate::sync::DropAbort(jh.abort_handle())));
+                bg_abort.replace(std::sync::Arc::new(crate::sync::DropAbort(Some(cancel_tx))));
             }
 
             Ok(Database {

libsql/src/sync.rs

@@ -6,7 +6,7 @@ use bytes::Bytes;
 use chrono::Utc;
 use http::{HeaderValue, StatusCode};
 use hyper::Body;
-use tokio::{io::AsyncWriteExt as _, task::AbortHandle};
+use tokio::io::AsyncWriteExt as _;
 use uuid::Uuid;
 
 #[cfg(test)]
@@ -81,11 +81,14 @@ pub struct PushResult {
     baton: Option<String>,
 }
 
-pub struct DropAbort(pub AbortHandle);
+pub struct DropAbort(pub Option<tokio::sync::oneshot::Sender<()>>);
 
 impl Drop for DropAbort {
     fn drop(&mut self) {
-        self.0.abort();
+        tracing::debug!("aborting");
+        if let Some(sender) = self.0.take() { // Replaces with None, returns owned value
+            let _ = sender.send(()); // Now we own the sender
+        }
     }
 }
 
@@ -846,8 +849,6 @@ async fn try_push(
 
     let mut frame_no = start_frame_no;
     while frame_no <= end_frame_no {
-        tokio::task::yield_now().await;
-
         let batch_size = sync_ctx.push_batch_size.min(end_frame_no - frame_no + 1);
         let mut frames = conn.wal_get_frame(frame_no, page_size)?;
         if batch_size > 1 {
@@ -895,8 +896,6 @@ pub async fn try_pull(
     let mut err = None;
 
     loop {
-        tokio::task::yield_now().await;
-
         let generation = sync_ctx.durable_generation();
         let frame_no = sync_ctx.durable_frame_num() + 1;
         match sync_ctx.pull_one_frame(generation, frame_no).await {