Set distribution key whe dispatching to content node

4 files changed, 24 insertions, 22 deletions

diff --git a/container-search/src/main/java/com/yahoo/fs4/DocumentInfo.java b/container-search/src/main/java/com/yahoo/fs4/DocumentInfo.java
index 075ec60c099..e5ab00fb139 100644
--- a/container-search/src/main/java/com/yahoo/fs4/DocumentInfo.java
+++ b/container-search/src/main/java/com/yahoo/fs4/DocumentInfo.java
@@ -28,9 +28,9 @@ public class DocumentInfo implements Cloneable {
     }
 
     public DocumentInfo(GlobalId globalId, int metric, int partId, int distributionKey) {
-        this.globalId=globalId;
-        this.metric=metric;
-        this.partId=partId;
+        this.globalId = globalId;
+        this.metric = metric;
+        this.partId = partId;
         this.distributionKey = distributionKey;
     }
 
diff --git a/container-search/src/main/java/com/yahoo/fs4/QueryResultPacket.java b/container-search/src/main/java/com/yahoo/fs4/QueryResultPacket.java
index b18d1e88a4b..239101184ea 100644
--- a/container-search/src/main/java/com/yahoo/fs4/QueryResultPacket.java
+++ b/container-search/src/main/java/com/yahoo/fs4/QueryResultPacket.java
@@ -125,7 +125,7 @@ public class QueryResultPacket extends Packet {
         soonActiveDocs = buffer.getLong();
         degradedReason = buffer.getInt();
 
-        decodeDocuments(buffer,documentCount);
+        decodeDocuments(buffer, documentCount);
         if (propsFeature) {
             int numMaps = buffer.getInt();
             propsArray = new FS4Properties[numMaps];
diff --git a/container-search/src/main/java/com/yahoo/fs4/mplex/Backend.java b/container-search/src/main/java/com/yahoo/fs4/mplex/Backend.java
index 81192be035a..68a13c06e32 100644
--- a/container-search/src/main/java/com/yahoo/fs4/mplex/Backend.java
+++ b/container-search/src/main/java/com/yahoo/fs4/mplex/Backend.java
@@ -61,7 +61,7 @@ public class Backend implements ConnectionFactory {
     private final ConnectionPool connectionPool;
     private final PacketDumper packetDumper;
     private final AtomicInteger connectionCount = new AtomicInteger(0);
-
+    private final Optional<Integer> distributionKey;
 
     /**
      * For unit testing.  do not use
@@ -74,9 +74,15 @@ public class Backend implements ConnectionFactory {
         packetDumper = null;
         address = null;
         connectionPool = new ConnectionPool();
+        distributionKey = Optional.empty();
     }
 
-    public Backend(String host, int port, String serverDiscriminator, ListenerPool listenerPool, ConnectionPool connectionPool) {
+    public Backend(String host,
+                   int port,
+                   String serverDiscriminator,
+                   ListenerPool listenerPool,
+                   ConnectionPool connectionPool,
+                   Optional<Integer> distributionKey) {
         String fileNamePattern = "qrs." + serverDiscriminator + '.' + host + ":" + port + ".%s" + ".dump";
         packetDumper = new PacketDumper(new File(Defaults.getDefaults().underVespaHome("logs/vespa/qrs/")),
                                         fileNamePattern);
@@ -86,6 +92,7 @@ public class Backend implements ConnectionFactory {
         this.port = port;
         address = new InetSocketAddress(host, port);
         this.connectionPool = connectionPool;
+        this.distributionKey = distributionKey;
     }
 
     private void logWarning(String attemptDescription, Exception e) {
@@ -96,11 +103,13 @@ public class Backend implements ConnectionFactory {
         log.log(Level.INFO, "Exception on " + attemptDescription + " '" + host + ":" + port + "': " + Exceptions.toMessageString(e));
     }
 
+    /** Returns the distribution key of the content node this represents, or empty if it is a dispatch node */
+    public Optional<Integer> distributionKey() { return distributionKey; }
+
     // ============================================================
     // ==== connection pool stuff
     // ============================================================
 
-
     /**
      * Fetch a connection from the connection pool.  If the pool
      * is empty we create a connection.
@@ -187,10 +196,7 @@ public class Backend implements ConnectionFactory {
     //==== channel management
     //============================================================
 
-    /**
-     * Open a new channel to fdispatch.  Analogous to the "Channel"
-     * concept as used in FS4.
-     */
+    /** Opens a new channel to fdispatch.  Analogous to the "Channel" concept as used in FS4. */
     public FS4Channel openChannel () {
         int cachedChannelId;
         synchronized (this) {
diff --git a/container-search/src/main/java/com/yahoo/fs4/mplex/FS4Channel.java b/container-search/src/main/java/com/yahoo/fs4/mplex/FS4Channel.java
index b0fce14e73f..237b0cdb8e2 100644
--- a/container-search/src/main/java/com/yahoo/fs4/mplex/FS4Channel.java
+++ b/container-search/src/main/java/com/yahoo/fs4/mplex/FS4Channel.java
@@ -1,10 +1,10 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-
 package com.yahoo.fs4.mplex;
 
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
@@ -16,23 +16,18 @@ import com.yahoo.fs4.ChannelTimeoutException;
 import com.yahoo.fs4.Packet;
 import com.yahoo.search.Query;
 
-
-
 /**
- *
  * This class is used to represent a "channel" in the FS4 protocol.
  * A channel represents a session between a client and the fdispatch.
  * Internally this class has a  response queue used by the backend
  * for queueing up FS4 packets that belong to this channel (or
  * <em>session</em>, which might be a more appropriate name for it).
- *
- * <P>
  * Outbound packets are handed off to the FS4Connection.
  *
  * @author Bjorn Borud
  */
-public class FS4Channel
-{
+public class FS4Channel {
+
     private static Logger log = Logger.getLogger(FS4Channel.class.getName());
 
     private Integer channelId;
@@ -67,13 +62,14 @@ public class FS4Channel
         return query;
     }
 
-    /**
-     * @return returns an Integer representing the (fs4) channel id
-     */
+    /** Returns the (fs4) channel id */
     public Integer getChannelId () {
         return channelId;
     }
 
+    /** Returns the distribution key of the content node this represents, or empty if it is a dispatch node */
+    public Optional<Integer> distributionKey() { return backend == null ? Optional.empty() : backend.distributionKey(); }
+
     /**
      * Closes the channel
      */