From 0c688742e102faebcb18bed1b84fe84f635b40d6 Mon Sep 17 00:00:00 2001
From: Brian Clozel <brian.clozel@broadcom.com>
Date: Tue, 10 Dec 2024 22:30:21 +0100
Subject: [PATCH] Fix custom scheduler support for @Scheduled methods

This commit fixes a regression introduced by gh-24560, when adding
execution metadata support for scheduled tasks. The
`OutcomeTrackingRunnable` would delegate to the actual runnable but
could also hide whether it implements the `SchedulingAwareRunnable`
contract.

This commit ensures that `OutcomeTrackingRunnable` always implements
that contract and delegates to the runnable if possible, or return
default values otherwise.

Fixes gh-34058
---
 .../scheduling/config/Task.java               | 21 ++++++++++++++++++-
 .../scheduling/config/TaskTests.java          | 16 ++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/spring-context/src/main/java/org/springframework/scheduling/config/Task.java b/spring-context/src/main/java/org/springframework/scheduling/config/Task.java
index d1a3373fa8d..ae14768a7d8 100644
--- a/spring-context/src/main/java/org/springframework/scheduling/config/Task.java
+++ b/spring-context/src/main/java/org/springframework/scheduling/config/Task.java
@@ -18,6 +18,8 @@ package org.springframework.scheduling.config;
 
 import java.time.Instant;
 
+import org.springframework.lang.Nullable;
+import org.springframework.scheduling.SchedulingAwareRunnable;
 import org.springframework.util.Assert;
 
 /**
@@ -68,7 +70,7 @@ public class Task {
 	}
 
 
-	private class OutcomeTrackingRunnable implements Runnable {
+	private class OutcomeTrackingRunnable implements SchedulingAwareRunnable {
 
 		private final Runnable runnable;
 
@@ -89,6 +91,23 @@ public class Task {
 			}
 		}
 
+		@Override
+		public boolean isLongLived() {
+			if (this.runnable instanceof SchedulingAwareRunnable sar) {
+				return sar.isLongLived();
+			}
+			return SchedulingAwareRunnable.super.isLongLived();
+		}
+
+		@Nullable
+		@Override
+		public String getQualifier() {
+			if (this.runnable instanceof SchedulingAwareRunnable sar) {
+				return sar.getQualifier();
+			}
+			return SchedulingAwareRunnable.super.getQualifier();
+		}
+
 		@Override
 		public String toString() {
 			return this.runnable.toString();
diff --git a/spring-context/src/test/java/org/springframework/scheduling/config/TaskTests.java b/spring-context/src/test/java/org/springframework/scheduling/config/TaskTests.java
index a122978625a..1540443e503 100644
--- a/spring-context/src/test/java/org/springframework/scheduling/config/TaskTests.java
+++ b/spring-context/src/test/java/org/springframework/scheduling/config/TaskTests.java
@@ -16,8 +16,12 @@
 
 package org.springframework.scheduling.config;
 
+import io.micrometer.observation.tck.TestObservationRegistry;
 import org.junit.jupiter.api.Test;
 
+import org.springframework.scheduling.SchedulingAwareRunnable;
+import org.springframework.scheduling.support.ScheduledMethodRunnable;
+
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
 import static org.assertj.core.api.Assertions.assertThatIllegalStateException;
@@ -72,6 +76,18 @@ class TaskTests {
 		assertThat(executionOutcome.throwable()).isInstanceOf(IllegalStateException.class);
 	}
 
+	@Test
+	void shouldDelegateToSchedulingAwareRunnable() throws Exception {
+		ScheduledMethodRunnable methodRunnable = new ScheduledMethodRunnable(new TestRunnable(),
+				TestRunnable.class.getMethod("run"), "myScheduler", TestObservationRegistry::create);
+		Task task = new Task(methodRunnable);
+
+		assertThat(task.getRunnable()).isInstanceOf(SchedulingAwareRunnable.class);
+		SchedulingAwareRunnable actual = (SchedulingAwareRunnable) task.getRunnable();
+		assertThat(actual.getQualifier()).isEqualTo(methodRunnable.getQualifier());
+		assertThat(actual.isLongLived()).isEqualTo(methodRunnable.isLongLived());
+	}
+
 
 	static class TestRunnable implements Runnable {