From 6679feb77b49ef3f779f1dcfb40b5ee25c5d9f76 Mon Sep 17 00:00:00 2001
From: Andy Wilkinson <awilkinson@gopivotal.com>
Date: Thu, 26 Sep 2013 15:08:34 +0100
Subject: [PATCH] Improve handling of missed heartbeats

Previously, when a broker heartbeat was mnissed, the STOMP connection
would be left in a semi-disconnected state such that, for example, the
read and write idle callbacks would still be active, even though
the underlying TCP connection had been nulled out.

As part of disconnecting the STOMP connection, this commit closes the
underlying TCP connection when a heartbeat's missed which cancels the
read and write idle callbacks. It also now copes with the underlying
TCP connection being null when sending a heartbeat to the broker. This
protects again a race condition between the write idle callback being
fired, such that a heartbeat needs to be sent, and the connection
being nulled out due to it being closed.
---
 .../simp/stomp/StompBrokerRelayMessageHandler.java   | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompBrokerRelayMessageHandler.java b/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompBrokerRelayMessageHandler.java
index 4596be012bf..b65a032d726 100644
--- a/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompBrokerRelayMessageHandler.java
+++ b/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompBrokerRelayMessageHandler.java
@@ -438,7 +438,12 @@ public class StompBrokerRelayMessageHandler extends AbstractBrokerMessageHandler
 
 		public void setDisconnected() {
 			this.readyConnection.set(null);
-			this.connection = null;
+
+			TcpConnection<Message<byte[]>, Message<byte[]>> localConnection = this.connection;
+			if (localConnection != null) {
+				localConnection.close();
+				this.connection = null;
+			}
 		}
 
 		@Override
@@ -499,7 +504,10 @@ public class StompBrokerRelayMessageHandler extends AbstractBrokerMessageHandler
 
 					@Override
 					public void run() {
-						stompConnection.connection.send(MessageBuilder.withPayload(heartbeatPayload).build());
+						TcpConnection<Message<byte[]>, Message<byte[]>> connection = stompConnection.connection;
+						if (connection != null) {
+							connection.send(MessageBuilder.withPayload(heartbeatPayload).build());
+						}
 					}
 
 				});