1 files changed, 7 insertions, 1 deletions
diff --git a/container-search/src/main/java/com/yahoo/search/dispatch/CloseableChannel.java b/container-search/src/main/java/com/yahoo/search/dispatch/CloseableChannel.java
index 838afa0c7fc..643b8f81318 100644
--- a/container-search/src/main/java/com/yahoo/search/dispatch/CloseableChannel.java
+++ b/container-search/src/main/java/com/yahoo/search/dispatch/CloseableChannel.java
@@ -17,9 +17,15 @@ import java.util.Optional;
  */
 public class CloseableChannel implements Closeable {
     private FS4Channel channel;
+    private final Optional<Integer> distributionKey;
 
     public CloseableChannel(Backend backend) {
+        this(backend, Optional.empty());
+    }
+
+    public CloseableChannel(Backend backend, Optional<Integer> distributionKey) {
         this.channel = backend.openChannel();
+        this.distributionKey = distributionKey;
     }
 
     public void setQuery(Query query) {
@@ -35,7 +41,7 @@ public class CloseableChannel implements Closeable {
     }
 
     public Optional<Integer> distributionKey() {
-        return channel.distributionKey();
+        return distributionKey;
     }
 
     @Override