From 31df09013280c62bf0ca49c1807e4e171340fd7b Mon Sep 17 00:00:00 2001
From: Rossen Stoyanchev <rstoyanchev@pivotal.io>
Date: Wed, 23 Nov 2016 20:49:24 -0500
Subject: [PATCH] Avoid locking in WebSocket session "close" callback

When processing a "close" notification from the server make an effort
to cancel any outstanding heartbeat but avoid going as far as acquiring
the responseLock since the server itself may already hold a lock of its
own leading to a potential deadlock.

The heartbeat task is now also further protected with an isClosed()
check in case the heartbeat does not get cancelled in a concurrent
scenario.

Issue: SPR-14917
---
 .../sockjs/transport/session/AbstractSockJsSession.java  | 9 +++++++--
 .../sockjs/transport/session/SockJsSessionTests.java     | 1 -
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/spring-websocket/src/main/java/org/springframework/web/socket/sockjs/transport/session/AbstractSockJsSession.java b/spring-websocket/src/main/java/org/springframework/web/socket/sockjs/transport/session/AbstractSockJsSession.java
index 7211c97049..515f55333f 100644
--- a/spring-websocket/src/main/java/org/springframework/web/socket/sockjs/transport/session/AbstractSockJsSession.java
+++ b/spring-websocket/src/main/java/org/springframework/web/socket/sockjs/transport/session/AbstractSockJsSession.java
@@ -396,7 +396,12 @@ public abstract class AbstractSockJsSession implements SockJsSession {
 		if (!isClosed()) {
 			try {
 				updateLastActiveTime();
-				cancelHeartbeat();
+				// Avoid cancelHeartbeat() and responseLock within server "close" callback
+				ScheduledFuture<?> future = this.heartbeatFuture;
+				if (future != null) {
+					this.heartbeatFuture = null;
+					future.cancel(false);
+				}
 			}
 			finally {
 				this.state = State.CLOSED;
@@ -446,7 +451,7 @@ public abstract class AbstractSockJsSession implements SockJsSession {
 		@Override
 		public void run() {
 			synchronized (responseLock) {
-				if (!this.expired) {
+				if (!this.expired && !isClosed()) {
 					try {
 						sendHeartbeat();
 					}
diff --git a/spring-websocket/src/test/java/org/springframework/web/socket/sockjs/transport/session/SockJsSessionTests.java b/spring-websocket/src/test/java/org/springframework/web/socket/sockjs/transport/session/SockJsSessionTests.java
index 8135623027..a2299ac6d5 100644
--- a/spring-websocket/src/test/java/org/springframework/web/socket/sockjs/transport/session/SockJsSessionTests.java
+++ b/spring-websocket/src/test/java/org/springframework/web/socket/sockjs/transport/session/SockJsSessionTests.java
@@ -144,7 +144,6 @@ public class SockJsSessionTests extends AbstractSockJsSessionTests<TestSockJsSes
 
 		assertClosed();
 		assertEquals(1, this.session.getNumberOfLastActiveTimeUpdates());
-		assertTrue(this.session.didCancelHeartbeat());
 		verify(this.webSocketHandler).afterConnectionClosed(this.session, CloseStatus.GOING_AWAY);
 	}