Closing now actually works

jean-philippe-martin · jean-philippe-martin · commit 00e0ab8df8d0 · 2016-11-18T17:33:49.000-08:00
diff --git a/google-cloud-contrib/google-cloud-nio/src/main/java/com/google/cloud/storage/contrib/nio/SeekableByteChannelPrefetcher.java b/google-cloud-contrib/google-cloud-nio/src/main/java/com/google/cloud/storage/contrib/nio/SeekableByteChannelPrefetcher.java
@@ -16,6 +16,7 @@
 
 package com.google.cloud.storage.contrib.nio;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.Futures;
 
 import java.io.Closeable;
@@ -72,6 +73,7 @@ public class SeekableByteChannelPrefetcher implements SeekableByteChannel {
   private final ExecutorService exec;
   private final Sorted buffers;
   private final ArrayList<Worker> idleWorkers;
+  private final int workerCount;
 
   private class Buffer {
     // index*bufferSize = file position. Set to -1 when we haven't yet decided.
@@ -139,6 +141,7 @@ public void init(long pos, Buffer buf) throws ExecutionException, InterruptedExc
 
     public ByteBuffer call() throws IOException, ExecutionException, InterruptedException {
       if (pos > chan.size()) {
+        reassignWorker(this);
         return null;
       }
       chan.position(pos);
@@ -163,10 +166,12 @@ public SeekableByteChannelPrefetcher(SeekableByteChannelPrefetcherOptions opts,
     this.idleWorkers = new ArrayList<>(this.prefetchingThreads + this.extraThreads);
     this.exec = Executors.newFixedThreadPool(prefetchingThreads + extraThreads);
     SeekableByteChannel chan = null;
-    for (SeekableByteChannel bc : channels) {
+    ImmutableList<SeekableByteChannel> underlyingChannels = ImmutableList.copyOf(channels);
+    for (SeekableByteChannel bc : underlyingChannels) {
       chan = bc;
       idleWorkers.add(new Worker(bc));
     }
+    workerCount = underlyingChannels.size();
     size = chan.size();
     position = 0;
   }
@@ -323,20 +328,27 @@ public boolean isOpen() {
   public void close() throws IOException {
     if (open) {
       closing = true;
-      while (true) {
-        synchronized (idleWorkers) {
-          if (idleWorkers.size() == prefetchingThreads + extraThreads) {
+      exec.shutdown();
+      try {
+        while (true) synchronized (idleWorkers) {
+          if (idleWorkers.size() >= workerCount) {
             // every thread is idle, we're done.
             break;
           }
+          idleWorkers.wait(60_000);
         }
+      } catch (InterruptedException e) {
+        System.out.println("Timed out while waiting for channels to close.");
       }
-      exec.shutdown();
       try {
         exec.awaitTermination(60, TimeUnit.SECONDS);
       } catch (InterruptedException e) {
         exec.shutdownNow();
       }
+      // Close all underlying channels
+      for (Worker w : idleWorkers) {
+        w.close();
+      }
       open = false;
     }
   }
@@ -417,7 +429,9 @@ private void reassignWorker(Worker w) throws ExecutionException, InterruptedExce
     long curIndex = index(position);
     if (!closing) {
       for (int i = 0; i < prefetchingThreads; i++) {
-        if (i > lastIndex) break;
+        if (i > lastIndex) {
+          break;
+        }
         if (buffers.get(curIndex + i) == null) {
           // work for you!
           Buffer buf = getEmptyBuffer();
@@ -440,7 +454,9 @@ private void startPrefetching(long position) throws ExecutionException, Interrup
     long lastIndex = index(size);
     long curIndex = index(position);
     for (int i = 0; i < prefetchingThreads; i++) {
-      if (i > lastIndex) break;
+      if (i > lastIndex) {
+        break;
+      }
       if (buffers.get(curIndex + i) == null) {
         // work available!
         Worker w = tryGetIdleWorker();
@@ -449,7 +465,6 @@ private void startPrefetching(long position) throws ExecutionException, Interrup
         }
         Buffer buf = getEmptyBuffer();
         sicWorker(w, bufferSize * (curIndex + i), buf);
-        return;
       }
     }
   }
diff --git a/google-cloud-contrib/google-cloud-nio/src/main/java/com/google/cloud/storage/contrib/nio/SeekableByteChannelPrefetcherOptions.java b/google-cloud-contrib/google-cloud-nio/src/main/java/com/google/cloud/storage/contrib/nio/SeekableByteChannelPrefetcherOptions.java
@@ -22,13 +22,13 @@
  */
 public class SeekableByteChannelPrefetcherOptions implements java.nio.file.OpenOption {
   // normal-case number of parallel reads.
-  public int prefetchingThreads = 4;
+  public int prefetchingThreads = 2;
   // in case the data we need isn't being prefetched, we can use up to this many
   // extra threads to fetch user-requested data.
   public int extraThreads = 1;
   // size in bytes for our buffer. Every fetcher grabs one buffer at a time.
   public int bufferSize = 50 * 1024 * 1024;
   // how many buffers we keep around. Should be at least prefetchingThreads + extraThreads.
   // bufferSize * bufferCount is how much memory this class'll allocate.
-  public int bufferCount = 6;
+  public int bufferCount = 4;
 }
