Properly close files and stream on exit

3 files changed, 32 insertions, 7 deletions

diff --git a/src/io/DataWriter.java b/src/io/DataWriter.java
index d06e283..ebbd129 100644
--- a/src/io/DataWriter.java
+++ b/src/io/DataWriter.java
@@ -1,5 +1,6 @@
 package io;
 
+import java.io.Closeable;
 import java.io.FileInputStream;
 import java.io.FileNotFoundException;
 import java.io.FileOutputStream;
@@ -22,7 +23,7 @@ import provider.ResultListener;
  *
  * @author Maurice Laveaux
  */
-public class DataWriter implements ResultListener {
+public class DataWriter implements ResultListener, Closeable {
 
     /**
      * The writer for the tweet stream.
diff --git a/src/main/TweetShell.java b/src/main/TweetShell.java
index be4124b..818c693 100644
--- a/src/main/TweetShell.java
+++ b/src/main/TweetShell.java
@@ -4,6 +4,7 @@ import io.CompressableDataWriter;
 import io.DataWriter;
 import io.OAuthRequester;
 import io.StreamImpl;
+import java.io.Closeable;
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.Map;
@@ -15,6 +16,7 @@ import java.util.logging.Level;
 import java.util.logging.Logger;
 import mining.Stream;
 import mining.TwitterApi;
+import org.apache.commons.io.IOUtils;
 import org.json.JSONException;
 import org.json.JSONObject;
 import provider.CompositeResultListener;
@@ -278,7 +280,7 @@ public class TweetShell implements TwitterApi.PinSupplier {
                 System.out.println("Unique users: " + tc.getUsers().size());
                 break;
             case close:
-                getStream().close();
+                safeClose();
                 break;
             case help:
                 for (String line : HELP) {
@@ -298,6 +300,7 @@ public class TweetShell implements TwitterApi.PinSupplier {
                 }
                 break;
             case exit:
+                safeClose();
                 throw new NoSuchElementException();
             case target:
                 if (params.length > 0) {
@@ -326,6 +329,18 @@ public class TweetShell implements TwitterApi.PinSupplier {
         }
     }
 
+    private void safeClose() {
+        if (stream_cached != null) {
+            try {
+                getStream().close();
+            } catch (IOException ex) {
+                // should not happen because getStream exists
+                System.err.println("getStream().close() -- " + ex.getMessage());
+            }
+        }
+        resultListeners.close();
+    }
+
     /**
      * @return All targets that can be disabled or enabled.
      */
@@ -430,8 +445,8 @@ public class TweetShell implements TwitterApi.PinSupplier {
 
         resultListeners.unregister(oldListener);
         // do we need to cleanup something?
-        if (oldListener instanceof DataWriter) {
-            ((DataWriter) oldListener).close();
+        if (oldListener instanceof Closeable) {
+            IOUtils.closeQuietly((Closeable) oldListener);
         }
         return true;
     }
diff --git a/src/provider/CompositeResultListener.java b/src/provider/CompositeResultListener.java
index 3db91da..bc8ad74 100644
--- a/src/provider/CompositeResultListener.java
+++ b/src/provider/CompositeResultListener.java
@@ -1,9 +1,9 @@
 package provider;
 
+import java.io.Closeable;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Set;
-import java.util.TreeSet;
+import org.apache.commons.io.IOUtils;
 import org.json.JSONObject;
 
 /**
@@ -11,7 +11,7 @@ import org.json.JSONObject;
  *
  * @author Peter Wu
  */
-public class CompositeResultListener implements ResultListener {
+public class CompositeResultListener implements ResultListener, Closeable {
 
     private final List<ResultListener> listeners;
 
@@ -53,4 +53,13 @@ public class CompositeResultListener implements ResultListener {
             rl.profileGenerated(obj);
         }
     }
+
+    @Override
+    public void close() {
+        for (ResultListener rl : listeners) {
+            if (rl instanceof Closeable) {
+                IOUtils.closeQuietly((Closeable) rl);
+            }
+        }
+    }
 }