Support heartbeat in SimpleBrokerMessageHandler

Issue: SPR-10954

spring-messaging/src/main/java/org/springframework/messaging/simp/SimpMessageHeaderAccessor.java

@@ -51,8 +51,6 @@ public class SimpMessageHeaderAccessor extends NativeMessageHeaderAccessor {
 
 	// SiMP header names
 
-	public static final String CONNECT_MESSAGE_HEADER = "simpConnectMessage";
-
 	public static final String DESTINATION_HEADER = "simpDestination";
 
 	public static final String MESSAGE_TYPE_HEADER = "simpMessageType";
@@ -65,6 +63,11 @@ public class SimpMessageHeaderAccessor extends NativeMessageHeaderAccessor {
 
 	public static final String USER_HEADER = "simpUser";
 
+	public static final String CONNECT_MESSAGE_HEADER = "simpConnectMessage";
+
+	public static final String HEART_BEAT_HEADER = "simpHeartbeat";
+
+
 	/**
 	 * For internal use.
 	 * <p>The original destination used by a client when subscribing. Such a
@@ -262,4 +265,8 @@ public class SimpMessageHeaderAccessor extends NativeMessageHeaderAccessor {
 		return (Principal) headers.get(USER_HEADER);
 	}
 
+	public static long[] getHeartbeat(Map<String, Object> headers) {
+		return (long[]) headers.get(HEART_BEAT_HEADER);
+	}
+
 }

spring-messaging/src/main/java/org/springframework/messaging/simp/SimpMessageType.java

@@ -29,14 +29,14 @@ public enum SimpMessageType {
 
 	CONNECT_ACK,
 
-	HEARTBEAT,
-
 	MESSAGE,
 
 	SUBSCRIBE,
 
 	UNSUBSCRIBE,
 
+	HEARTBEAT,
+
 	DISCONNECT,
 
 	DISCONNECT_ACK,

spring-messaging/src/main/java/org/springframework/messaging/simp/broker/SimpleBrokerMessageHandler.java

@@ -16,7 +16,11 @@
 
 package org.springframework.messaging.simp.broker;
 
+import java.security.Principal;
 import java.util.Collection;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ScheduledFuture;
 
 import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageChannel;
@@ -27,6 +31,7 @@ import org.springframework.messaging.simp.SimpMessageType;
 import org.springframework.messaging.support.MessageBuilder;
 import org.springframework.messaging.support.MessageHeaderAccessor;
 import org.springframework.messaging.support.MessageHeaderInitializer;
+import org.springframework.scheduling.TaskScheduler;
 import org.springframework.util.Assert;
 import org.springframework.util.MultiValueMap;
 import org.springframework.util.PathMatcher;
@@ -43,10 +48,18 @@ public class SimpleBrokerMessageHandler extends AbstractBrokerMessageHandler {
 
 	private static final byte[] EMPTY_PAYLOAD = new byte[0];
 
+	private final Map<String, SessionInfo> sessions = new ConcurrentHashMap<String, SessionInfo>();
+
 	private SubscriptionRegistry subscriptionRegistry;
 
 	private PathMatcher pathMatcher;
 
+	private TaskScheduler taskScheduler;
+
+	private long[] heartbeatValue;
+
+	private ScheduledFuture<?> heartbeatFuture;
+
 	private MessageHeaderInitializer headerInitializer;
 
 
@@ -100,6 +113,49 @@ public class SimpleBrokerMessageHandler extends AbstractBrokerMessageHandler {
 		initPathMatcherToUse();
 	}
 
+	/**
+	 * Configure the {@link org.springframework.scheduling.TaskScheduler} to
+	 * use for providing heartbeat support. Setting this property also sets the
+	 * {@link #setHeartbeatValue heartbeatValue} to "10000, 10000".
+	 * <p>By default this is not set.
+	 * @since 4.2
+	 */
+	public void setTaskScheduler(TaskScheduler taskScheduler) {
+		Assert.notNull(taskScheduler);
+		this.taskScheduler = taskScheduler;
+		if (this.heartbeatValue == null) {
+			this.heartbeatValue = new long[] {10000, 10000};
+		}
+	}
+
+	/**
+	 * Return the configured TaskScheduler.
+	 */
+	public TaskScheduler getTaskScheduler() {
+		return this.taskScheduler;
+	}
+
+	/**
+	 * Configure the value for the heart-beat settings. The first number
+	 * represents how often the server will write or send a heartbeat.
+	 * The second is how often the client should write. 0 means no heartbeats.
+	 * <p>By default this is set to "0, 0" unless the {@link #setTaskScheduler
+	 * taskScheduler} in which case the default becomes "10000,10000"
+	 * (in milliseconds).
+	 * @since 4.2
+	 */
+	public void setHeartbeatValue(long[] heartbeat) {
+		Assert.notNull(heartbeat);
+		this.heartbeatValue = heartbeat;
+	}
+
+	/**
+	 * The configured value for the heart-beat settings.
+	 */
+	public long[] getHeartbeatValue() {
+		return this.heartbeatValue;
+	}
+
 	/**
 	 * Configure a {@link MessageHeaderInitializer} to apply to the headers
 	 * of all messages sent to the client outbound channel.
@@ -120,11 +176,37 @@ public class SimpleBrokerMessageHandler extends AbstractBrokerMessageHandler {
 	@Override
 	public void startInternal() {
 		publishBrokerAvailableEvent();
+		if (getTaskScheduler() != null) {
+			long interval = initHeartbeatTaskDelay();
+			if (interval > 0) {
+				this.heartbeatFuture = this.taskScheduler.scheduleWithFixedDelay(new HeartbeatTask(), interval);
+			}
+		}
+		else {
+			Assert.isTrue(getHeartbeatValue() == null ||
+					(getHeartbeatValue()[0] == 0 && getHeartbeatValue()[1] == 0),
+					"Heartbeat values configured but no TaskScheduler is provided.");
+		}
+	}
+
+	private long initHeartbeatTaskDelay() {
+		if (getHeartbeatValue() == null) {
+			return 0;
+		}
+		else if (getHeartbeatValue()[0] > 0 && getHeartbeatValue()[1] > 0) {
+			return Math.min(getHeartbeatValue()[0], getHeartbeatValue()[1]);
+		}
+		else {
+			return (getHeartbeatValue()[0] > 0 ? getHeartbeatValue()[0] : getHeartbeatValue()[1]);
+		}
 	}
 
 	@Override
 	public void stopInternal() {
 		publishBrokerUnavailableEvent();
+		if (this.heartbeatFuture != null) {
+			this.heartbeatFuture.cancel(true);
+		}
 	}
 
 	@Override
@@ -133,6 +215,9 @@ public class SimpleBrokerMessageHandler extends AbstractBrokerMessageHandler {
 		SimpMessageType messageType = SimpMessageHeaderAccessor.getMessageType(headers);
 		String destination = SimpMessageHeaderAccessor.getDestination(headers);
 		String sessionId = SimpMessageHeaderAccessor.getSessionId(headers);
+		Principal user = SimpMessageHeaderAccessor.getUser(headers);
+
+		updateSessionReadTime(sessionId);
 
 		if (!checkDestinationPrefix(destination)) {
 			return;
@@ -150,23 +235,21 @@ public class SimpleBrokerMessageHandler extends AbstractBrokerMessageHandler {
 		}
 		else if (SimpMessageType.CONNECT.equals(messageType)) {
 			logMessage(message);
+			long[] clientHeartbeat = SimpMessageHeaderAccessor.getHeartbeat(headers);
+			long[] serverHeartbeat = getHeartbeatValue();
+			this.sessions.put(sessionId, new SessionInfo(sessionId, user, clientHeartbeat, serverHeartbeat));
 			SimpMessageHeaderAccessor connectAck = SimpMessageHeaderAccessor.create(SimpMessageType.CONNECT_ACK);
 			initHeaders(connectAck);
 			connectAck.setSessionId(sessionId);
 			connectAck.setUser(SimpMessageHeaderAccessor.getUser(headers));
 			connectAck.setHeader(SimpMessageHeaderAccessor.CONNECT_MESSAGE_HEADER, message);
+			connectAck.setHeader(SimpMessageHeaderAccessor.HEART_BEAT_HEADER, serverHeartbeat);
 			Message<byte[]> messageOut = MessageBuilder.createMessage(EMPTY_PAYLOAD, connectAck.getMessageHeaders());
 			getClientOutboundChannel().send(messageOut);
 		}
 		else if (SimpMessageType.DISCONNECT.equals(messageType)) {
 			logMessage(message);
-			this.subscriptionRegistry.unregisterAllSubscriptions(sessionId);
-			SimpMessageHeaderAccessor disconnectAck = SimpMessageHeaderAccessor.create(SimpMessageType.DISCONNECT_ACK);
-			initHeaders(disconnectAck);
-			disconnectAck.setSessionId(sessionId);
-			disconnectAck.setUser(SimpMessageHeaderAccessor.getUser(headers));
-			Message<byte[]> messageOut = MessageBuilder.createMessage(EMPTY_PAYLOAD, disconnectAck.getMessageHeaders());
-			getClientOutboundChannel().send(messageOut);
+			handleDisconnect(sessionId, user);
 		}
 		else if (SimpMessageType.SUBSCRIBE.equals(messageType)) {
 			logMessage(message);
@@ -178,6 +261,15 @@ public class SimpleBrokerMessageHandler extends AbstractBrokerMessageHandler {
 		}
 	}
 
+	private void updateSessionReadTime(String sessionId) {
+		if (sessionId != null) {
+			SessionInfo info = this.sessions.get(sessionId);
+			if (info != null) {
+				info.setLastReadTime(System.currentTimeMillis());
+			}
+		}
+	}
+
 	private void logMessage(Message<?> message) {
 		if (logger.isDebugEnabled()) {
 			SimpMessageHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, SimpMessageHeaderAccessor.class);
@@ -192,11 +284,23 @@ public class SimpleBrokerMessageHandler extends AbstractBrokerMessageHandler {
 		}
 	}
 
+	private void handleDisconnect(String sessionId, Principal user) {
+		this.sessions.remove(sessionId);
+		this.subscriptionRegistry.unregisterAllSubscriptions(sessionId);
+		SimpMessageHeaderAccessor accessor = SimpMessageHeaderAccessor.create(SimpMessageType.DISCONNECT_ACK);
+		accessor.setSessionId(sessionId);
+		accessor.setUser(user);
+		initHeaders(accessor);
+		Message<byte[]> message = MessageBuilder.createMessage(EMPTY_PAYLOAD, accessor.getMessageHeaders());
+		getClientOutboundChannel().send(message);
+	}
+
 	protected void sendMessageToSubscribers(String destination, Message<?> message) {
 		MultiValueMap<String,String> subscriptions = this.subscriptionRegistry.findSubscriptions(message);
 		if (!subscriptions.isEmpty() && logger.isDebugEnabled()) {
 			logger.debug("Broadcasting to " + subscriptions.size() + " sessions.");
 		}
+		long now = System.currentTimeMillis();
 		for (String sessionId : subscriptions.keySet()) {
 			for (String subscriptionId : subscriptions.get(sessionId)) {
 				SimpMessageHeaderAccessor headerAccessor = SimpMessageHeaderAccessor.create(SimpMessageType.MESSAGE);
@@ -212,6 +316,12 @@ public class SimpleBrokerMessageHandler extends AbstractBrokerMessageHandler {
 				catch (Throwable ex) {
 					logger.error("Failed to send " + message, ex);
 				}
+				finally {
+					SessionInfo info = this.sessions.get(sessionId);
+					if (info != null) {
+						info.setLastWriteTime(now);
+					}
+				}
 			}
 		}
 	}
@@ -221,4 +331,93 @@ public class SimpleBrokerMessageHandler extends AbstractBrokerMessageHandler {
 		return "SimpleBroker[" + this.subscriptionRegistry + "]";
 	}
 
+
+	private static class SessionInfo {
+
+		/* STOMP spec: receiver SHOULD take into account an error margin */
+		private static final long HEARTBEAT_MULTIPLIER = 3;
+
+
+		private final String sessiondId;
+
+		private final Principal user;
+
+		private final long readInterval;
+
+		private final long writeInterval;
+
+		private volatile long lastReadTime;
+
+		private volatile long lastWriteTime;
+
+
+		public SessionInfo(String sessiondId, Principal user, long[] clientHeartbeat, long[] serverHeartbeat) {
+			this.sessiondId = sessiondId;
+			this.user = user;
+			if (clientHeartbeat != null && serverHeartbeat != null) {
+				this.readInterval = (clientHeartbeat[0] > 0 && serverHeartbeat[1] > 0 ?
+						Math.max(clientHeartbeat[0], serverHeartbeat[1]) * HEARTBEAT_MULTIPLIER : 0);
+				this.writeInterval = (clientHeartbeat[1] > 0 && serverHeartbeat[0] > 0 ?
+						Math.max(clientHeartbeat[1], serverHeartbeat[0]) : 0);
+			}
+			else {
+				this.readInterval = 0;
+				this.writeInterval = 0;
+			}
+			this.lastReadTime = this.lastWriteTime = System.currentTimeMillis();
+		}
+
+		public String getSessiondId() {
+			return this.sessiondId;
+		}
+
+		public Principal getUser() {
+			return this.user;
+		}
+
+		public long getReadInterval() {
+			return this.readInterval;
+		}
+
+		public long getWriteInterval() {
+			return this.writeInterval;
+		}
+
+		public long getLastReadTime() {
+			return this.lastReadTime;
+		}
+
+		public void setLastReadTime(long lastReadTime) {
+			this.lastReadTime = lastReadTime;
+		}
+
+		public long getLastWriteTime() {
+			return this.lastWriteTime;
+		}
+
+		public void setLastWriteTime(long lastWriteTime) {
+			this.lastWriteTime = lastWriteTime;
+		}
+	}
+
+	private class HeartbeatTask implements Runnable {
+
+		@Override
+		public void run() {
+			long now = System.currentTimeMillis();
+			for (SessionInfo info : sessions.values()) {
+				if (info.getReadInterval() > 0 && (now - info.getLastReadTime()) > info.getReadInterval()) {
+					handleDisconnect(info.getSessiondId(), info.getUser());
+				}
+				if (info.getWriteInterval() > 0 && (now - info.getLastWriteTime()) > info.getWriteInterval()) {
+					SimpMessageHeaderAccessor accessor = SimpMessageHeaderAccessor.create(SimpMessageType.HEARTBEAT);
+					accessor.setSessionId(info.getSessiondId());
+					accessor.setUser(info.getUser());
+					initHeaders(accessor);
+					MessageHeaders headers = accessor.getMessageHeaders();
+					getClientOutboundChannel().send(MessageBuilder.createMessage(EMPTY_PAYLOAD, headers));
+				}
+			}
+		}
+	}
 }

