Move code and make sure we synchronize while doing all operations

2 files changed, 11 insertions, 19 deletions

diff --git a/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileDownloader.java b/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileDownloader.java
index 67fcff2224b..f7b951e89d8 100644
--- a/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileDownloader.java
+++ b/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileDownloader.java
@@ -119,14 +119,7 @@ public class FileDownloader implements AutoCloseable {
 
     /** Download, the future returned will be complete()d by receiving method in {@link FileReceiver} */
     private synchronized Future<Optional<File>> download(FileReferenceDownload fileReferenceDownload) {
-        FileReference fileReference = fileReferenceDownload.fileReference();
-        FileReferenceDownload inProgress = fileReferenceDownloader.getDownloadInProgress(fileReference);
-        return (inProgress == null) ? queueForDownload(fileReferenceDownload) : inProgress.future();
-    }
-
-    private Future<Optional<File>> queueForDownload(FileReferenceDownload fileReferenceDownload) {
-        fileReferenceDownloader.addToDownloadQueue(fileReferenceDownload);
-        return fileReferenceDownload.future();
+        return fileReferenceDownloader.download(fileReferenceDownload);
     }
 
     public FileReferenceDownloader fileReferenceDownloader() {
diff --git a/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceDownloader.java b/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceDownloader.java
index 71b6a08998c..815fdb194cf 100644
--- a/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceDownloader.java
+++ b/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceDownloader.java
@@ -6,6 +6,8 @@ import com.yahoo.config.FileReference;
 import com.yahoo.jrt.Int32Value;
 import com.yahoo.jrt.Request;
 import com.yahoo.jrt.StringValue;
+
+import java.util.concurrent.Future;
 import java.util.logging.Level;
 import com.yahoo.vespa.config.Connection;
 import com.yahoo.vespa.config.ConnectionPool;
@@ -20,7 +22,6 @@ import java.util.Optional;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
-import java.util.logging.Level;
 import java.util.logging.Logger;
 
 /**
@@ -79,14 +80,18 @@ public class FileReferenceDownloader {
         }
     }
 
-    void addToDownloadQueue(FileReferenceDownload fileReferenceDownload) {
-        FileReference fileReference = fileReferenceDownload.fileReference();
-        log.log(Level.FINE, () -> "Will download file reference '" + fileReference.value() + "' with timeout " + downloadTimeout);
+    Future<Optional<File>> download(FileReferenceDownload fileReferenceDownload) {
         synchronized (downloads) {
+            FileReference fileReference = fileReferenceDownload.fileReference();
+            FileReferenceDownload inProgress = downloads.get(fileReference);
+            if (inProgress != null) return inProgress.future();
+
+            log.log(Level.FINE, () -> "Will download file reference '" + fileReference.value() + "' with timeout " + downloadTimeout);
             downloads.put(fileReference, fileReferenceDownload);
             downloadStatus.put(fileReference, 0.0);
+            downloadExecutor.submit(() -> startDownload(fileReferenceDownload));
+            return fileReferenceDownload.future();
         }
-        downloadExecutor.submit(() -> startDownload(fileReferenceDownload));
     }
 
     void completedDownloading(FileReference fileReference, File file) {
@@ -143,12 +148,6 @@ public class FileReferenceDownloader {
         }
     }
 
-    FileReferenceDownload getDownloadInProgress(FileReference fileReference) {
-        synchronized (downloads) {
-            return downloads.get(fileReference);
-        }
-    }
-
     private boolean validateResponse(Request request) {
         if (request.isError()) {
             return false;