Merge pull request #4221 from vespa-engine/hmusum/download-from-another-config-server

Add support for downloading from another config server

5 files changed, 93 insertions, 12 deletions

diff --git a/configserver/pom.xml b/configserver/pom.xml
index 4ebb76bd5fe..66359e0bcc0 100644
--- a/configserver/pom.xml
+++ b/configserver/pom.xml
@@ -126,6 +126,11 @@
     </dependency>
     <dependency>
       <groupId>com.yahoo.vespa</groupId>
+      <artifactId>filedistribution</artifactId>
+      <version>6-SNAPSHOT</version>
+    </dependency>
+    <dependency>
+      <groupId>com.yahoo.vespa</groupId>
       <artifactId>filedistributionmanager</artifactId>
       <version>${project.version}</version>
     </dependency>
diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/filedistribution/FileServer.java b/configserver/src/main/java/com/yahoo/vespa/config/server/filedistribution/FileServer.java
index a504cd120ee..9dc94c9fe93 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/filedistribution/FileServer.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/filedistribution/FileServer.java
@@ -4,7 +4,10 @@ package com.yahoo.vespa.config.server.filedistribution;
 import com.google.inject.Inject;
 import com.yahoo.config.FileReference;
 import com.yahoo.config.model.api.FileDistribution;
+import com.yahoo.config.subscription.ConfigSourceSet;
 import com.yahoo.io.IOUtils;
+import com.yahoo.vespa.config.JRTConnectionPool;
+import com.yahoo.vespa.filedistribution.FileDownloader;
 
 import java.io.File;
 import java.io.IOException;