spring-messaging/src/main/java/org/springframework/messaging/simp/config/SimpleBrokerRegistration.java

@@ -19,6 +19,7 @@ package org.springframework.messaging.simp.config;
 import org.springframework.messaging.MessageChannel;
 import org.springframework.messaging.SubscribableChannel;
 import org.springframework.messaging.simp.broker.SimpleBrokerMessageHandler;
+import org.springframework.scheduling.TaskScheduler;
 
 /**
  * Registration class for configuring a {@link SimpleBrokerMessageHandler}.
@@ -28,14 +29,54 @@ import org.springframework.messaging.simp.broker.SimpleBrokerMessageHandler;
  */
 public class SimpleBrokerRegistration extends AbstractBrokerRegistration {
 
+	private TaskScheduler taskScheduler;
+
+	private long[] heartbeat;
+
+
 	public SimpleBrokerRegistration(SubscribableChannel inChannel, MessageChannel outChannel, String[] prefixes) {
 		super(inChannel, outChannel, prefixes);
 	}
 
+
+	/**
+	 * Configure the {@link org.springframework.scheduling.TaskScheduler} to
+	 * use for providing heartbeat support. Setting this property also sets the
+	 * {@link #setHeartbeatValue heartbeatValue} to "10000, 10000".
+	 * <p>By default this is not set.
+	 * @since 4.2
+	 */
+	public SimpleBrokerRegistration setTaskScheduler(TaskScheduler taskScheduler) {
+		this.taskScheduler = taskScheduler;
+		return this;
+	}
+
+	/**
+	 * Configure the value for the heartbeat settings. The first number
+	 * represents how often the server will write or send a heartbeat.
+	 * The second is how often the client should write. 0 means no heartbeats.
+	 * <p>By default this is set to "0, 0" unless the {@link #setTaskScheduler
+	 * taskScheduler} in which case the default becomes "10000,10000"
+	 * (in milliseconds).
+	 * @since 4.2
+	 */
+	public SimpleBrokerRegistration setHeartbeatValue(long[] heartbeat) {
+		this.heartbeat = heartbeat;
+		return this;
+	}
+
+
 	@Override
 	protected SimpleBrokerMessageHandler getMessageHandler(SubscribableChannel brokerChannel) {
-		return new SimpleBrokerMessageHandler(getClientInboundChannel(),
+		SimpleBrokerMessageHandler handler = new SimpleBrokerMessageHandler(getClientInboundChannel(),
 				getClientOutboundChannel(), brokerChannel, getDestinationPrefixes());
+		if (this.taskScheduler != null) {
+			handler.setTaskScheduler(this.taskScheduler);
+		}
+		if (this.heartbeat != null) {
+			handler.setHeartbeatValue(this.heartbeat);
+		}
+		return handler;
 	}
 
 }

spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/DefaultStompSession.java

@@ -389,7 +389,7 @@ public class DefaultStompSession implements ConnectionHandlingStompSession {
 					}
 				}
 				else if (StompCommand.CONNECTED.equals(command)) {
-					initHeartbeats(stompHeaders);
+					initHeartbeatTasks(stompHeaders);
 					this.sessionFuture.set(this);
 					this.sessionHandler.afterConnected(this, stompHeaders);
 				}
