From 607d591e3c7e17eb5497a408a1ec3224cfd492b3 Mon Sep 17 00:00:00 2001
From: fhanik <fhanik@13f79535-47bb-0310-9956-ffa450edef68>
Date: Fri, 24 Oct 2008 16:14:07 +0000
Subject: [PATCH] During a read or write timeout, we were cancelling the main
 key, this is fine for regular IO, but for Comet IO, this will cause events to
 stop since we remove the socket from the main selector

git-svn-id: https://svn.apache.org/repos/asf/tomcat/trunk@707670 13f79535-47bb-0310-9956-ffa450edef68
---
 .../tomcat/util/net/NioBlockingSelector.java       | 40 ++++++++++++++++------
 1 file changed, 30 insertions(+), 10 deletions(-)

diff --git a/java/org/apache/tomcat/util/net/NioBlockingSelector.java b/java/org/apache/tomcat/util/net/NioBlockingSelector.java
index 292bbdf45..85dcd7027 100644
--- a/java/org/apache/tomcat/util/net/NioBlockingSelector.java
+++ b/java/org/apache/tomcat/util/net/NioBlockingSelector.java
@@ -81,6 +81,7 @@ public class NioBlockingSelector {
     public int write(ByteBuffer buf, NioChannel socket, long writeTimeout,MutableInteger lastWrite) throws IOException {
         SelectionKey key = socket.getIOChannel().keyFor(socket.getPoller().getSelector());
         if ( key == null ) throw new IOException("Key no longer registered");
+        KeyReference reference = new KeyReference();
         KeyAttachment att = (KeyAttachment) key.attachment();
         int written = 0;
         boolean timedout = false;
@@ -101,7 +102,7 @@ public class NioBlockingSelector {
                 }
                 try {
                     if ( att.getWriteLatch()==null || att.getWriteLatch().getCount()==0) att.startWriteLatch(1);
-                    poller.add(att,SelectionKey.OP_WRITE);
+                    poller.add(att,SelectionKey.OP_WRITE,reference);
                     att.awaitWriteLatch(writeTimeout,TimeUnit.MILLISECONDS);
                 }catch (InterruptedException ignore) {
                     Thread.interrupted();
@@ -122,9 +123,10 @@ public class NioBlockingSelector {
                 throw new SocketTimeoutException();
         } finally {
             poller.remove(att,SelectionKey.OP_WRITE);
-            if (timedout && key != null) {
-                poller.cancelKey(socket, key);
+            if (timedout && reference.key!=null) {
+                poller.cancelKey(reference.key);
             }
+            reference.key = null;
         }
         return written;
     }
@@ -145,6 +147,7 @@ public class NioBlockingSelector {
     public int read(ByteBuffer buf, NioChannel socket, long readTimeout) throws IOException {
         SelectionKey key = socket.getIOChannel().keyFor(socket.getPoller().getSelector());
         if ( key == null ) throw new IOException("Key no longer registered");
+        KeyReference reference = new KeyReference();
         KeyAttachment att = (KeyAttachment) key.attachment();
         int read = 0;
         boolean timedout = false;
@@ -162,7 +165,7 @@ public class NioBlockingSelector {
                 }
                 try {
                     if ( att.getReadLatch()==null || att.getReadLatch().getCount()==0) att.startReadLatch(1);
-                    poller.add(att,SelectionKey.OP_READ);
+                    poller.add(att,SelectionKey.OP_READ, reference);
                     att.awaitReadLatch(readTimeout,TimeUnit.MILLISECONDS);
                 }catch (InterruptedException ignore) {
                     Thread.interrupted();
@@ -182,9 +185,10 @@ public class NioBlockingSelector {
                 throw new SocketTimeoutException();
         } finally {
             poller.remove(att,SelectionKey.OP_READ);
-            if (timedout && key != null) {
-                poller.cancelKey(socket,key);
+            if (timedout && reference.key!=null) {
+                poller.cancelKey(reference.key);
             }
+            reference.key = null;
         }
         return read;
     }
@@ -196,7 +200,7 @@ public class NioBlockingSelector {
         protected ConcurrentLinkedQueue<Runnable> events = new ConcurrentLinkedQueue<Runnable>();
         public void disable() { run = false; selector.wakeup();}
         protected AtomicInteger wakeupCounter = new AtomicInteger(0);
-        public void cancelKey(final NioChannel socket, final SelectionKey key) {
+        public void cancelKey(final SelectionKey key) {
             Runnable r = new Runnable() {
                 public void run() {
                     key.cancel();
@@ -219,7 +223,7 @@ public class NioBlockingSelector {
             }
         }
         
-        public void add(final KeyAttachment key, final int ops) {
+        public void add(final KeyAttachment key, final int ops, final KeyReference ref) {
             Runnable r = new Runnable() {
                 public void run() {
                     if ( key == null ) return;
@@ -231,6 +235,7 @@ public class NioBlockingSelector {
                     try {
                         if (sk == null) {
                             sk = ch.register(selector, ops, key);
+                            ref.key = sk;
                         } else if (!sk.isValid()) {
                             cancel(sk,key,ops);
                         } else {
@@ -360,15 +365,30 @@ public class NioBlockingSelector {
             }catch( Exception ignore ) {
                 if (log.isDebugEnabled())log.debug("",ignore);
             }
+            try {
+                selector.close();//Close the connector
+            }catch( Exception ignore ) {
+                if (log.isDebugEnabled())log.debug("",ignore);
+            }
         }
         
         public void countDown(CountDownLatch latch) {
             if ( latch == null ) return;
             latch.countDown();
         }
+    }
+    
+    public class KeyReference {
+        SelectionKey key = null;
         
-        
-        
+        @Override
+        public void finalize() {
+            if (key!=null && key.isValid()) {
+                log.warn("Possible key leak, cancelling key in the finalizer.");
+                try {key.cancel();}catch (Exception ignore){}
+            }
+            key = null;
+        }
     }
 
 }
\ No newline at end of file
-- 
2.11.0