Narrow down routing of Document API GetDocument messages.

4 files changed, 79 insertions, 0 deletions

diff --git a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/DocumentRouteSelectorPolicy.java b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/DocumentRouteSelectorPolicy.java
index dbf4a6cc593..c3204c7b5e2 100755
--- a/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/DocumentRouteSelectorPolicy.java
+++ b/documentapi/src/main/java/com/yahoo/documentapi/messagebus/protocol/DocumentRouteSelectorPolicy.java
@@ -3,6 +3,7 @@ package com.yahoo.documentapi.messagebus.protocol;
 
 import com.yahoo.config.subscription.ConfigSubscriber;
 import com.yahoo.document.Document;
+import com.yahoo.document.DocumentGet;
 import com.yahoo.document.DocumentPut;
 import com.yahoo.document.DocumentType;
 import com.yahoo.document.select.DocumentSelector;
@@ -153,6 +154,15 @@ public class DocumentRouteSelectorPolicy
                 return true;
             }
         }
+        case DocumentProtocol.MESSAGE_GETDOCUMENT: {
+            GetDocumentMessage getMsg = (GetDocumentMessage)msg;
+            if (getMsg.getDocumentId().hasDocType()) {
+                DocumentGet getOp = new DocumentGet(getMsg.getDocumentId());
+                return selector.accepts(getOp) != Result.FALSE;
+            } else {
+                return true;
+            }
+        }
 
         default:
             return true;
diff --git a/documentapi/src/test/java/com/yahoo/documentapi/messagebus/protocol/test/PolicyTestCase.java b/documentapi/src/test/java/com/yahoo/documentapi/messagebus/protocol/test/PolicyTestCase.java
index 1192f8bc7ad..e2f1c9cd937 100755
--- a/documentapi/src/test/java/com/yahoo/documentapi/messagebus/protocol/test/PolicyTestCase.java
+++ b/documentapi/src/test/java/com/yahoo/documentapi/messagebus/protocol/test/PolicyTestCase.java
@@ -378,6 +378,28 @@ public class PolicyTestCase {
         frame.assertSelect(Arrays.asList("testdoc-route", "other-route"));
     }
 
