From 2a2b9f026ea0a180d2edca804d868bd9cbfcdb0a Mon Sep 17 00:00:00 2001
From: Nikita <abracham.mitchell@gmail.com>
Date: Sun, 9 Sep 2012 12:49:27 +0400
Subject: [PATCH] Heartbeat race condition fixed

---
 .../corundumstudio/socketio/HeartbeatHandler.java   | 13 ++++++++-----
 .../socketio/scheduler/CancelableScheduler.java     |  2 +-
 2 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/src/main/java/com/corundumstudio/socketio/HeartbeatHandler.java b/src/main/java/com/corundumstudio/socketio/HeartbeatHandler.java
index b54dfc5..9ecda18 100644
--- a/src/main/java/com/corundumstudio/socketio/HeartbeatHandler.java
+++ b/src/main/java/com/corundumstudio/socketio/HeartbeatHandler.java
@@ -44,21 +44,24 @@ public class HeartbeatHandler implements Disconnectable {
             return;
         }
 
-        scheduler.cancel(new SchedulerKey(Type.HEARBEAT_TIMEOUT, client.getSessionId()));
+        final SchedulerKey key = new SchedulerKey(Type.HEARBEAT_TIMEOUT, client.getSessionId());
+        // cancel heartbeat check because the client answered
+        scheduler.cancel(key);
         scheduler.schedule(new Runnable() {
             public void run() {
                 client.send(new Packet(PacketType.HEARTBEAT));
-                scheduleClientHeartbeatCheck(client);
+                scheduleClientHeartbeatCheck(client, key);
             }
         }, configuration.getHeartbeatInterval(), TimeUnit.SECONDS);
     }
 
-    private void scheduleClientHeartbeatCheck(final BaseClient client) {
-        SchedulerKey key = new SchedulerKey(Type.HEARBEAT_TIMEOUT, client.getSessionId());
+    private void scheduleClientHeartbeatCheck(final BaseClient client, SchedulerKey key) {
+        // cancel previous heartbeat check
+        scheduler.cancel(key);
         scheduler.schedule(key, new Runnable() {
             public void run() {
                 client.disconnect();
-                log.debug("Client with sessionId: {} disconnected due to heartbeat timeout", client.getSessionId());
+                log.debug("Client with sessionId: {} disconnected due to heartbeat timeout, for {}", client.getSessionId());
             }
         }, configuration.getHeartbeatTimeout(), TimeUnit.SECONDS);
     }
diff --git a/src/main/java/com/corundumstudio/socketio/scheduler/CancelableScheduler.java b/src/main/java/com/corundumstudio/socketio/scheduler/CancelableScheduler.java
index ede3d42..71c68c8 100644
--- a/src/main/java/com/corundumstudio/socketio/scheduler/CancelableScheduler.java
+++ b/src/main/java/com/corundumstudio/socketio/scheduler/CancelableScheduler.java
@@ -24,7 +24,7 @@ import java.util.concurrent.TimeUnit;
 
 public class CancelableScheduler {
 
-    private final Map<Object, Future<?>> scheduledFutures = new ConcurrentHashMap<Object, Future<?>>();
+    private final Map<SchedulerKey, Future<?>> scheduledFutures = new ConcurrentHashMap<SchedulerKey, Future<?>>();
     private final ScheduledExecutorService executorService;
 
     public CancelableScheduler(int threadPoolSize) {