Dispatch get and visitor messages to correct distributor stripe.

3 files changed, 21 insertions, 0 deletions

diff --git a/storage/src/vespa/storage/distributor/distributor.cpp b/storage/src/vespa/storage/distributor/distributor.cpp
index 368a25315fb..65945b2c6ae 100644
--- a/storage/src/vespa/storage/distributor/distributor.cpp
+++ b/storage/src/vespa/storage/distributor/distributor.cpp
@@ -24,6 +24,7 @@
 #include <vespa/storage/config/distributorconfiguration.h>
 #include <vespa/storage/distributor/maintenance/simplebucketprioritydatabase.h>
 #include <vespa/storageapi/message/persistence.h>
+#include <vespa/storageapi/message/visitor.h>
 #include <vespa/storageframework/generic/status/xmlstatusreporter.h>
 #include <vespa/vdslib/distribution/distribution.h>
 #include <vespa/vespalib/util/memoryusage.h>
@@ -350,6 +351,12 @@ get_bucket_id_for_striping(const api::StorageMessage& msg, const DistributorNode
                     return reply.getBucketId();
                 }
             }
+            case api::MessageType::GET_ID:
+                return node_ctx.bucket_id_factory().getBucketId(dynamic_cast<const api::GetCommand&>(msg).getDocumentId());
+            case api::MessageType::VISITOR_CREATE_ID:
+                return dynamic_cast<const api::CreateVisitorCommand&>(msg).super_bucket_id();
+            case api::MessageType::VISITOR_CREATE_REPLY_ID:
+                return dynamic_cast<const api::CreateVisitorReply&>(msg).super_bucket_id();
             default:
                 return msg.getBucketId();
         }
diff --git a/storageapi/src/vespa/storageapi/message/visitor.cpp b/storageapi/src/vespa/storageapi/message/visitor.cpp
index d87f65a72cf..ec7fb1dc2d4 100644
--- a/storageapi/src/vespa/storageapi/message/visitor.cpp
+++ b/storageapi/src/vespa/storageapi/message/visitor.cpp
@@ -74,6 +74,16 @@ CreateVisitorCommand::getBucket() const
     return document::Bucket(_bucketSpace, document::BucketId());
 }
 
+document::BucketId
+CreateVisitorCommand::super_bucket_id() const
+{
+    if (_buckets.empty()) {
+        // TODO STRIPE: Is this actually an error situation? Should be fixed elsewhere.
+        return document::BucketId();
+    }
+    return _buckets[0];
+}
+
 void
 CreateVisitorCommand::print(std::ostream& out, bool verbose,
                             const std::string& indent) const
@@ -120,6 +130,7 @@ CreateVisitorCommand::print(std::ostream& out, bool verbose,
 
 CreateVisitorReply::CreateVisitorReply(const CreateVisitorCommand& cmd)
     : StorageReply(cmd),
+      _super_bucket_id(cmd.super_bucket_id()),
       _lastBucket(document::BucketId(INT_MAX))
 {
 }
diff --git a/storageapi/src/vespa/storageapi/message/visitor.h b/storageapi/src/vespa/storageapi/message/visitor.h
index 8440591ecde..1313b275e95 100644
--- a/storageapi/src/vespa/storageapi/message/visitor.h
+++ b/storageapi/src/vespa/storageapi/message/visitor.h
@@ -79,6 +79,7 @@ public:
     uint32_t getVisitorCmdId() const { return _visitorCmdId; }
     document::BucketSpace getBucketSpace() const { return _bucketSpace; }
     document::Bucket getBucket() const override;
+    document::BucketId super_bucket_id() const;
     const vespalib::string & getLibraryName() const { return _libName; }
     const vespalib::string & getInstanceId() const { return _instanceId; }
     const vespalib::string & getControlDestination() const { return _controlDestination; }
@@ -114,6 +115,7 @@ public:
  */
 class CreateVisitorReply : public StorageReply {
 private:
+    document::BucketId _super_bucket_id;
     document::BucketId _lastBucket;
     vdslib::VisitorStatistics _visitorStatistics;
 
@@ -124,6 +126,7 @@ public:
 
     void setLastBucket(const document::BucketId& lastBucket) { _lastBucket = lastBucket; }
 
+    const document::BucketId& super_bucket_id() const { return _super_bucket_id; }
     const document::BucketId& getLastBucket() const { return _lastBucket; }
 
     void setVisitorStatistics(const vdslib::VisitorStatistics& stats) { _visitorStatistics = stats; }