Allow set of distributable files to change after initial deploy.

This is necessary as bootstrapping the configserver with a new version might change set of files to distribute.

14 files changed, 114 insertions, 39 deletions

diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/filedistribution/FileDBRegistry.java b/configserver/src/main/java/com/yahoo/vespa/config/server/filedistribution/FileDBRegistry.java
index 1aecc4cde7d..5ec83529601 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/filedistribution/FileDBRegistry.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/filedistribution/FileDBRegistry.java
@@ -7,12 +7,16 @@ import com.yahoo.net.HostName;
 import com.yahoo.text.Utf8;
 import net.jpountz.xxhash.XXHashFactory;
 
+import java.io.BufferedReader;
+import java.io.IOException;
+import java.io.Reader;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.regex.Pattern;
 
 /**
  * @author Tony Vaagenes
@@ -20,19 +24,51 @@ import java.util.Optional;
 public class FileDBRegistry implements FileRegistry {
 
     private final AddFileInterface manager;
-    private final List<Entry> entries = new ArrayList<>();
     private final Map<String, FileReference> fileReferenceCache = new HashMap<>();
+    private static final String entryDelimiter = "\t";
+    private static final Pattern entryDelimiterPattern = Pattern.compile(entryDelimiter, Pattern.LITERAL);
 
     public FileDBRegistry(AddFileInterface manager) {
         this.manager = manager;
     }
+    public static FileRegistry create(AddFileInterface manager, Reader persistedState) {
+        try (BufferedReader reader = new BufferedReader(persistedState)) {
+            String ignoredFileSourceHost = reader.readLine();
+            if (ignoredFileSourceHost == null)
+                throw new RuntimeException("No file source host");
+            return new FileDBRegistry(manager, decode(reader));
+        } catch (IOException e) {
+            throw new RuntimeException("Error while reading pre-generated file registry", e);
+        }
+    }
+
+    static Map<String, FileReference> decode(BufferedReader reader) {
+        Map<String, FileReference> refs = new HashMap<>();
+        try {
+            String line;
+            while ((line = reader.readLine()) != null) {
+                String[] parts = entryDelimiterPattern.split(line);
+                if (parts.length < 2)
+                    throw new IllegalArgumentException("Cannot split '" + line + "' into two parts");
+                refs.put(parts[0], new FileReference(parts[1]));
+            }
+        } catch (IOException e) {
+            throw new RuntimeException("Error while reading pre-generated file registry", e);
+        }
+        return refs;
+    }
+    private FileDBRegistry(AddFileInterface manager, Map<String, FileReference> knownReferences) {
+        this.manager = manager;
+        for (Map.Entry<String, FileReference> e : knownReferences.entrySet()) {
+            fileReferenceCache.put(e.getKey(), e.getValue());
+        }
+    }
 
     @Override
     public synchronized FileReference addFile(String relativePath) {
         Optional<FileReference> cachedReference = Optional.ofNullable(fileReferenceCache.get(relativePath));
         return cachedReference.orElseGet(() -> {
             FileReference newRef = manager.addFile(relativePath);
-            entries.add(new Entry(relativePath, newRef));
             fileReferenceCache.put(relativePath, newRef);
             return newRef;
         });
@@ -44,7 +80,6 @@ public class FileDBRegistry implements FileRegistry {
         Optional<FileReference> cachedReference = Optional.ofNullable(fileReferenceCache.get(uri));
         return cachedReference.orElseGet(() -> {
             FileReference newRef = manager.addUri(uri, relativePath);
-            entries.add(new Entry(uri, newRef));
             fileReferenceCache.put(uri, newRef);
             return newRef;
         });
@@ -58,7 +93,6 @@ public class FileDBRegistry implements FileRegistry {
             Optional<FileReference> cachedReference = Optional.ofNullable(fileReferenceCache.get(blobName));
             return cachedReference.orElseGet(() -> {
                 FileReference newRef = manager.addBlob(blob, relativePath);
-                entries.add(new Entry(blobName, newRef));
                 fileReferenceCache.put(blobName, newRef);
                 return newRef;
             });
@@ -66,12 +100,11 @@ public class FileDBRegistry implements FileRegistry {
     }
 
     @Override
-    public String fileSourceHost() {
-        return HostName.getLocalhost();
-    }
-
-    @Override
     public synchronized List<Entry> export() {
+        List<Entry> entries = new ArrayList<>();
+        for (Map.Entry<String, FileReference> entry : fileReferenceCache.entrySet()) {
+            entries.add(new Entry(entry.getKey(), entry.getValue()));
+        }
         return entries;
     }
 
diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/filedistribution/FileDistributionProvider.java b/configserver/src/main/java/com/yahoo/vespa/config/server/filedistribution/FileDistributionProvider.java
index b3f3214793c..914d4f48571 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/filedistribution/FileDistributionProvider.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/filedistribution/FileDistributionProvider.java
@@ -5,6 +5,7 @@ import com.yahoo.config.model.api.FileDistribution;
 import com.yahoo.config.application.api.FileRegistry;
 
 import java.io.File;
+import java.io.Reader;
 
 /**
  * Provides file distribution registry and invoker.
@@ -15,13 +16,23 @@ public class FileDistributionProvider {
 
     private final FileRegistry fileRegistry;
     private final FileDistribution fileDistribution;
+    private final AddFileInterface fileManager;
 
     public FileDistributionProvider(File applicationDir, FileDistribution fileDistribution) {
-        this(new FileDBRegistry(new ApplicationFileManager(applicationDir, new FileDirectory(fileDistribution.getFileReferencesDir()))), fileDistribution);
+        this(new ApplicationFileManager(applicationDir, new FileDirectory(fileDistribution.getFileReferencesDir())), fileDistribution);
         ensureDirExists(fileDistribution.getFileReferencesDir());
     }
 
     FileDistributionProvider(FileRegistry fileRegistry, FileDistribution fileDistribution) {
+       this(null, fileRegistry, fileDistribution);
+    }
+
+    private FileDistributionProvider(AddFileInterface fileManager, FileDistribution fileDistribution) {
+        this(fileManager, new FileDBRegistry(fileManager), fileDistribution);
+    }
+
+    private FileDistributionProvider(AddFileInterface fileManager, FileRegistry fileRegistry, FileDistribution fileDistribution) {
+        this.fileManager = fileManager;
         this.fileRegistry = fileRegistry;
         this.fileDistribution = fileDistribution;
     }
@@ -34,6 +45,10 @@ public class FileDistributionProvider {
         return fileDistribution;
     }
 
+    public FileRegistry createPregeneratedFileRegistry(Reader persistedState) {
+        return FileDBRegistry.create(fileManager, persistedState);
+    }
+
     private static void ensureDirExists(File dir) {
         if (!dir.exists()) {
             boolean success = dir.mkdirs();
diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionPreparer.java b/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionPreparer.java
index c2fc64c615c..b9f79763a41 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionPreparer.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionPreparer.java
@@ -104,6 +104,10 @@ public class SessionPreparer {
         this.flagSource = flagSource;
     }
 
+    public FileDistributionFactory getFileDistributionFactory() {
+        return fileDistributionFactory;
+    }
+
     /**
      * Prepares a session (validates, builds model, writes to zookeeper and distributes files)
      *
diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionRepository.java b/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionRepository.java
index 72e7391098d..398695a5d6e 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionRepository.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionRepository.java
@@ -31,6 +31,7 @@ import com.yahoo.vespa.config.server.application.TenantApplications;
 import com.yahoo.vespa.config.server.configchange.ConfigChangeActions;
 import com.yahoo.vespa.config.server.deploy.TenantFileSystemDirs;
 import com.yahoo.vespa.config.server.filedistribution.FileDirectory;
+import com.yahoo.vespa.config.server.filedistribution.FileDistributionFactory;
 import com.yahoo.vespa.config.server.http.UnknownVespaVersionException;
 import com.yahoo.vespa.config.server.modelfactory.ActivatedModelsBuilder;
 import com.yahoo.vespa.config.server.modelfactory.ModelFactoryRegistry;
@@ -103,6 +104,7 @@ public class SessionRepository {
     private final Clock clock;
     private final Curator curator;
     private final Executor zkWatcherExecutor;
+    private final FileDistributionFactory fileDistributionFactory;
     private final PermanentApplicationPackage permanentApplicationPackage;
     private final FlagSource flagSource;
     private final TenantFileSystemDirs tenantFileSystemDirs;
@@ -129,6 +131,7 @@ public class SessionRepository {
                              Curator curator,
                              Metrics metrics,
                              StripedExecutor<TenantName> zkWatcherExecutor,
+                             FileDistributionFactory fileDistributionFactory,
                              PermanentApplicationPackage permanentApplicationPackage,
                              FlagSource flagSource,
                              ExecutorService zkCacheExecutor,
@@ -148,6 +151,7 @@ public class SessionRepository {
         this.curator = curator;
         this.sessionLifetime = Duration.ofSeconds(configserverConfig.sessionLifetime());
         this.zkWatcherExecutor = command -> zkWatcherExecutor.execute(tenantName, command);
+        this.fileDistributionFactory = fileDistributionFactory;
         this.permanentApplicationPackage = permanentApplicationPackage;
         this.flagSource = flagSource;
         this.tenantFileSystemDirs = new TenantFileSystemDirs(configServerDB, tenantName);
@@ -781,7 +785,8 @@ public class SessionRepository {
 
     private SessionZooKeeperClient createSessionZooKeeperClient(long sessionId) {
         String serverId = configserverConfig.serverId();
-        return new SessionZooKeeperClient(curator, tenantName, sessionId, serverId, maxNodeSize);
+        return new SessionZooKeeperClient(curator, tenantName, sessionId, serverId,
+                fileDistributionFactory.createProvider(getSessionAppDir(sessionId)), maxNodeSize);
     }
 
     private File getAndValidateExistingSessionAppDir(long sessionId) {
diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionZooKeeperClient.java b/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionZooKeeperClient.java
index 1b23c3bf37a..2d2712b3002 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionZooKeeperClient.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/session/SessionZooKeeperClient.java
@@ -21,6 +21,7 @@ import com.yahoo.transaction.Transaction;
 import com.yahoo.vespa.config.server.UserConfigDefinitionRepo;
 import com.yahoo.vespa.config.server.deploy.ZooKeeperClient;
 import com.yahoo.vespa.config.server.deploy.ZooKeeperDeployer;
+import com.yahoo.vespa.config.server.filedistribution.FileDistributionProvider;
 import com.yahoo.vespa.config.server.tenant.OperatorCertificateSerializer;
 import com.yahoo.vespa.config.server.tenant.TenantRepository;
 import com.yahoo.vespa.config.server.tenant.TenantSecretStoreSerializer;
@@ -68,18 +69,20 @@ public class SessionZooKeeperClient {
     private final Path sessionStatusPath;
     private final String serverId;  // hostname
     private final int maxNodeSize;
+    private final FileDistributionProvider fileDistributionProvider;
 
-    public SessionZooKeeperClient(Curator curator, TenantName tenantName, long sessionId, String serverId, int maxNodeSize) {
+    public SessionZooKeeperClient(Curator curator, TenantName tenantName, long sessionId, String serverId, FileDistributionProvider fileDistributionProvider, int maxNodeSize) {
         this.curator = curator;
         this.tenantName = tenantName;
         this.sessionPath = getSessionPath(tenantName, sessionId);
         this.serverId = serverId;
         this.sessionStatusPath = sessionPath.append(ZKApplication.SESSIONSTATE_ZK_SUBPATH);
         this.maxNodeSize = maxNodeSize;
+        this.fileDistributionProvider = fileDistributionProvider;
     }
 
-    public SessionZooKeeperClient(Curator curator, TenantName tenantName, long sessionId, String serverId) {
-        this(curator, tenantName, sessionId, serverId, 10 * 1024 * 1024);
+    public SessionZooKeeperClient(Curator curator, TenantName tenantName, long sessionId, FileDistributionProvider fileDistributionProvider, String serverId) {
+        this(curator, tenantName, sessionId, serverId, fileDistributionProvider, 10 * 1024 * 1024);
     }
 
     public void writeStatus(Session.Status sessionStatus) {
@@ -138,7 +141,7 @@ public class SessionZooKeeperClient {
     }
 
     public ApplicationPackage loadApplicationPackage() {
-        return new ZKApplicationPackage(curator, sessionPath, maxNodeSize);
+        return new ZKApplicationPackage(fileDistributionProvider, curator, sessionPath, maxNodeSize);
     }
 
     public ConfigDefinitionRepo getUserConfigDefinitions() {
diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/tenant/TenantRepository.java b/configserver/src/main/java/com/yahoo/vespa/config/server/tenant/TenantRepository.java
index b10a90de714..3b16e9aa6c1 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/tenant/TenantRepository.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/tenant/TenantRepository.java
@@ -350,6 +350,7 @@ public class TenantRepository {
                                                                     curator,
                                                                     metrics,
                                                                     zkSessionWatcherExecutor,
+                                                                    fileDistributionFactory,
                                                                     permanentApplicationPackage,
                                                                     flagSource,
                                                                     zkCacheExecutor,
@@ -564,6 +565,10 @@ public class TenantRepository {
         return ImmutableSet.copyOf(tenants.values());
     }
 
+    public FileDistributionFactory getFileDistributionFactory() {
+        return fileDistributionFactory;
+    }
+
     /**
      * Gets zookeeper path for tenant data
      *
diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/zookeeper/ZKApplicationPackage.java b/configserver/src/main/java/com/yahoo/vespa/config/server/zookeeper/ZKApplicationPackage.java
index c88b0cf7099..38bed8ee05d 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/zookeeper/ZKApplicationPackage.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/zookeeper/ZKApplicationPackage.java
@@ -10,7 +10,6 @@ import com.yahoo.config.application.api.ComponentInfo;
 import com.yahoo.config.application.api.FileRegistry;
 import com.yahoo.config.application.api.UnparsedConfigDefinition;
 import com.yahoo.config.codegen.DefParser;
-import com.yahoo.config.model.application.provider.PreGeneratedFileRegistry;
 import com.yahoo.config.provision.AllocatedHosts;
 import com.yahoo.config.provision.ApplicationId;
 import com.yahoo.config.provision.serialization.AllocatedHostsSerializer;
@@ -20,6 +19,7 @@ import com.yahoo.path.Path;
 import com.yahoo.vespa.config.ConfigDefinition;
 import com.yahoo.vespa.config.ConfigDefinitionBuilder;
 import com.yahoo.vespa.config.ConfigDefinitionKey;
+import com.yahoo.vespa.config.server.filedistribution.FileDistributionProvider;
 import com.yahoo.vespa.config.util.ConfigUtils;
 import com.yahoo.vespa.curator.Curator;
 
@@ -54,17 +54,17 @@ public class ZKApplicationPackage implements ApplicationPackage {
     public static final String allocatedHostsNode = "allocatedHosts";
     private final ApplicationMetaData metaData;
 
-    public ZKApplicationPackage(Curator curator, Path sessionPath, int maxNodeSize) {
+    public ZKApplicationPackage(FileDistributionProvider fileDistributionProvider, Curator curator, Path sessionPath, int maxNodeSize) {
         verifyAppPath(curator, sessionPath);
         zkApplication = new ZKApplication(curator, sessionPath, maxNodeSize);
         metaData = readMetaDataFromLiveApp(zkApplication);
-        importFileRegistries();
+        importFileRegistries(fileDistributionProvider);
         allocatedHosts = importAllocatedHosts();
     }
 
     // For testing
-    ZKApplicationPackage(Curator curator, Path sessionPath) {
-        this(curator, sessionPath, 10 * 1024 * 1024);
+    ZKApplicationPackage(FileDistributionProvider fileDistributionProvider, Curator curator, Path sessionPath) {
+        this(fileDistributionProvider, curator, sessionPath, 10 * 1024 * 1024);
     }
 
     private Optional<AllocatedHosts> importAllocatedHosts() {
@@ -85,17 +85,17 @@ public class ZKApplicationPackage implements ApplicationPackage {
         }
     }
 
-    private void importFileRegistries() {
+    private void importFileRegistries(FileDistributionProvider fileDistributionProvider) {
         List<String> perVersionFileRegistryNodes = zkApplication.getChildren(Path.fromString(fileRegistryNode));
         perVersionFileRegistryNodes
                 .forEach(version ->
                                  fileRegistryMap.put(Version.fromString(version),
-                                                     importFileRegistry(Joiner.on("/").join(fileRegistryNode, version))));
+                                                     importFileRegistry(fileDistributionProvider, Joiner.on("/").join(fileRegistryNode, version))));
     }
 
-    private FileRegistry importFileRegistry(String fileRegistryNode) {
+    private FileRegistry importFileRegistry(FileDistributionProvider fileDistributionProvider, String fileRegistryNode) {
         try {
-            return PreGeneratedFileRegistry.importRegistry(zkApplication.getDataReader(Path.fromString(fileRegistryNode)));
+            return fileDistributionProvider.createPregeneratedFileRegistry(zkApplication.getDataReader(Path.fromString(fileRegistryNode)));
         } catch (Exception e) {
             throw new RuntimeException("Could not determine which files to distribute", e);
         }
diff --git a/configserver/src/test/java/com/yahoo/vespa/config/server/ApplicationRepositoryTest.java b/configserver/src/test/java/com/yahoo/vespa/config/server/ApplicationRepositoryTest.java
index 1db8b9a2741..665b603acfc 100644
--- a/configserver/src/test/java/com/yahoo/vespa/config/server/ApplicationRepositoryTest.java
+++ b/configserver/src/test/java/com/yahoo/vespa/config/server/ApplicationRepositoryTest.java
@@ -473,6 +473,7 @@ public class ApplicationRepositoryTest {
                                                       new SessionZooKeeperClient(curator,
                                                                                  tenantName,
                                                                                  sessionId,
+                                                                                 tenantRepository.getFileDistributionFactory().createProvider(testApp),
                                                                                  ConfigUtils.getCanonicalHostName()));
         sessionRepository.addLocalSession(localSession2);
         assertEquals(2, sessionRepository.getLocalSessions().size());
diff --git a/configserver/src/test/java/com/yahoo/vespa/config/server/filedistribution/MockFileRegistry.java b/configserver/src/test/java/com/yahoo/vespa/config/server/filedistribution/MockFileRegistry.java
index 723adc1400b..b11308214a3 100644
--- a/configserver/src/test/java/com/yahoo/vespa/config/server/filedistribution/MockFileRegistry.java
+++ b/configserver/src/test/java/com/yahoo/vespa/config/server/filedistribution/MockFileRegistry.java
@@ -37,9 +37,6 @@ public class MockFileRegistry implements FileRegistry {
         return fileReference;
     }
 
-    @Override
-    public String fileSourceHost() { return HostName.getLocalhost(); }
-
     public List<Entry> export() { return entries; }
 
     @Override
diff --git a/configserver/src/test/java/com/yahoo/vespa/config/server/http/v2/SessionCreateHandlerTest.java b/configserver/src/test/java/com/yahoo/vespa/config/server/http/v2/SessionCreateHandlerTest.java
index 5d7322070e7..3ecf76544c4 100644
--- a/configserver/src/test/java/com/yahoo/vespa/config/server/http/v2/SessionCreateHandlerTest.java
+++ b/configserver/src/test/java/com/yahoo/vespa/config/server/http/v2/SessionCreateHandlerTest.java
@@ -71,11 +71,14 @@ public class SessionCreateHandlerTest extends SessionHandlerTest {
     @Rule
     public TemporaryFolder temporaryFolder = new TemporaryFolder();
 
+    private ConfigserverConfig configserverConfig;
+
     @Before
     public void setupRepo() throws IOException {
-        ConfigserverConfig configserverConfig = new ConfigserverConfig.Builder()
+        configserverConfig = new ConfigserverConfig.Builder()
                 .configServerDBDir(temporaryFolder.newFolder().getAbsolutePath())
                 .configDefinitionsDir(temporaryFolder.newFolder().getAbsolutePath())
+                .fileReferencesDir(temporaryFolder.newFolder().getAbsolutePath())
                 .build();
         TenantRepository tenantRepository = new TestTenantRepository.Builder()
                 .withConfigserverConfig(configserverConfig)
@@ -196,7 +199,7 @@ public class SessionCreateHandlerTest extends SessionHandlerTest {
     private SessionCreateHandler createHandler() {
         return new SessionCreateHandler(SessionCreateHandler.testOnlyContext(),
                                         applicationRepository,
-                                        new ConfigserverConfig.Builder().build());
+                                        configserverConfig);
     }
 
     private HttpRequest post() throws FileNotFoundException {
diff --git a/configserver/src/test/java/com/yahoo/vespa/config/server/session/SessionPreparerTest.java b/configserver/src/test/java/com/yahoo/vespa/config/server/session/SessionPreparerTest.java
index a9e7ffd5b5a..cb3f94100d1 100644
--- a/configserver/src/test/java/com/yahoo/vespa/config/server/session/SessionPreparerTest.java
+++ b/configserver/src/test/java/com/yahoo/vespa/config/server/session/SessionPreparerTest.java
@@ -188,7 +188,7 @@ public class SessionPreparerTest {
         HostRegistry hostValidator = new HostRegistry();
         hostValidator.update(applicationId("foo"), Collections.singletonList("mytesthost"));
         preparer.prepare(hostValidator, new BaseDeployLogger(), new PrepareParams.Builder().applicationId(applicationId("default")).build(),
-                         Optional.empty(), Instant.now(), app.getAppDir(), app, createSessionZooKeeperClient());
+                         Optional.empty(), Instant.now(), app.getAppDir(), app, createSessionZooKeeperClient(app.getAppDir()));
     }
     
     @Test
@@ -203,7 +203,7 @@ public class SessionPreparerTest {
         hostValidator.update(applicationId, Collections.singletonList("mytesthost"));
         preparer.prepare(hostValidator, logger, new PrepareParams.Builder().applicationId(applicationId).build(),
                          Optional.empty(), Instant.now(), app.getAppDir(), app,
-                         createSessionZooKeeperClient());
+                         createSessionZooKeeperClient(app.getAppDir()));
         assertEquals(logged.toString(), "");
     }
 
@@ -212,7 +212,7 @@ public class SessionPreparerTest {
         PrepareParams params = new PrepareParams.Builder().applicationId(applicationId()).build();
         int sessionId = 1;
         prepare(testApp, params);
-        assertThat(createSessionZooKeeperClient(sessionId).readApplicationId().get(), is(applicationId()));
+        assertThat(createSessionZooKeeperClient(testApp, sessionId).readApplicationId().get(), is(applicationId()));
     }
 
     @Test
@@ -345,7 +345,7 @@ public class SessionPreparerTest {
         FilesApplicationPackage applicationPackage = getApplicationPackage(app);
         return preparer.prepare(new HostRegistry(), getLogger(), params,
                                 Optional.empty(), Instant.now(), applicationPackage.getAppDir(),
-                                applicationPackage, createSessionZooKeeperClient(sessionId));
+                                applicationPackage, createSessionZooKeeperClient(applicationPackage.getAppDir(), sessionId));
     }
 
     private FilesApplicationPackage getApplicationPackage(File testFile) throws IOException {
@@ -369,12 +369,12 @@ public class SessionPreparerTest {
                                   ApplicationName.from(applicationName), InstanceName.defaultName());
     }
 
-    private SessionZooKeeperClient createSessionZooKeeperClient() {
-        return createSessionZooKeeperClient(1);
+    private SessionZooKeeperClient createSessionZooKeeperClient(File appDir) {
+        return createSessionZooKeeperClient(appDir, 1);
     }
 
-    private SessionZooKeeperClient createSessionZooKeeperClient(long sessionId) {
-        return new SessionZooKeeperClient(curator, applicationId().tenant(), sessionId, ConfigUtils.getCanonicalHostName());
+    private SessionZooKeeperClient createSessionZooKeeperClient(File appDir, long sessionId) {
+        return new SessionZooKeeperClient(curator, applicationId().tenant(), sessionId, preparer.getFileDistributionFactory().createProvider(appDir), ConfigUtils.getCanonicalHostName());
     }
 
     private Path sessionPath(long sessionId) {
diff --git a/configserver/src/test/java/com/yahoo/vespa/config/server/session/SessionRepositoryTest.java b/configserver/src/test/java/com/yahoo/vespa/config/server/session/SessionRepositoryTest.java
index c6105898794..850c24fd3e6 100644
--- a/configserver/src/test/java/com/yahoo/vespa/config/server/session/SessionRepositoryTest.java
+++ b/configserver/src/test/java/com/yahoo/vespa/config/server/session/SessionRepositoryTest.java
@@ -286,6 +286,7 @@ public class SessionRepositoryTest {
         SessionZooKeeperClient zkc = new SessionZooKeeperClient(curator,
                                                                 tenantName,
                                                                 sessionId,
+                                                                tenantRepository.getFileDistributionFactory().createProvider(new File("dummy")),
                                                                 ConfigUtils.getCanonicalHostName());
         zkc.createNewSession(Instant.now());
         if (wait) {
diff --git a/configserver/src/test/java/com/yahoo/vespa/config/server/session/SessionZooKeeperClientTest.java b/configserver/src/test/java/com/yahoo/vespa/config/server/session/SessionZooKeeperClientTest.java
index 57ef55cc890..b8b7abf2d72 100644
--- a/configserver/src/test/java/com/yahoo/vespa/config/server/session/SessionZooKeeperClientTest.java
+++ b/configserver/src/test/java/com/yahoo/vespa/config/server/session/SessionZooKeeperClientTest.java
@@ -8,6 +8,8 @@ import com.yahoo.config.provision.ApplicationId;
 import com.yahoo.config.provision.TenantName;
 import com.yahoo.path.Path;
 import com.yahoo.text.Utf8;
+import com.yahoo.vespa.config.server.filedistribution.FileDistributionProvider;
+import com.yahoo.vespa.config.server.filedistribution.MockFileDistributionProvider;
 import com.yahoo.vespa.config.server.tenant.TenantRepository;
 import com.yahoo.vespa.config.util.ConfigUtils;
 import com.yahoo.vespa.curator.Curator;
@@ -17,6 +19,7 @@ import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
 
+import java.io.File;
 import java.time.Instant;
 import java.util.List;
 import java.util.Optional;
@@ -35,6 +38,7 @@ public class SessionZooKeeperClientTest {
     private static final TenantName tenantName = TenantName.defaultName();
 
     private Curator curator;
+    private FileDistributionProvider fileDistributionProvider;
 
     @Rule
     public ExpectedException expectedException = ExpectedException.none();
@@ -43,6 +47,7 @@ public class SessionZooKeeperClientTest {
     public void setup() {
         curator = new MockCurator();
         curator.create(sessionsPath());
+        fileDistributionProvider = new MockFileDistributionProvider(new File("appdir"), new File("refdir"));
     }
 
     @Test
@@ -171,6 +176,7 @@ public class SessionZooKeeperClientTest {
         SessionZooKeeperClient zkc = new SessionZooKeeperClient(curator,
                                                                 tenantName,
                                                                 sessionId,
+                                                                fileDistributionProvider,
                                                                 ConfigUtils.getCanonicalHostName());
         zkc.createNewSession(Instant.now());
         return zkc;
diff --git a/configserver/src/test/java/com/yahoo/vespa/config/server/zookeeper/ZKApplicationPackageTest.java b/configserver/src/test/java/com/yahoo/vespa/config/server/zookeeper/ZKApplicationPackageTest.java
index 4b4b4b92627..fc365016450 100644
--- a/configserver/src/test/java/com/yahoo/vespa/config/server/zookeeper/ZKApplicationPackageTest.java
+++ b/configserver/src/test/java/com/yahoo/vespa/config/server/zookeeper/ZKApplicationPackageTest.java
@@ -14,6 +14,8 @@ import com.yahoo.config.provisioning.FlavorsConfig;
 import com.yahoo.io.IOUtils;
 import com.yahoo.path.Path;
 import com.yahoo.text.Utf8;
+import com.yahoo.vespa.config.server.filedistribution.MockFileDistributionFactory;
+import com.yahoo.vespa.config.server.filedistribution.MockFileDistributionProvider;
 import com.yahoo.vespa.curator.Curator;
 import com.yahoo.vespa.curator.mock.MockCurator;
 import org.junit.Before;
@@ -75,7 +77,8 @@ public class ZKApplicationPackageTest {
     @Test
     public void testBasicZKFeed() throws IOException {
         feed(curator, new File(APP));
-        ZKApplicationPackage zkApp = new ZKApplicationPackage(curator, Path.fromString("/0"));
+        MockFileDistributionProvider fileDistributionProvider = new MockFileDistributionProvider(new File(APP), new File("references"));
+        ZKApplicationPackage zkApp = new ZKApplicationPackage(fileDistributionProvider, curator, Path.fromString("/0"));
         assertTrue(Pattern.compile(".*<slobroks>.*",Pattern.MULTILINE+Pattern.DOTALL).matcher(IOUtils.readAll(zkApp.getServices())).matches());
         assertTrue(Pattern.compile(".*<alias>.*",Pattern.MULTILINE+Pattern.DOTALL).matcher(IOUtils.readAll(zkApp.getHosts())).matches());
         assertTrue(Pattern.compile(".*<slobroks>.*",Pattern.MULTILINE+Pattern.DOTALL).matcher(IOUtils.readAll(zkApp.getFile(Path.fromString("services.xml")).createReader())).matches());
@@ -96,7 +99,6 @@ public class ZKApplicationPackageTest {
         Version goodVersion = new Version(3, 0, 0);
         assertTrue(zkApp.getFileRegistries().containsKey(goodVersion));
         assertFalse(zkApp.getFileRegistries().containsKey(new Version(0, 0, 0)));
-        assertThat(zkApp.getFileRegistries().get(goodVersion).fileSourceHost(), is("dummyfiles"));
         AllocatedHosts readInfo = zkApp.getAllocatedHosts().get();
         assertEquals(Utf8.toString(toJson(ALLOCATED_HOSTS)), Utf8.toString(toJson(readInfo)));
         assertEquals(TEST_FLAVOR.get().resources(), readInfo.getHosts().iterator().next().advertisedResources());