@@ -16,6 +19,7 @@ public class FileServer {
     private static final Logger log = Logger.getLogger(FileServer.class.getName());
     private final FileDirectory root;
     private final ExecutorService executor;
+    private final FileDownloader downloader = new FileDownloader(new JRTConnectionPool(ConfigSourceSet.createDefault()));
 
     public static class ReplayStatus {
         private final int code;
@@ -86,4 +90,8 @@ public class FileServer {
         // TODO remove once verified in system tests.
         log.info("Done serving reference '" + reference.toString() + "' with file '" + file.getAbsolutePath() + "'");
     }
+
+    public void download(FileReference fileReference) {
+        downloader.getFile(fileReference);
+    }
 }
diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/rpc/RpcServer.java b/configserver/src/main/java/com/yahoo/vespa/config/server/rpc/RpcServer.java
index 7c5adb3b932..d17cdf722ea 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/rpc/RpcServer.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/rpc/RpcServer.java
@@ -21,7 +21,6 @@ import com.yahoo.jrt.StringValue;
 import com.yahoo.jrt.Supervisor;
 import com.yahoo.jrt.Target;
 import com.yahoo.jrt.Transport;
-import com.yahoo.jrt.Value;
 import com.yahoo.log.LogLevel;
 import com.yahoo.vespa.config.ErrorCode;
 import com.yahoo.vespa.config.JRTMethods;
@@ -249,7 +248,6 @@ public class RpcServer implements Runnable, ReloadListener, TenantListener {
         }
 
         for (int i = 0; i < responsesSent; i++) {
-
             try {
                 completionService.take();
             } catch (InterruptedException e) {
@@ -469,6 +467,8 @@ public class RpcServer implements Runnable, ReloadListener, TenantListener {
                     : FileApiErrorCodes.NOT_FOUND;
             if (result == FileApiErrorCodes.OK) {
                 fileServer.startFileServing(fileReference, new FileReceiver(request.target()));
+            } else {
+                fileServer.download(new FileReference(fileReference));
             }
         } catch (IllegalArgumentException e) {
             result = FileApiErrorCodes.NOT_FOUND;
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 fbadddd624a..08595662f36 100644
--- a/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceDownloader.java
+++ b/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceDownloader.java
@@ -5,6 +5,7 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.yahoo.concurrent.DaemonThreadFactory;
 import com.yahoo.config.FileReference;
+import com.yahoo.jrt.ErrorCode;
 import com.yahoo.jrt.Request;
 import com.yahoo.jrt.StringValue;
 import com.yahoo.log.LogLevel;
@@ -35,7 +36,6 @@ import java.util.stream.Collectors;
  *
  * @author hmusum
  */
-// TODO: Add retries when a config server does not have a file reference
 // TODO: Handle shutdown of executors
 class FileReferenceDownloader {
 
@@ -66,10 +66,20 @@ class FileReferenceDownloader {
             throws ExecutionException, InterruptedException, TimeoutException {
         downloads.put(fileReference, fileReferenceDownload);
         setDownloadStatus(fileReference.value(), 0.0);
-        if (startDownloadRpc(fileReference))
+
+        int numAttempts = 0;
+        boolean downloadStarted = false;
+        do {
+            if (startDownloadRpc(fileReference))
+                downloadStarted = true;
+            else
+                Thread.sleep(100);
+        } while (!downloadStarted && ++numAttempts <= 10);  // TODO: How long/many times to retry?
+
+        if (downloadStarted) {
             return fileReferenceDownload.future().get(timeout.toMillis(), TimeUnit.MILLISECONDS);
-        else {
-            fileReferenceDownload.future().setException(new RuntimeException("Failed getting file"));
+        } else {
+            fileReferenceDownload.future().setException(new RuntimeException("Failed getting file reference '" + fileReference.value() + "'"));
             downloads.remove(fileReference);
             return Optional.empty();
         }
@@ -118,15 +128,17 @@ class FileReferenceDownloader {
         execute(request, connection);
         if (validateResponse(request)) {
             log.log(LogLevel.DEBUG, "Request callback, OK. Req: " + request + "\nSpec: " + connection);
-            if (request.returnValues().get(0).asInt32() == 0)
+            if (request.returnValues().get(0).asInt32() == 0) {
                 log.log(LogLevel.INFO, "Found file reference '" + fileReference.value() + "' available at " + connection.getAddress());
-            else
+                return true;
+            } else {
                 log.log(LogLevel.INFO, "File reference '" + fileReference.value() + "' not found for " + connection.getAddress());
-            return true;
+                return false;
+            }
         } else {
             log.log(LogLevel.WARNING, "Request failed. Req: " + request + "\nSpec: " + connection.getAddress());
-            connection.setError(request.errorCode());
-            // TODO: Retry with another config server
+            if (request.isError() && request.errorCode() == ErrorCode.CONNECTION)
+                connection.setError(request.errorCode());
             return false;
         }
     }
diff --git a/filedistribution/src/test/java/com/yahoo/vespa/filedistribution/FileDownloaderTest.java b/filedistribution/src/test/java/com/yahoo/vespa/filedistribution/FileDownloaderTest.java
index 738b0888956..278c46dab8b 100644
--- a/filedistribution/src/test/java/com/yahoo/vespa/filedistribution/FileDownloaderTest.java
+++ b/filedistribution/src/test/java/com/yahoo/vespa/filedistribution/FileDownloaderTest.java
@@ -27,6 +27,7 @@ import java.util.Arrays;
 import java.util.List;
 import java.util.Optional;
 
+import static com.yahoo.jrt.ErrorCode.CONNECTION;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
@@ -45,7 +46,7 @@ public class FileDownloaderTest {
         try {
             downloadDir = Files.createTempDirectory("filedistribution").toFile();
             connection = new MockConnection();
-            fileDownloader = new FileDownloader(connection, downloadDir, Duration.ofMillis(3000));
+            fileDownloader = new FileDownloader(connection, downloadDir, Duration.ofMillis(2000));
         } catch (IOException e) {
             e.printStackTrace();
             fail(e.getMessage());
@@ -115,6 +116,38 @@ public class FileDownloaderTest {
     }
 
     @Test
+    public void getFileWhenConnectionError() throws IOException {
+        fileDownloader = new FileDownloader(connection, downloadDir, Duration.ofMillis(3000));
+        File downloadDir = fileDownloader.downloadDirectory();
+
+        int timesToFail = 2;
+        MockConnection.ConnectionErrorResponseHandler responseHandler = new MockConnection.ConnectionErrorResponseHandler(timesToFail);
+        connection.setResponseHandler(responseHandler);
+
+        FileReference fileReference = new FileReference("fileReference");
+        File fileReferenceFullPath = fileReferenceFullPath(downloadDir, fileReference);
+        assertFalse(fileReferenceFullPath.getAbsolutePath(), fileDownloader.getFile(fileReference).isPresent());
+
+        // Verify download status
+        assertDownloadStatus(fileDownloader, fileReference, 0.0);
+
+        // Receives fileReference, should return and make it available to caller
+        String filename = "abc.jar";
+        receiveFile(fileReference, filename, "some other content");
+        Optional<File> downloadedFile = fileDownloader.getFile(fileReference);
+
+        assertTrue(downloadedFile.isPresent());
+        File downloadedFileFullPath = new File(fileReferenceFullPath, filename);
+        assertEquals(downloadedFileFullPath.getAbsolutePath(), downloadedFile.get().getAbsolutePath());
+        assertEquals("some other content", IOUtils.readFile(downloadedFile.get()));
+
+        // Verify download status when downloaded
+        assertDownloadStatus(fileDownloader, fileReference, 100.0);
+
+        assertEquals(timesToFail, responseHandler.failedTimes);
+    }
+
+    @Test
     public void setFilesToDownload() throws IOException {
         Duration timeout = Duration.ofMillis(200);
         File downloadDir = Files.createTempDirectory("filedistribution").toFile();
@@ -271,7 +304,30 @@ public class FileDownloaderTest {
                     request.returnValues().add(new Int32Value(0));
                     request.returnValues().add(new StringValue("OK"));
                 }
+            }
+        }
+
+        static class ConnectionErrorResponseHandler implements MockConnection.ResponseHandler {
 
+            private final int timesToFail;
+            private int failedTimes = 0;
+
+            ConnectionErrorResponseHandler(int timesToFail) {
+                super();
+                this.timesToFail = timesToFail;
+            }
+
+            @Override
+            public void request(Request request) {
+                if (request.methodName().equals("filedistribution.serveFile")) {
+                    if (failedTimes < timesToFail) {
+                        request.setError(CONNECTION, "Connection error");
+                        failedTimes++;
+                    } else {
+                        request.returnValues().add(new Int32Value(0));
+                        request.returnValues().add(new StringValue("OK"));
+                    }
+                }
             }
         }
     }