Rework guard to detect when dropped without notifying. Add to log collection to tests to ensure that the message fires.

Author: VegetarianOrc

crates/sdk-core/src/core_tests/workers.rs

@@ -62,6 +62,7 @@ use temporalio_common::{
             },
         },
     },
+    telemetry::{CoreTelemetry, Logger, TelemetryOptions, telemetry_init},
     worker::WorkerTaskTypes,
 };
 use tokio::sync::{Barrier, Notify, watch};
@@ -1334,6 +1335,16 @@ async fn graceful_shutdown_sends_shutdown_worker_rpc_during_initiate() {
 /// is able to trigger PollShutdown.
 #[tokio::test]
 async fn nexus_shutdown_does_not_hang_when_pending_completion_is_cancelled() {
+    let telemetry = telemetry_init(
+        TelemetryOptions::builder()
+            .logging(Logger::Forward {
+                filter: "OFF,temporalio_sdk_core::worker::nexus=ERROR".to_string(),
+            })
+            .build(),
+    )
+    .unwrap();
+    let _guard = tracing::subscriber::set_default(telemetry.trace_subscriber().unwrap().clone());
+
     let mut client = mock_manual_worker_client();
     let completion_rpc_started = Arc::new(Barrier::new(2));
     let completion_rpc_started_clone = completion_rpc_started.clone();
@@ -1398,6 +1409,18 @@ async fn nexus_shutdown_does_not_hang_when_pending_completion_is_cancelled() {
     }
     drop(completion);
 
+    let logs = telemetry.fetch_buffered_logs();
+    assert!(
+        logs.iter().any(|log| {
+            log.level == tracing::Level::ERROR
+                && log.target == "temporalio_sdk_core::worker::nexus"
+                && log
+                    .message
+                    .starts_with("TaskCompletedGuard triggered notify on drop")
+        }),
+        "expected TaskCompletedGuard error log, got {logs:#?}"
+    );
+
     // Polling again should now return PollError::ShutDown because the outstanding task map is empty
     // and waiters should have been notified.
     let poll_result = tokio::time::timeout(Duration::from_secs(1), poll_future.as_mut())

crates/sdk-core/src/worker/nexus.rs

@@ -120,8 +120,7 @@ impl NexusManager {
     ) -> Result<(), CompleteNexusError> {
         let removed = self.outstanding_task_map.lock().remove(&tt);
         if let Some(task_info) = removed {
-            let task_completed_notify =
-                TaskCompletedNotify::new(self.task_completed_notify.clone());
+            let task_completed_notify = TaskCompletedGuard::new(self.task_completed_notify.clone());
             self.metrics
                 .nexus_task_execution_latency(task_info.start_time.elapsed());
             task_info.timeout_task.inspect(|jh| jh.abort());
@@ -331,31 +330,34 @@ impl NexusManager {
 /// TaskCompleteNotify is used to ensure that waiters are notified when a task
 /// is removed from the outstanding task map even in the event that the running
 /// future is dropped.
-struct TaskCompletedNotify {
+struct TaskCompletedGuard {
     inner: Option<Arc<Notify>>,
 }
 
-impl TaskCompletedNotify {
+impl TaskCompletedGuard {
     fn new(notify: Arc<Notify>) -> Self {
         Self {
             inner: Some(notify),
         }
     }
 
-    fn notify_inner(&mut self) {
+    fn notify_waiters(mut self) {
         if let Some(notify) = self.inner.take() {
             notify.notify_waiters();
         }
     }
-
-    fn notify_waiters(mut self) {
-        self.notify_inner();
-    }
 }
 
-impl Drop for TaskCompletedNotify {
+impl Drop for TaskCompletedGuard {
     fn drop(&mut self) {
-        self.notify_inner();
+        if let Some(notify) = self.inner.take() {
+            error!(
+                "TaskCompletedGuard triggered notify on drop. This indicates that the caller has \
+                dropped the future for `complete_task`. \
+                This should not happen. Please file a bug report."
+            );
+            notify.notify_waiters();
+        }
     }
 }