+    @Test
+    public void get_document_messages_are_sent_to_the_route_handling_the_given_document_type() {
+        PolicyTestFrame frame = createFrameWithTwoRoutes();
+
+        frame.setMessage(createGet("id:ns:testdoc::1"));
+        frame.assertSelect(Arrays.asList("testdoc-route"));
+
+        frame.setMessage(createGet("id:ns:other::1"));
+        frame.assertSelect(Arrays.asList("other-route"));
+    }
+
+    @Test
+    public void get_document_messages_with_legacy_document_ids_are_sent_to_all_routes() {
+        PolicyTestFrame frame = createFrameWithTwoRoutes();
+
+        frame.setMessage(createGet("userdoc:testdoc:1234:1"));
+        frame.assertSelect(Arrays.asList("testdoc-route", "other-route"));
+
+        frame.setMessage(createGet("userdoc:other:1234:1"));
+        frame.assertSelect(Arrays.asList("testdoc-route", "other-route"));
+    }
+
     private PolicyTestFrame createFrameWithTwoRoutes() {
         PolicyTestFrame result = new PolicyTestFrame(manager);
         result.setHop(new HopSpec("test", createDocumentRouteSelectorConfigWithTwoRoutes())
@@ -400,6 +422,10 @@ public class PolicyTestCase {
         return new RemoveDocumentMessage(new DocumentId(docId));
     }
 
+    private GetDocumentMessage createGet(String docId) {
+        return new GetDocumentMessage(new DocumentId(docId));
+    }
+
     @Test
     public void testSubsetService() {
         PolicyTestFrame frame = new PolicyTestFrame("docproc/cluster.default", manager);
diff --git a/documentapi/src/tests/policies/policies_test.cpp b/documentapi/src/tests/policies/policies_test.cpp
index c01cdfde30c..7c93e186338 100644
--- a/documentapi/src/tests/policies/policies_test.cpp
+++ b/documentapi/src/tests/policies/policies_test.cpp
@@ -66,6 +66,8 @@ public:
     void testDocumentRouteSelectorIgnore();
     void remove_document_messages_are_sent_to_the_route_handling_the_given_document_type();
     void remove_document_messages_with_legacy_document_ids_are_sent_to_all_routes();
+    void get_document_messages_are_sent_to_the_route_handling_the_given_document_type();
+    void get_document_messages_with_legacy_document_ids_are_sent_to_all_routes();
     void testExternSend();
     void testExternMultipleSlobroks();
     void testLoadBalancer();
@@ -108,6 +110,8 @@ Test::Main() {
     testDocumentRouteSelectorIgnore();  TEST_FLUSH();
     remove_document_messages_are_sent_to_the_route_handling_the_given_document_type(); TEST_FLUSH();
     remove_document_messages_with_legacy_document_ids_are_sent_to_all_routes();        TEST_FLUSH();
+    get_document_messages_are_sent_to_the_route_handling_the_given_document_type(); TEST_FLUSH();
+    get_document_messages_with_legacy_document_ids_are_sent_to_all_routes();        TEST_FLUSH();
     testExternSend();                   TEST_FLUSH();
     testExternMultipleSlobroks();       TEST_FLUSH();
     testLoadBalancer();                 TEST_FLUSH();
@@ -685,6 +689,12 @@ makeRemove(vespalib::string docId)
     return std::make_unique<RemoveDocumentMessage>(DocumentId(docId));
 }
 
+std::unique_ptr<GetDocumentMessage>
+makeGet(vespalib::string docId)
+{
+    return std::make_unique<GetDocumentMessage>(DocumentId(docId));
+}
+
 }
 
 void
@@ -711,6 +721,30 @@ Test::remove_document_messages_with_legacy_document_ids_are_sent_to_all_routes()
     EXPECT_TRUE(frame->testSelect({"testdoc-route", "other-route"}));
 }
 
+void
+Test::get_document_messages_are_sent_to_the_route_handling_the_given_document_type()
+{
+    auto frame = createFrameWithTwoRoutes(_repo);
+
+    frame->setMessage(makeGet("id:ns:testdoc::1"));
+    EXPECT_TRUE(frame->testSelect({"testdoc-route"}));
+
+    frame->setMessage(makeGet("id:ns:other::1"));
+    EXPECT_TRUE(frame->testSelect({"other-route"}));
+}
+
+void
+Test::get_document_messages_with_legacy_document_ids_are_sent_to_all_routes()
+{
+    auto frame = createFrameWithTwoRoutes(_repo);
+
+    frame->setMessage(makeGet("userdoc:testdoc:1234:1"));
+    EXPECT_TRUE(frame->testSelect({"testdoc-route", "other-route"}));
+
+    frame->setMessage(makeGet("userdoc:other:1234:1"));
+    EXPECT_TRUE(frame->testSelect({"testdoc-route", "other-route"}));
+}
+
 namespace {
     string getDefaultDistributionConfig(
                     uint16_t redundancy = 2, uint16_t nodeCount = 10,
diff --git a/documentapi/src/vespa/documentapi/messagebus/policies/documentrouteselectorpolicy.cpp b/documentapi/src/vespa/documentapi/messagebus/policies/documentrouteselectorpolicy.cpp
index 62135150d98..378f12b5138 100644
--- a/documentapi/src/vespa/documentapi/messagebus/policies/documentrouteselectorpolicy.cpp
+++ b/documentapi/src/vespa/documentapi/messagebus/policies/documentrouteselectorpolicy.cpp
@@ -4,6 +4,7 @@
 #include <vespa/document/bucket/bucketidfactory.h>
 #include <vespa/document/select/parser.h>
 #include <vespa/documentapi/messagebus/documentprotocol.h>
+#include <vespa/documentapi/messagebus/messages/getdocumentmessage.h>
 #include <vespa/documentapi/messagebus/messages/putdocumentmessage.h>
 #include <vespa/documentapi/messagebus/messages/updatedocumentmessage.h>
 #include <vespa/documentapi/messagebus/messages/documentignoredreply.h>
@@ -137,6 +138,14 @@ DocumentRouteSelectorPolicy::select(mbus::RoutingContext &context, const vespali
             return true;
         }
     }
+    case DocumentProtocol::MESSAGE_GETDOCUMENT: {
+        const GetDocumentMessage &getMsg = static_cast<const GetDocumentMessage &>(msg);
+        if (getMsg.getDocumentId().hasDocType()) {
+            return it->second->contains(getMsg.getDocumentId()) != Result::False;
+        } else {
+            return true;
+        }
+    }
     default:
         return true;
     }