Fixing a Thread Leak in Database Client Code (Related to #29)

src/main/java/com/google/firebase/database/tubesock/WebSocket.java

@@ -16,6 +16,8 @@
 
 package com.google.firebase.database.tubesock;
 
+import static com.google.common.base.Preconditions.checkState;
+
 import java.io.DataInputStream;
 import java.io.IOException;
 import java.io.OutputStream;
@@ -64,7 +66,7 @@ public void setName(Thread t, String name) {
   private final WebSocketWriter writer;
   private final WebSocketHandshake handshake;
   private final int clientId = clientCount.incrementAndGet();
-  private final Thread innerThread;
+  private Thread innerThread;
   private volatile State state = State.NONE;
   private volatile Socket socket = null;
   private WebSocketEventHandler eventHandler = null;
@@ -99,15 +101,6 @@ public WebSocket(URI url, String protocol) {
    *     if not extra headers are requested
    */
   public WebSocket(URI url, String protocol, Map<String, String> extraHeaders) {
-    innerThread =
-        getThreadFactory()
-            .newThread(
-                new Runnable() {
-                  @Override
-                  public void run() {
-                    runReader();
-                  }
-                });
     this.url = url;
     handshake = new WebSocketHandshake(url, protocol, extraHeaders);
     receiver = new WebSocketReceiver(this);
@@ -150,9 +143,23 @@ public synchronized void connect() {
       close();
       return;
     }
-    getIntializer().setName(getInnerThread(), THREAD_BASE_NAME + "Reader-" + clientId);
     state = State.CONNECTING;
-    getInnerThread().start();
+    start();
+  }
+
+  private synchronized void start() {
+    checkState(innerThread == null, "Inner thread already started");
+    innerThread =
+        getThreadFactory()
+            .newThread(
+                new Runnable() {
+                  @Override
+                  public void run() {
+                    runReader();
+                  }
+                });
+    getIntializer().setName(innerThread, THREAD_BASE_NAME + "Reader-" + clientId);
+    innerThread.start();
   }
 
   /**
@@ -312,13 +319,15 @@ private Socket createSocket() {
    * convenience method to make sure everything shuts down, if desired.
    */
   public void blockClose() throws InterruptedException {
-    // If the thread is new, it will never run, since we closed the connection before we
-    // actually
-    // connected
-    if (writer.getInnerThread().getState() != Thread.State.NEW) {
-      writer.getInnerThread().join();
+    writer.waitForTermination();
+    Thread thread;
+    synchronized (this) {
+      if (innerThread == null) {
+        return;
+      }
+      thread = innerThread;
     }
-    getInnerThread().join();
+    thread.join();
   }
 
   private void runReader() {
@@ -389,7 +398,7 @@ private void runReader() {
       writer.setOutput(output);
       receiver.setInput(input);
       state = WebSocket.State.CONNECTED;
-      writer.getInnerThread().start();
+      writer.start();
       eventHandler.onOpen();
       receiver.run();
     } catch (WebSocketException wse) {
@@ -402,10 +411,6 @@ private void runReader() {
     }
   }
 
-  Thread getInnerThread() {
-    return innerThread;
-  }
-
   private enum State {
     NONE,
     CONNECTING,

src/main/java/com/google/firebase/database/tubesock/WebSocketReceiver.java

@@ -47,6 +47,9 @@ void run() {
     this.eventHandler = websocket.getEventHandler();
     while (!stop) {
       try {
+        if (Thread.interrupted()) {
+          throw new InterruptedException();
+        }
         int offset = 0;
         offset += read(inputHeader, offset, 1);
         boolean fin = (inputHeader[0] & 0x80) != 0;
@@ -94,6 +97,8 @@ void run() {
         continue;
       } catch (IOException ioe) {
         handleError(new WebSocketException("IO Error", ioe));
+      } catch (InterruptedException e) {
+        handleError(new WebSocketException("Receiver interrupted", e));
       } catch (WebSocketException e) {
         handleError(e);
       }

src/main/java/com/google/firebase/database/tubesock/WebSocketWriter.java

@@ -16,6 +16,8 @@
 
 package com.google.firebase.database.tubesock;
 
+import static com.google.common.base.Preconditions.checkState;
+
 import java.io.IOException;
 import java.io.OutputStream;
 import java.nio.ByteBuffer;
@@ -32,26 +34,18 @@
 class WebSocketWriter {
 
   private final Random random = new Random();
-  private final Thread innerThread;
+  private final String threadName;
+
+  private Thread innerThread;
   private BlockingQueue<ByteBuffer> pendingBuffers;
   private volatile boolean stop = false;
   private boolean closeSent = false;
   private WebSocket websocket;
   private WritableByteChannel channel;
 
   WebSocketWriter(WebSocket websocket, String threadBaseName, int clientId) {
-    innerThread =
-        WebSocket.getThreadFactory()
-            .newThread(
-                new Runnable() {
-                  @Override
-                  public void run() {
-                    runWriter();
-                  }
-                });
-
-    WebSocket.getIntializer().setName(getInnerThread(), threadBaseName + "Writer-" + clientId);
     this.websocket = websocket;
+    this.threadName = threadBaseName + "Writer-" + clientId;
     pendingBuffers = new LinkedBlockingQueue<>();
   }
 
@@ -166,7 +160,31 @@ private void runWriter() {
     }
   }
 
-  Thread getInnerThread() {
-    return innerThread;
+  synchronized void start() {
+    checkState(innerThread == null, "Inner thread already started");
+    innerThread =
+        WebSocket.getThreadFactory()
+            .newThread(
+                new Runnable() {
+                  @Override
+                  public void run() {
+                    runWriter();
+                  }
+                });
+    WebSocket.getIntializer().setName(innerThread, threadName);
+    innerThread.start();
+  }
+
+  void waitForTermination() throws InterruptedException {
+    // If the thread is null, it will never be instantiated, since we closed the connection
+    // before we actually connected.
+    Thread thread;
+    synchronized (this) {
+      if (innerThread == null) {
+        return;
+      }
+      thread = innerThread;
+    }
+    thread.join();
   }
 }