Add heartbeat lock to SockJS server sessions

Even before this change SockJS sessions always cancelled the heartbeat
task first prior to sending messages. However when the heartbeat task
is already in progress, cancellation of it is not enough and we must
wait until the heartbeat is sent.

This commit adds a heartbeat write lock which is obtained and held
during the sending of a heartbeat. Now when sessions send a message
they still cancel the heartbeat task but if that fails they also wait
for the heartbeat write lock.

Issue: SPR-14356

Author: rstoyanchev

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2015 the original author or authors.
+ * Copyright 2002-2016 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -27,6 +27,8 @@
 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,6 +108,8 @@ private enum State {NEW, OPEN, CLOSED}
 
 	private volatile ScheduledFuture<?> heartbeatTask;
 
+	private final Lock heartbeatLock = new ReentrantLock();
+
 	private volatile boolean heartbeatDisabled;
 
 
@@ -246,8 +250,15 @@ public void disableHeartbeat() {
 
 	public void sendHeartbeat() throws SockJsTransportFailureException {
 		if (isActive()) {
-			writeFrame(SockJsFrame.heartbeatFrame());
-			scheduleHeartbeat();
+			if (heartbeatLock.tryLock()) {
+				try {
+					writeFrame(SockJsFrame.heartbeatFrame());
+					scheduleHeartbeat();
+				}
+				finally {
+					heartbeatLock.unlock();
+				}
+			}
 		}
 	}
 
@@ -282,13 +293,26 @@ protected void cancelHeartbeat() {
 		try {
 			ScheduledFuture<?> task = this.heartbeatTask;
 			this.heartbeatTask = null;
+			if (task == null || task.isCancelled()) {
+				return;
+			}
 
-			if ((task != null) && !task.isDone()) {
-				if (logger.isTraceEnabled()) {
-					logger.trace("Cancelling heartbeat in session " + getId());
-				}
-				task.cancel(false);
+			if (logger.isTraceEnabled()) {
+				logger.trace("Cancelling heartbeat in session " + getId());
+			}
+			if (task.cancel(false)) {
+				return;
+			}
+
+			if (logger.isTraceEnabled()) {
+				logger.trace("Failed to cancel heartbeat, acquiring heartbeat write lock.");
+			}
+			this.heartbeatLock.lock();
+
+			if (logger.isTraceEnabled()) {
+				logger.trace("Releasing heartbeat lock.");
 			}
+			this.heartbeatLock.unlock();
 		}
 		catch (Throwable ex) {
 			logger.debug("Failure while cancelling heartbeat in session " + getId(), ex);

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2014 the original author or authors.
+ * Copyright 2002-2016 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -287,11 +287,12 @@ public void scheduleAndCancelHeartbeat() throws Exception {
 		verify(this.taskScheduler).schedule(any(Runnable.class), any(Date.class));
 		verifyNoMoreInteractions(this.taskScheduler);
 
-		given(task.isDone()).willReturn(false);
+		given(task.isCancelled()).willReturn(false);
+		given(task.cancel(false)).willReturn(true);
 
 		this.session.cancelHeartbeat();
 
-		verify(task).isDone();
+		verify(task).isCancelled();
 		verify(task).cancel(false);
 		verifyNoMoreInteractions(task);
 	}