WIP for adding bucket space to stat-related messages in Java

5 files changed, 81 insertions, 6 deletions

diff --git a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/DocumentProtocol.java b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/DocumentProtocol.java
index 7cebdbb9abf..d8ee2a39ecd 100755
--- a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/DocumentProtocol.java
+++ b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/DocumentProtocol.java
@@ -357,8 +357,9 @@ public class DocumentProtocol implements Protocol {
         putRoutableFactory(MESSAGE_REMOVEDOCUMENT, new RoutableFactories52.RemoveDocumentMessageFactory(), from52);
 
         // 6.x serialization
-        // TODO stat, listbuckets
         putRoutableFactory(MESSAGE_CREATEVISITOR, new RoutableFactories60.CreateVisitorMessageFactory(), from6);
+        putRoutableFactory(MESSAGE_STATBUCKET, new RoutableFactories60.StatBucketMessageFactory(), from6);
+        putRoutableFactory(MESSAGE_GETBUCKETLIST, new RoutableFactories60.GetBucketListMessageFactory(), from6);
     }
 
     /**
diff --git a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/GetBucketListMessage.java b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/GetBucketListMessage.java
index 74629018e55..5e5d79f86b7 100755
--- a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/GetBucketListMessage.java
+++ b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/GetBucketListMessage.java
@@ -6,13 +6,19 @@ import com.yahoo.document.BucketId;
 public class GetBucketListMessage extends DocumentMessage {
 
     private BucketId bucketId;
+    private String bucketSpace = "default";
 
     GetBucketListMessage() {
         // must be deserialized into
     }
 
     public GetBucketListMessage(BucketId bucketId) {
+        this(bucketId, "default");
+    }
+
+    public GetBucketListMessage(BucketId bucketId, String bucketSpace) {
         this.bucketId = bucketId;
+        this.bucketSpace = bucketSpace;
     }
 
     public BucketId getBucketId() {
@@ -23,6 +29,14 @@ public class GetBucketListMessage extends DocumentMessage {
         bucketId = id;
     }
 
+    public String getBucketSpace() {
+        return bucketSpace;
+    }
+
+    public void setBucketSpace(String bucketSpace) {
+        this.bucketSpace = bucketSpace;
+    }
+
     @Override
     public DocumentReply createReply() {
         return new StatBucketReply();
@@ -30,7 +44,7 @@ public class GetBucketListMessage extends DocumentMessage {
 
     @Override
     public int getApproxSize() {
-        return super.getApproxSize() + 8;
+        return super.getApproxSize() + 8 + bucketSpace.length();
     }
 
     @Override
diff --git a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/RoutableFactories50.java b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/RoutableFactories50.java
index 7417e63480e..0ea82ba0d83 100755
--- a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/RoutableFactories50.java
+++ b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/RoutableFactories50.java
@@ -392,18 +392,28 @@ public abstract class RoutableFactories50 {
 
     public static class GetBucketListMessageFactory extends DocumentMessageFactory {
 
+        protected String decodeBucketSpace(Deserializer deserializer) {
+            return "default"; // TODO fixed space repo
+        }
+
         @Override
         protected DocumentMessage doDecode(DocumentDeserializer buf) {
             GetBucketListMessage msg = new GetBucketListMessage();
             msg.setBucketId(new BucketId(buf.getLong(null)));
+            msg.setBucketSpace(decodeBucketSpace(buf));
             return msg;
         }
 
+        protected boolean encodeBucketSpace(String bucketSpace, DocumentSerializer buf) {
+            // TODO fixed space repo
+            return "default".equals(bucketSpace);
+        }
+
         @Override
         protected boolean doEncode(DocumentMessage obj, DocumentSerializer buf) {
             GetBucketListMessage msg = (GetBucketListMessage)obj;
             buf.putLong(null, msg.getBucketId().getRawId());
-            return true;
+            return encodeBucketSpace(msg.getBucketSpace(), buf);
         }
     }
 
@@ -831,20 +841,30 @@ public abstract class RoutableFactories50 {
 
     public static class StatBucketMessageFactory extends DocumentMessageFactory {
 
+        protected String decodeBucketSpace(Deserializer deserializer) {
+            return "default"; // TODO fixed space repo
+        }
+
         @Override
         protected DocumentMessage doDecode(DocumentDeserializer buf) {
             StatBucketMessage msg = new StatBucketMessage();
             msg.setBucketId(new BucketId(buf.getLong(null)));
             msg.setDocumentSelection(decodeString(buf));
+            msg.setBucketSpace(decodeBucketSpace(buf));
             return msg;
         }
 
+        protected boolean encodeBucketSpace(String bucketSpace, DocumentSerializer buf) {
+            // TODO fixed space repo
+            return "default".equals(bucketSpace);
+        }
+
         @Override
         protected boolean doEncode(DocumentMessage obj, DocumentSerializer buf) {
             StatBucketMessage msg = (StatBucketMessage)obj;
             buf.putLong(null, msg.getBucketId().getRawId());
             encodeString(msg.getDocumentSelection(), buf);
-            return true;
+            return encodeBucketSpace(msg.getBucketSpace(), buf);
         }
     }
 
diff --git a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/RoutableFactories60.java b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/RoutableFactories60.java
index ef63fe8b295..723760df117 100644
--- a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/RoutableFactories60.java
+++ b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/RoutableFactories60.java
@@ -19,4 +19,30 @@ public class RoutableFactories60 extends RoutableFactories52 {
         }
     }
 
+    public static class StatBucketMessageFactory extends RoutableFactories52.StatBucketMessageFactory {
+        @Override
+        protected String decodeBucketSpace(Deserializer deserializer) {
+            return decodeString(deserializer);
+        }
+
+        @Override
+        protected boolean encodeBucketSpace(String bucketSpace, DocumentSerializer buf) {
+            encodeString(bucketSpace, buf);
+            return true;
+        }
+    }
+
+    public static class GetBucketListMessageFactory extends RoutableFactories52.GetBucketListMessageFactory {
+        @Override
+        protected String decodeBucketSpace(Deserializer deserializer) {
+            return decodeString(deserializer);
+        }
+
+        @Override
+        protected boolean encodeBucketSpace(String bucketSpace, DocumentSerializer buf) {
+            encodeString(bucketSpace, buf);
+            return true;
+        }
+    }
+
 }
diff --git a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/StatBucketMessage.java b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/StatBucketMessage.java
index 9b2176d85f6..8fb4ceb6d41 100755
--- a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/StatBucketMessage.java
+++ b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/StatBucketMessage.java
@@ -6,6 +6,7 @@ import com.yahoo.document.BucketId;
 public class StatBucketMessage extends DocumentMessage {
 
     private BucketId bucketId;
+    private String bucketSpace = "default";
     private String documentSelection;
 
     StatBucketMessage() {
@@ -13,7 +14,12 @@ public class StatBucketMessage extends DocumentMessage {
     }
 
     public StatBucketMessage(BucketId bucket, String documentSelection) {
-        this.bucketId = bucket;
+        this(bucket, "default", documentSelection);
+    }
+
+    public StatBucketMessage(BucketId bucketId, String bucketSpace, String documentSelection) {
+        this.bucketId = bucketId;
+        this.bucketSpace = bucketSpace;
         this.documentSelection = documentSelection;
     }
 
@@ -33,6 +39,14 @@ public class StatBucketMessage extends DocumentMessage {
         this.documentSelection = documentSelection;
     }
 
+    public String getBucketSpace() {
+        return bucketSpace;
+    }
+
+    public void setBucketSpace(String bucketSpace) {
+        this.bucketSpace = bucketSpace;
+    }
+
     @Override
     public DocumentReply createReply() {
         return new StatBucketReply();
@@ -40,7 +54,7 @@ public class StatBucketMessage extends DocumentMessage {
 
     @Override
     public int getApproxSize() {
-        return super.getApproxSize() + 8 + documentSelection.length();
+        return super.getApproxSize() + 8 + bucketSpace.length() + documentSelection.length();
     }
 
     @Override