Wait for IO completion

build.gradle

@@ -4,9 +4,9 @@ version '0.1.4'
 apply plugin: 'java'
 apply plugin: 'application'
 
+
 sourceCompatibility = 1.8
 mainClassName = "me.davidsargent.stubjars.Main"
-
 repositories {
     mavenCentral()
 }

src/main/java/me/davidsargent/stubjars/StubJars.java

@@ -91,7 +91,15 @@ public void createSourceFiles() {
             sourceFiles.append(file.getAbsolutePath()).append(System.lineSeparator());
             writer.write();
         }
+
         writerThread.done();
+        try {
+            writerThread.waitForCompletion();
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+            return;
+        }
+
         Writer sourcesList = new Writer(SOURCES_LIST_FILE);
         try {
             sourcesList.write(sourceFiles.toString());

src/main/java/me/davidsargent/stubjars/components/writer/WriterThread.java

@@ -14,12 +14,14 @@
 package me.davidsargent.stubjars.components.writer;
 
 import java.util.concurrent.ArrayBlockingQueue;
+import java.util.concurrent.Semaphore;
 
 /**
  * @see Writer
  */
 public class WriterThread extends Thread implements Runnable {
     private final ArrayBlockingQueue<Writer> writersToProcess;
+    private Thread runningThread = null;
     private volatile boolean stop = false;
 
     public WriterThread() {
@@ -31,6 +33,10 @@ public void done() {
         stop = true;
     }
 
+    public void kill() {
+        runningThread.interrupt();
+    }
+
     void addWriter(Writer writer) {
         try {
             writersToProcess.put(writer);
@@ -39,21 +45,28 @@ void addWriter(Writer writer) {
         }
     }
 
-    @Override
-    public void run() {
+    public void waitForCompletion() throws InterruptedException {
+        if (runningThread == null) return;
+        runningThread.join();
+    }
+
+    private void internalRun() {
         while (!Thread.currentThread().isInterrupted() && !(writersToProcess.isEmpty() && stop)) {
-            Writer writer = writersToProcess.poll();
-            if (writer == null) {
-                try {
-                    Thread.sleep(10);
-                } catch (InterruptedException e) {
-                    Thread.currentThread().interrupt();
-                }
-
-                continue;
+            Writer writer;
+            try {
+                writer = writersToProcess.take();
+            } catch (InterruptedException e) {
+                return;
             }
 
             writer._threadWrite();
         }
     }
+
+    @Override
+    public void run() {
+        runningThread = Thread.currentThread();
+        internalRun();
+        runningThread = null;
+    }
 }