Fix accordingly to code review

Author: reda-alaoui

spring-webmvc/src/main/java/org/springframework/web/servlet/config/annotation/AsyncSupportConfigurer.java

@@ -16,6 +16,7 @@
 
 package org.springframework.web.servlet.config.annotation;
 
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -32,6 +33,7 @@
  * Helps with configuring options for asynchronous request processing.
  *
  * @author Rossen Stoyanchev
+ * @author Réda Housni Alaoui
  * @since 3.2
  */
 public class AsyncSupportConfigurer {
@@ -44,6 +46,8 @@ public class AsyncSupportConfigurer {
 
 	private final List<DeferredResultProcessingInterceptor> deferredResultInterceptors = new ArrayList<>();
 
+	private @Nullable Duration sseHeartbeatPeriod;
+
 
 	/**
 	 * The provided task executor is used for the following:
@@ -99,6 +103,14 @@ public AsyncSupportConfigurer registerDeferredResultInterceptors(
 		return this;
 	}
 
+	/**
+	 * Configure the SSE heartbeat period.
+	 * @param sseHeartbeatPeriod The SSE heartbeat period
+	 */
+	public AsyncSupportConfigurer setSseHeartbeatPeriod(Duration sseHeartbeatPeriod) {
+		this.sseHeartbeatPeriod = sseHeartbeatPeriod;
+		return this;
+	}
 
 	protected @Nullable AsyncTaskExecutor getTaskExecutor() {
 		return this.taskExecutor;
@@ -116,4 +128,8 @@ protected List<DeferredResultProcessingInterceptor> getDeferredResultInterceptor
 		return this.deferredResultInterceptors;
 	}
 
+	protected @Nullable Duration getSseHeartbeatPeriod() {
+		return this.sseHeartbeatPeriod;
+	}
+
 }

spring-webmvc/src/main/java/org/springframework/web/servlet/mvc/method/annotation/DefaultSseEmitterHeartbeatExecutor.java

@@ -17,11 +17,13 @@
 package org.springframework.web.servlet.mvc.method.annotation;
 
 import java.lang.reflect.Method;
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ConcurrentHashMap;
@@ -54,6 +56,8 @@
 import org.springframework.http.converter.HttpMessageConverter;
 import org.springframework.http.converter.StringHttpMessageConverter;
 import org.springframework.http.converter.support.AllEncompassingFormHttpMessageConverter;
+import org.springframework.scheduling.TaskScheduler;
+import org.springframework.scheduling.concurrent.SimpleAsyncTaskScheduler;
 import org.springframework.ui.ModelMap;
 import org.springframework.util.Assert;
 import org.springframework.util.ClassUtils;
@@ -202,8 +206,9 @@ public class RequestMappingHandlerAdapter extends AbstractHandlerMethodAdapter
 
 	private final Map<ControllerAdviceBean, Set<Method>> modelAttributeAdviceCache = new LinkedHashMap<>();
 
-	@Nullable
-	private SseEmitterHeartbeatExecutor sseEmitterHeartbeatExecutor;
+	private TaskScheduler taskScheduler = new SimpleAsyncTaskScheduler();
+
+	private @Nullable Duration sseHeartbeatPeriod;
 
 	/**
 	 * Provide resolvers for custom argument types. Custom resolvers are ordered
@@ -530,10 +535,17 @@ public void setParameterNameDiscoverer(ParameterNameDiscoverer parameterNameDisc
 	}
 
 	/**
-	 * Set the {@link SseEmitterHeartbeatExecutor} that will be used to periodically prob the SSE connection health
+	 * Set the {@link TaskScheduler}
+	 */
+	public void setTaskScheduler(TaskScheduler taskScheduler) {
+		this.taskScheduler = taskScheduler;
+	}
+
+	/**
+	 * Sets the heartbeat period that will be used to periodically prob the SSE connection health
 	 */
-	public void setSseEmitterHeartbeatExecutor(@Nullable SseEmitterHeartbeatExecutor sseEmitterHeartbeatExecutor) {
-		this.sseEmitterHeartbeatExecutor = sseEmitterHeartbeatExecutor;
+	public void setSseHeartbeatPeriod(@Nullable Duration sseHeartbeatPeriod) {
+		this.sseHeartbeatPeriod = sseHeartbeatPeriod;
 	}
 
 	/**
@@ -743,9 +755,12 @@ private List<HandlerMethodReturnValueHandler> getDefaultReturnValueHandlers() {
 		handlers.add(new ModelAndViewMethodReturnValueHandler());
 		handlers.add(new ModelMethodProcessor());
 		handlers.add(new ViewMethodReturnValueHandler());
+
+		SseEmitterHeartbeatExecutor sseEmitterHeartbeatExecutor = Optional.ofNullable(sseHeartbeatPeriod)
+				.map(period -> new SseEmitterHeartbeatExecutor(taskScheduler, period)).orElse(null);
 		handlers.add(new ResponseBodyEmitterReturnValueHandler(getMessageConverters(),
 				this.reactiveAdapterRegistry, this.taskExecutor, this.contentNegotiationManager,
-				initViewResolvers(), initLocaleResolver(), this.sseEmitterHeartbeatExecutor));
+				initViewResolvers(), initLocaleResolver(), sseEmitterHeartbeatExecutor));
 		handlers.add(new StreamingResponseBodyReturnValueHandler());
 		handlers.add(new HttpEntityMethodProcessor(getMessageConverters(),
 				this.contentNegotiationManager, this.requestResponseBodyAdvice, this.errorResponseInterceptors));

spring-webmvc/src/main/java/org/springframework/web/servlet/mvc/method/annotation/RequestMappingHandlerAdapter.java

@@ -18,14 +18,18 @@
 
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
+import java.time.Duration;
 import java.util.Collections;
 import java.util.LinkedHashSet;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
 
 import org.jspecify.annotations.Nullable;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.MediaType;
 import org.springframework.http.server.ServerHttpResponse;
@@ -41,17 +45,22 @@
  * @author Juergen Hoeller
  * @author Sam Brannen
  * @author Brian Clozel
+ * @author Réda Housni Alaoui
  * @since 4.2
  */
 public class SseEmitter extends ResponseBodyEmitter {
 
+	private static final Logger LOGGER = LoggerFactory.getLogger(SseEmitter.class);
+
 	private static final MediaType TEXT_PLAIN = new MediaType("text", "plain", StandardCharsets.UTF_8);
 
 	/**
 	 * Guards access to write operations on the response.
 	 */
 	private final Lock writeLock = new ReentrantLock();
 
+	private volatile @Nullable Long lastEmissionNanoTime;
+
 	/**
 	 * Create a new SseEmitter instance.
 	 */
@@ -134,12 +143,31 @@ public void send(SseEventBuilder builder) throws IOException {
 		this.writeLock.lock();
 		try {
 			super.send(dataToSend);
+			this.lastEmissionNanoTime = System.nanoTime();
 		}
 		finally {
 			this.writeLock.unlock();
 		}
 	}
 
+	void notifyOfHeartbeatTick(Duration heartbeatPeriod) {
+		boolean skip = Optional.ofNullable(lastEmissionNanoTime)
+				.map(lastEmissionNanoTime -> System.nanoTime() - lastEmissionNanoTime)
+				.map(nanoTimeElapsedSinceLastEmission -> nanoTimeElapsedSinceLastEmission < heartbeatPeriod.toNanos())
+				.orElse(false);
+		if (skip) {
+			return;
+		}
+		LOGGER.trace("Sending heartbeat to {}", this);
+		SseEmitter.SseEventBuilder eventBuilder = SseEmitter.event().name("ping").data("ping", MediaType.TEXT_PLAIN);
+		try {
+			send(eventBuilder);
+		} catch (IOException | RuntimeException e) {
+			// According to SseEmitter's Javadoc, the container itself will call SseEmitter#completeWithError
+			LOGGER.debug(e.getMessage());
+		}
+	}
+
 	@Override
 	public String toString() {
 		return "SseEmitter@" + ObjectUtils.getIdentityHexString(this);

spring-webmvc/src/main/java/org/springframework/web/servlet/mvc/method/annotation/SseEmitter.java

@@ -16,10 +16,74 @@
 
 package org.springframework.web.servlet.mvc.method.annotation;
 
+
+import java.time.Duration;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ScheduledFuture;
+
+import org.jspecify.annotations.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.scheduling.TaskScheduler;
+
 /**
  * @author Réda Housni Alaoui
  */
-public interface SseEmitterHeartbeatExecutor {
+class SseEmitterHeartbeatExecutor {
+
+	private static final Logger LOGGER = LoggerFactory.getLogger(SseEmitterHeartbeatExecutor.class);
+
+	private final TaskScheduler taskScheduler;
+	private final Set<SseEmitter> emitters = ConcurrentHashMap.newKeySet();
+
+	private final Object lifecycleMonitor = new Object();
+
+	private final Duration period;
+
+	@Nullable
+	private volatile ScheduledFuture<?> taskFuture;
+
+	public SseEmitterHeartbeatExecutor(TaskScheduler taskScheduler, Duration period) {
+		this.taskScheduler = taskScheduler;
+		this.period = period;
+	}
+
+	public void register(SseEmitter emitter) {
+		startIfNeeded();
+
+		Runnable closeCallback = () -> emitters.remove(emitter);
+		emitter.onCompletion(closeCallback);
+		emitter.onError(t -> closeCallback.run());
+		emitter.onTimeout(closeCallback);
+
+		emitters.add(emitter);
+	}
+
+	boolean isRegistered(SseEmitter emitter) {
+		return emitters.contains(emitter);
+	}
+
+	private void startIfNeeded() {
+		if (taskFuture != null) {
+			return;
+		}
+		synchronized (lifecycleMonitor) {
+			if (taskFuture != null) {
+				return;
+			}
+			taskFuture = taskScheduler.scheduleAtFixedRate(this::notifyEmitters, period);
+		}
+	}
+
+	private void notifyEmitters() {
+		LOGGER.atDebug().log(() -> "Notifying %s emitter(s)".formatted(emitters.size()));
 
-	void register(SseEmitter emitter);
+		for (SseEmitter emitter : emitters) {
+			if (Thread.currentThread().isInterrupted()) {
+				return;
+			}
+			emitter.notifyOfHeartbeatTick(period);
+		}
+	}
 }