@@ -420,20 +420,18 @@ public class DefaultStompSession implements ConnectionHandlingStompSession {
 		handler.handleFrame(stompHeaders, object);
 	}
 
-	private void initHeartbeats(StompHeaders connectedHeaders) {
-		long clientRead = this.connectHeaders.getHeartbeat()[0];
-		long serverWrite = connectedHeaders.getHeartbeat()[1];
-
-		if (clientRead > 0 && serverWrite > 0) {
-			long interval = Math.max(clientRead,  serverWrite);
+	private void initHeartbeatTasks(StompHeaders connectedHeaders) {
+		long[] connect = this.connectHeaders.getHeartbeat();
+		long[] connected = connectedHeaders.getHeartbeat();
+		if (connect == null || connected == null) {
+			return;
+		}
+		if (connect[0] > 0 && connected[1] > 0) {
+			long interval = Math.max(connect[0],  connected[1]);
 			this.connection.onWriteInactivity(new WriteInactivityTask(), interval);
 		}
-
-		long clientWrite = this.connectHeaders.getHeartbeat()[1];
-		long serverRead = connectedHeaders.getHeartbeat()[0];
-
-		if (clientWrite > 0 && serverRead > 0) {
-			final long interval = Math.max(clientWrite, serverRead) * HEARTBEAT_MULTIPLIER;
+		if (connect[1] > 0 && connected[0] > 0) {
+			final long interval = Math.max(connect[1], connected[0]) * HEARTBEAT_MULTIPLIER;
 			this.connection.onReadInactivity(new ReadInactivityTask(), interval);
 		}
 	}

spring-messaging/src/test/java/org/springframework/messaging/simp/broker/SimpleBrokerMessageHandlerTests.java

@@ -16,12 +16,15 @@
 
 package org.springframework.messaging.simp.broker;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertTrue;
-import static org.mockito.Mockito.times;
-import static org.mockito.Mockito.verify;
+import static org.junit.Assert.*;
+import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.*;
 
+import java.security.Principal;
 import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.ScheduledFuture;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -29,13 +32,16 @@ import org.mockito.ArgumentCaptor;
 import org.mockito.Captor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+
 import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageChannel;
+import org.springframework.messaging.MessageHeaders;
 import org.springframework.messaging.SubscribableChannel;
 import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
 import org.springframework.messaging.simp.SimpMessageType;
 import org.springframework.messaging.simp.TestPrincipal;
 import org.springframework.messaging.support.MessageBuilder;
+import org.springframework.scheduling.TaskScheduler;
 
 /**
  * Unit tests for SimpleBrokerMessageHandler.
@@ -43,6 +49,7 @@ import org.springframework.messaging.support.MessageBuilder;
  * @author Rossen Stoyanchev
  * @since 4.0
  */
+@SuppressWarnings("unchecked")
 public class SimpleBrokerMessageHandlerTests {
 
 	private SimpleBrokerMessageHandler messageHandler;
@@ -56,6 +63,9 @@ public class SimpleBrokerMessageHandlerTests {
 	@Mock
 	private SubscribableChannel brokerChannel;
 
+	@Mock
+	private TaskScheduler taskScheduler;
+
 	@Captor
 	ArgumentCaptor<Message<?>> messageCaptor;
 
@@ -133,11 +143,11 @@ public class SimpleBrokerMessageHandlerTests {
 	@Test
 	public void connect() {
 
-		String sess1 = "sess1";
-
 		this.messageHandler.start();
 
-		Message<String> connectMessage = createConnectMessage(sess1);
+		String id = "sess1";
+		Message<String> connectMessage = createConnectMessage(id, new TestPrincipal("joe"), null);
+		this.messageHandler.setTaskScheduler(this.taskScheduler);
 		this.messageHandler.handleMessage(connectMessage);
 
 		verify(this.clientOutboundChannel, times(1)).send(this.messageCaptor.capture());
@@ -145,10 +155,150 @@ public class SimpleBrokerMessageHandlerTests {
 
 		SimpMessageHeaderAccessor connectAckHeaders = SimpMessageHeaderAccessor.wrap(connectAckMessage);
 		assertEquals(connectMessage, connectAckHeaders.getHeader(SimpMessageHeaderAccessor.CONNECT_MESSAGE_HEADER));
-		assertEquals(sess1, connectAckHeaders.getSessionId());
+		assertEquals(id, connectAckHeaders.getSessionId());
 		assertEquals("joe", connectAckHeaders.getUser().getName());
+		assertArrayEquals(new long[] {10000, 10000},
+				SimpMessageHeaderAccessor.getHeartbeat(connectAckHeaders.getMessageHeaders()));
 	}
 
+	@Test
+	public void heartbeatValueWithAndWithoutTaskScheduler() throws Exception {
+
+		assertNull(this.messageHandler.getHeartbeatValue());
+
+		this.messageHandler.setTaskScheduler(this.taskScheduler);
+
+		assertNotNull(this.messageHandler.getHeartbeatValue());
+		assertArrayEquals(new long[] {10000, 10000}, this.messageHandler.getHeartbeatValue());
+	}
+
+	@Test(expected = IllegalArgumentException.class)
+	public void startWithHeartbeatValueWithoutTaskScheduler() throws Exception {
+		this.messageHandler.setHeartbeatValue(new long[] {10000, 10000});
+		this.messageHandler.start();
+	}
+
+	@SuppressWarnings("unchecked")
+	@Test
+	public void startAndStopWithHeartbeatValue() throws Exception {
+
+		ScheduledFuture future = mock(ScheduledFuture.class);
+		when(this.taskScheduler.scheduleWithFixedDelay(any(Runnable.class), eq(15000L))).thenReturn(future);
+
+		this.messageHandler.setTaskScheduler(this.taskScheduler);
+		this.messageHandler.setHeartbeatValue(new long[] {15000, 16000});
+		this.messageHandler.start();
+
+		verify(this.taskScheduler).scheduleWithFixedDelay(any(Runnable.class), eq(15000L));
+		verifyNoMoreInteractions(this.taskScheduler, future);
+
+		this.messageHandler.stop();
+
+		verify(future).cancel(true);
+		verifyNoMoreInteractions(future);
+	}
+
+	@SuppressWarnings("unchecked")
+	@Test
+	public void startWithOneZeroHeartbeatValue() throws Exception {
+
+		this.messageHandler.setTaskScheduler(this.taskScheduler);
+		this.messageHandler.setHeartbeatValue(new long[] {0, 10000});
+		this.messageHandler.start();
+
+		verify(this.taskScheduler).scheduleWithFixedDelay(any(Runnable.class), eq(10000L));
+	}
+
+	@Test
+	public void readInactivity() throws Exception {
+
+		this.messageHandler.setHeartbeatValue(new long[] {0, 1});
+		this.messageHandler.setTaskScheduler(this.taskScheduler);
+		this.messageHandler.start();
+
+		ArgumentCaptor<Runnable> taskCaptor = ArgumentCaptor.forClass(Runnable.class);
+		verify(this.taskScheduler).scheduleWithFixedDelay(taskCaptor.capture(), eq(1L));
+		Runnable heartbeatTask = taskCaptor.getValue();
+		assertNotNull(heartbeatTask);
+
+		String id = "sess1";
+		TestPrincipal user = new TestPrincipal("joe");
+		Message<String> connectMessage = createConnectMessage(id, user, new long[] {1, 0});
+		this.messageHandler.handleMessage(connectMessage);
+
+		Thread.sleep(10);
+		heartbeatTask.run();
+
+		verify(this.clientOutboundChannel, atLeast(2)).send(this.messageCaptor.capture());
+		List<Message<?>> messages = this.messageCaptor.getAllValues();
+		assertEquals(2, messages.size());
+
+		MessageHeaders headers = messages.get(0).getHeaders();
+		assertEquals(SimpMessageType.CONNECT_ACK, headers.get(SimpMessageHeaderAccessor.MESSAGE_TYPE_HEADER));
+		headers = messages.get(1).getHeaders();
+		assertEquals(SimpMessageType.DISCONNECT_ACK, headers.get(SimpMessageHeaderAccessor.MESSAGE_TYPE_HEADER));
+		assertEquals(id, headers.get(SimpMessageHeaderAccessor.SESSION_ID_HEADER));
+		assertEquals(user, headers.get(SimpMessageHeaderAccessor.USER_HEADER));
+	}
+
+	@Test
+	public void writeInactivity() throws Exception {
+
+		this.messageHandler.setHeartbeatValue(new long[] {1, 0});
+		this.messageHandler.setTaskScheduler(this.taskScheduler);
+		this.messageHandler.start();
+
+		ArgumentCaptor<Runnable> taskCaptor = ArgumentCaptor.forClass(Runnable.class);
+		verify(this.taskScheduler).scheduleWithFixedDelay(taskCaptor.capture(), eq(1L));
+		Runnable heartbeatTask = taskCaptor.getValue();
+		assertNotNull(heartbeatTask);
+
+		String id = "sess1";
+		TestPrincipal user = new TestPrincipal("joe");
+		Message<String> connectMessage = createConnectMessage(id, user, new long[] {0, 1});
+		this.messageHandler.handleMessage(connectMessage);
+
+		Thread.sleep(10);
+		heartbeatTask.run();
+
+		verify(this.clientOutboundChannel, times(2)).send(this.messageCaptor.capture());
+		List<Message<?>> messages = this.messageCaptor.getAllValues();
+		assertEquals(2, messages.size());
+
+		MessageHeaders headers = messages.get(0).getHeaders();
+		assertEquals(SimpMessageType.CONNECT_ACK, headers.get(SimpMessageHeaderAccessor.MESSAGE_TYPE_HEADER));
+		headers = messages.get(1).getHeaders();
+		assertEquals(SimpMessageType.HEARTBEAT, headers.get(SimpMessageHeaderAccessor.MESSAGE_TYPE_HEADER));
+		assertEquals(id, headers.get(SimpMessageHeaderAccessor.SESSION_ID_HEADER));
+		assertEquals(user, headers.get(SimpMessageHeaderAccessor.USER_HEADER));
+	}
+
+	@Test
+	public void readWriteIntervalCalculation() throws Exception {
+
+		this.messageHandler.setHeartbeatValue(new long[] {1, 1});
+		this.messageHandler.setTaskScheduler(this.taskScheduler);
+		this.messageHandler.start();
+
+		ArgumentCaptor<Runnable> taskCaptor = ArgumentCaptor.forClass(Runnable.class);
+		verify(this.taskScheduler).scheduleWithFixedDelay(taskCaptor.capture(), eq(1L));
+		Runnable heartbeatTask = taskCaptor.getValue();
+		assertNotNull(heartbeatTask);
+
+		String id = "sess1";
+		TestPrincipal user = new TestPrincipal("joe");
+		Message<String> connectMessage = createConnectMessage(id, user, new long[] {10000, 10000});
+		this.messageHandler.handleMessage(connectMessage);
+
+		Thread.sleep(10);
+		heartbeatTask.run();
+
+		verify(this.clientOutboundChannel, times(1)).send(this.messageCaptor.capture());
+		List<Message<?>> messages = this.messageCaptor.getAllValues();
+		assertEquals(1, messages.size());
+		assertEquals(SimpMessageType.CONNECT_ACK,
+				messages.get(0).getHeaders().get(SimpMessageHeaderAccessor.MESSAGE_TYPE_HEADER));
+	}
 
 	private Message<String> createSubscriptionMessage(String sessionId, String subcriptionId, String destination) {
 		SimpMessageHeaderAccessor headers = SimpMessageHeaderAccessor.create(SimpMessageType.SUBSCRIBE);
@@ -158,17 +308,18 @@ public class SimpleBrokerMessageHandlerTests {
 		return MessageBuilder.createMessage("", headers.getMessageHeaders());
 	}
 
-	private Message<String> createConnectMessage(String sessionId) {
-		SimpMessageHeaderAccessor headers = SimpMessageHeaderAccessor.create(SimpMessageType.CONNECT);
-		headers.setSessionId(sessionId);
-		headers.setUser(new TestPrincipal("joe"));
-		return MessageBuilder.createMessage("", headers.getMessageHeaders());
+	private Message<String> createConnectMessage(String sessionId, Principal user, long[] heartbeat) {
+		SimpMessageHeaderAccessor accessor = SimpMessageHeaderAccessor.create(SimpMessageType.CONNECT);
+		accessor.setSessionId(sessionId);
+		accessor.setUser(user);
+		accessor.setHeader(SimpMessageHeaderAccessor.HEART_BEAT_HEADER, heartbeat);
+		return MessageBuilder.createMessage("", accessor.getMessageHeaders());
 	}
 
 	private Message<String> createMessage(String destination, String payload) {
 		SimpMessageHeaderAccessor headers = SimpMessageHeaderAccessor.create(SimpMessageType.MESSAGE);
 		headers.setDestination(destination);
-		return MessageBuilder.createMessage("", headers.getMessageHeaders());
+		return MessageBuilder.createMessage(payload, headers.getMessageHeaders());
 	}
 
 	private boolean messageCaptured(String sessionId, String subcriptionId, String destination) {