Harden synchronization around SockJS heartbeats

Create an explicit heartbeat task with an experiration flag so that
it can be cancelled reliably vs relying on the ScheduledFutureTask
cancel method which may return true even if the task is already
running.

Issue: SPR-14356

Author: rstoyanchev

spring-websocket/src/main/java/org/springframework/web/socket/sockjs/transport/session/AbstractSockJsSession.java

@@ -27,8 +27,6 @@
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ScheduledFuture;
-import java.util.concurrent.locks.Lock;
-import java.util.concurrent.locks.ReentrantLock;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -106,9 +104,11 @@ private enum State {NEW, OPEN, CLOSED}
 
 	private volatile long timeLastActive = this.timeCreated;
 
-	private volatile ScheduledFuture<?> heartbeatTask;
+	private ScheduledFuture<?> heartbeatFuture;
 
-	private final Lock heartbeatLock = new ReentrantLock();
+	private HeartbeatTask heartbeatTask;
+
+	private final Object heartbeatLock = new Object();
 
 	private volatile boolean heartbeatDisabled;
 
@@ -249,19 +249,10 @@ public void disableHeartbeat() {
 	}
 
 	public void sendHeartbeat() throws SockJsTransportFailureException {
-		if (isActive()) {
-			if (heartbeatLock.tryLock()) {
-				try {
-					if (this.heartbeatTask == null) {
-						// Cancelled while waiting to acquire the lock
-						return;
-					}
-					writeFrame(SockJsFrame.heartbeatFrame());
-					scheduleHeartbeat();
-				}
-				finally {
-					heartbeatLock.unlock();
-				}
+		synchronized (this.heartbeatLock) {
+			if (isActive() && !this.heartbeatDisabled) {
+				writeFrame(SockJsFrame.heartbeatFrame());
+				scheduleHeartbeat();
 			}
 		}
 	}
@@ -270,56 +261,33 @@ protected void scheduleHeartbeat() {
 		if (this.heartbeatDisabled) {
 			return;
 		}
-
-		Assert.state(this.config.getTaskScheduler() != null, "Expected SockJS TaskScheduler");
-		cancelHeartbeat();
-		if (!isActive()) {
-			return;
-		}
-
-		Date time = new Date(System.currentTimeMillis() + this.config.getHeartbeatTime());
-		this.heartbeatTask = this.config.getTaskScheduler().schedule(new Runnable() {
-			public void run() {
-				try {
-					sendHeartbeat();
-				}
-				catch (Throwable ex) {
-					// ignore
-				}
-			}
-		}, time);
-		if (logger.isTraceEnabled()) {
-			logger.trace("Scheduled heartbeat in session " + getId());
-		}
-	}
-
-	protected void cancelHeartbeat() {
-		try {
-			ScheduledFuture<?> task = this.heartbeatTask;
-			this.heartbeatTask = null;
-			if (task == null || task.isCancelled()) {
+		synchronized (this.heartbeatLock) {
+			cancelHeartbeat();
+			if (!isActive()) {
 				return;
 			}
-
+			Date time = new Date(System.currentTimeMillis() + this.config.getHeartbeatTime());
+			this.heartbeatTask = new HeartbeatTask();
+			this.heartbeatFuture = this.config.getTaskScheduler().schedule(this.heartbeatTask, time);
 			if (logger.isTraceEnabled()) {
-				logger.trace("Cancelling heartbeat in session " + getId());
-			}
-			if (task.cancel(false)) {
-				return;
+				logger.trace("Scheduled heartbeat in session " + getId());
 			}
+		}
+	}
 
-			if (logger.isTraceEnabled()) {
-				logger.trace("Failed to cancel heartbeat, acquiring heartbeat write lock.");
+	protected void cancelHeartbeat() {
+		synchronized (this.heartbeatLock) {
+			if (this.heartbeatFuture != null) {
+				if (logger.isTraceEnabled()) {
+					logger.trace("Cancelling heartbeat in session " + getId());
+				}
+				this.heartbeatFuture.cancel(false);
+				this.heartbeatFuture = null;
 			}
-			this.heartbeatLock.lock();
-
-			if (logger.isTraceEnabled()) {
-				logger.trace("Releasing heartbeat lock.");
+			if (this.heartbeatTask != null) {
+				this.heartbeatTask.cancel();
+				this.heartbeatTask = null;
 			}
-			this.heartbeatLock.unlock();
-		}
-		catch (Throwable ex) {
-			logger.debug("Failure while cancelling heartbeat in session " + getId(), ex);
 		}
 	}
 
@@ -469,4 +437,28 @@ public String toString() {
 		return getClass().getSimpleName() + "[id=" + getId() + "]";
 	}
 
+
+	private class HeartbeatTask implements Runnable {
+
+		private boolean expired;
+
+		@Override
+		public void run() {
+			synchronized (heartbeatLock) {
+				if (!this.expired) {
+					try {
+						sendHeartbeat();
+					}
+					finally {
+						this.expired = true;
+					}
+				}
+			}
+		}
+
+		void cancel() {
+			this.expired = true;
+		}
+	}
+
 }

spring-websocket/src/test/java/org/springframework/web/socket/sockjs/transport/session/SockJsSessionTests.java

@@ -288,6 +288,7 @@ public void scheduleHeartbeatNotActive() throws Exception {
 	@Test
 	public void sendHeartbeatWhenDisabled() throws Exception {
 		this.session.disableHeartbeat();
+		this.session.setActive(true);
 		this.session.sendHeartbeat();
 
 		assertEquals(Collections.emptyList(), this.session.getSockJsFramesWritten());
@@ -310,7 +311,6 @@ public void scheduleAndCancelHeartbeat() throws Exception {
 
 		this.session.cancelHeartbeat();
 
-		verify(task).isCancelled();
 		verify(task).cancel(false);
 		verifyNoMoreInteractions(task);
 	}