Align sleep timing tests with SDK parity

Author: manuschillerdev

src/context.rs

@@ -4,7 +4,7 @@ use crate::task_result::{
     fetch_task_result_snapshot,
 };
 use crate::types::{ClaimedTask, Json, duration_seconds_ceil, validate_queue_name};
-use chrono::{DateTime, Utc};
+use chrono::{DateTime, Timelike, Utc};
 use deadpool_postgres::Pool;
 use serde::{Serialize, de::DeserializeOwned};
 use serde_json::{Map, Value};
@@ -204,22 +204,32 @@ impl TaskContext {
     }
 
     pub async fn sleep_for(&mut self, name: impl AsRef<str>, duration: Duration) -> Result<()> {
-        self.sleep_until(
-            name,
-            Utc::now()
-                + chrono::Duration::from_std(duration).map_err(|_| {
-                    Error::Config("sleep duration is too large for chrono".to_string())
-                })?,
-        )
-        .await
+        let duration = chrono::Duration::from_std(duration)
+            .map_err(|_| Error::Config("sleep duration is too large for chrono".to_string()))?;
+        let name = name.as_ref().to_string();
+        let now = self.current_time().await?;
+        self.sleep_until_at_current_time(&name, now + duration, now)
+            .await
     }
 
     pub async fn sleep_until(
         &mut self,
         name: impl AsRef<str>,
         wake_at: DateTime<Utc>,
     ) -> Result<()> {
-        let checkpoint_name = self.next_checkpoint_name(name.as_ref());
+        let name = name.as_ref().to_string();
+        let now = self.current_time().await?;
+        self.sleep_until_at_current_time(&name, wake_at, now).await
+    }
+
+    async fn sleep_until_at_current_time(
+        &mut self,
+        name: &str,
+        wake_at: DateTime<Utc>,
+        now: DateTime<Utc>,
+    ) -> Result<()> {
+        let checkpoint_name = self.next_checkpoint_name(name);
+        let wake_at = postgres_timestamp_precision(wake_at);
         let actual_wake_at = if let Some(cached) = self.lookup_checkpoint(&checkpoint_name).await? {
             serde_json::from_value(cached)?
         } else {
@@ -228,7 +238,7 @@ impl TaskContext {
             wake_at
         };
 
-        if Utc::now() < actual_wake_at {
+        if now < actual_wake_at {
             self.schedule_run(actual_wake_at).await?;
             return Err(Error::Suspended);
         }
@@ -490,13 +500,29 @@ impl TaskContext {
         Ok(())
     }
 
+    async fn current_time(&self) -> Result<DateTime<Utc>> {
+        let client = self.pool.get().await?;
+        let row = client
+            .query_one("SELECT absurd.current_time()", &[])
+            .await
+            .map_err(map_database_error)?;
+        Ok(row.get(0))
+    }
+
     fn notify_lease_extended(&self, duration: Duration) {
         if let Some(tx) = &self.lease_tx {
             let _ = tx.send(duration);
         }
     }
 }
 
+fn postgres_timestamp_precision(timestamp: DateTime<Utc>) -> DateTime<Utc> {
+    let truncated_nanos = timestamp.timestamp_subsec_micros() * 1_000;
+    timestamp
+        .with_nanosecond(truncated_nanos)
+        .unwrap_or(timestamp)
+}
+
 async fn extend_claim(
     pool: &Pool,
     queue_name: &str,
@@ -515,6 +541,34 @@ async fn extend_claim(
     Ok(())
 }
 
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn postgres_timestamp_precision_truncates_sub_microsecond_nanos() -> Result<()> {
+        let timestamp = DateTime::from_timestamp(1_777_000_000, 123_456_789)
+            .ok_or_else(|| Error::Config("valid timestamp".to_string()))?;
+
+        let normalized = postgres_timestamp_precision(timestamp);
+
+        assert_eq!(normalized.timestamp(), timestamp.timestamp());
+        assert_eq!(normalized.timestamp_subsec_nanos(), 123_456_000);
+        Ok(())
+    }
+
+    #[test]
+    fn postgres_timestamp_precision_preserves_microsecond_values() -> Result<()> {
+        let timestamp = DateTime::from_timestamp(1_777_000_000, 123_456_000)
+            .ok_or_else(|| Error::Config("valid timestamp".to_string()))?;
+
+        let normalized = postgres_timestamp_precision(timestamp);
+
+        assert_eq!(normalized, timestamp);
+        Ok(())
+    }
+}
+
 #[derive(Clone, Debug, Default)]
 #[non_exhaustive]
 pub struct AwaitEventOptions {

tests/integration.rs

@@ -1486,44 +1486,78 @@ async fn repeated_step_names_are_numbered() -> Result<()> {
 
 #[tokio::test]
 #[ignore = "requires a Postgres database initialized with Absurd SQL"]
-async fn sleep_until_suspends_then_resumes_from_checkpoint() -> Result<()> {
-    let (queue, client) = test_client().await?;
+async fn sleep_for_suspends_until_duration_elapses() -> Result<()> {
+    let (queue, client, pool) = test_client_with_single_connection_pool().await?;
+    let base = utc("2024-05-05T10:00:00Z")?;
+    set_fake_now(&pool, Some(base)).await?;
+
+    let task = Task::<(), Value>::new("sleep-for").queue(&queue);
+    client.register(&task, |(), mut ctx| async move {
+        ctx.sleep_for("wait-for", Duration::from_secs(60)).await?;
+        Ok(json!({ "resumed": true }))
+    })?;
+
+    let spawned = client.spawn(&task, (), Default::default()).await?;
+    assert_eq!(client.work_batch(WorkBatchOptions::new()).await?, 1);
+
+    let task = fetch_task(&queue, spawned.task_id).await?;
+    let run = fetch_run(&queue, spawned.run_id).await?;
+    let wake_at = base + ChronoDuration::seconds(60);
+    assert_eq!(task.state, "sleeping");
+    assert_eq!(run.state, "sleeping");
+    assert_eq!(run.available_at, Some(wake_at));
+
+    set_fake_now(&pool, Some(wake_at + ChronoDuration::seconds(5))).await?;
+    assert_eq!(client.work_batch(WorkBatchOptions::new()).await?, 1);
+
+    let task = fetch_task(&queue, spawned.task_id).await?;
+    assert_eq!(task.state, "completed");
+    assert_eq!(task.completed_payload, Some(json!({ "resumed": true })));
+
+    client.drop_queue().await?;
+    Ok(())
+}
+
+#[tokio::test]
+#[ignore = "requires a Postgres database initialized with Absurd SQL"]
+async fn sleep_until_checkpoint_prevents_rescheduling() -> Result<()> {
+    let (queue, client, pool) = test_client_with_single_connection_pool().await?;
+    let base = utc("2024-05-06T09:00:00Z")?;
+    set_fake_now(&pool, Some(base)).await?;
+    let wake_at = base + ChronoDuration::minutes(5);
     let executions = Arc::new(AtomicUsize::new(0));
 
-    let task = Task::<(), Value>::new("sleepy").queue(&queue);
+    let task = Task::<(), Value>::new("sleep-until").queue(&queue);
     client.register(&task, {
         let executions = Arc::clone(&executions);
         move |(), mut ctx| {
             let executions = Arc::clone(&executions);
             async move {
                 let execution = executions.fetch_add(1, Ordering::SeqCst) + 1;
-                let wake_at = Utc::now() + ChronoDuration::seconds(1);
-                ctx.sleep_until("pause", wake_at).await?;
+                ctx.sleep_until("sleep-step", wake_at).await?;
                 Ok(json!({ "executions": execution }))
             }
         }
     })?;
 
     let spawned = client.spawn(&task, (), Default::default()).await?;
-
     assert_eq!(client.work_batch(WorkBatchOptions::new()).await?, 1);
 
+    let checkpoints = fetch_checkpoints(&queue, spawned.task_id).await?;
+    assert_eq!(
+        checkpoints,
+        vec![("sleep-step".to_string(), serde_json::to_value(wake_at)?)],
+    );
+
     let task = fetch_task(&queue, spawned.task_id).await?;
     let run = fetch_run(&queue, spawned.run_id).await?;
     assert_eq!(task.state, "sleeping");
     assert_eq!(run.state, "sleeping");
     assert!(run.wake_event.is_none());
     assert!(run.failure_reason.is_none());
+    assert_eq!(run.available_at, Some(wake_at));
 
-    let checkpoints = fetch_checkpoints(&queue, spawned.task_id).await?;
-    assert_eq!(checkpoints.len(), 1);
-    assert_eq!(checkpoints[0].0, "pause");
-    let checkpoint_wake: DateTime<Utc> = serde_json::from_value(checkpoints[0].1.clone())?;
-    assert_eq!(run.available_at, Some(checkpoint_wake));
-
-    assert_eq!(client.work_batch(WorkBatchOptions::new()).await?, 0);
-
-    tokio::time::sleep(Duration::from_millis(1_100)).await;
+    set_fake_now(&pool, Some(wake_at)).await?;
     assert_eq!(client.work_batch(WorkBatchOptions::new()).await?, 1);
 
     let task = fetch_task(&queue, spawned.task_id).await?;