Merge pull request #29518 from vespa-engine/balder/ensure-distributionkey-is-always-set

Avoid timeout during grouping leaving distributionKey unset. Populate…

4 files changed, 15 insertions, 26 deletions

diff --git a/searchcore/src/tests/grouping/grouping_test.cpp b/searchcore/src/tests/grouping/grouping_test.cpp
index fb46fb65f3b..951b0bf2857 100644
--- a/searchcore/src/tests/grouping/grouping_test.cpp
+++ b/searchcore/src/tests/grouping/grouping_test.cpp
@@ -301,7 +301,7 @@ TEST_F("testGroupingSession", DoomFixture()) {
     RankedHit hit;
     hit._docId = 0;
     GroupingManager &manager(session.getGroupingManager());
-    manager.groupInRelevanceOrder(&hit, 1);
+    manager.groupInRelevanceOrder(7, &hit, 1);
     CheckAttributeReferences attrCheck_after;
     GroupingList &gl3(initContext.getGroupingList());
     for (auto & grouping : gl3) {
@@ -360,7 +360,7 @@ TEST_F("testEmptySessionId", DoomFixture()) {
     RankedHit hit;
     hit._docId = 0;
     GroupingManager &manager(session.getGroupingManager());
-    manager.groupInRelevanceOrder(&hit, 1);
+    manager.groupInRelevanceOrder(8, &hit, 1);
     EXPECT_EQUAL(id, session.getSessionId());
     ASSERT_TRUE(!session.getGroupingManager().empty());
     ASSERT_TRUE(session.finished() && session.getSessionId().empty());
@@ -427,7 +427,7 @@ void doGrouping(GroupingContext &ctx,
     hits.emplace_back(doc1, rank1);
     hits.emplace_back(doc2, rank2);
     hits.emplace_back(doc3, rank3);
-    man.groupInRelevanceOrder(&hits[0], 3);
+    man.groupInRelevanceOrder(9, &hits[0], 3);
 }
 
 TEST_F("test grouping fork/join", DoomFixture()) {
diff --git a/searchcore/src/vespa/searchcore/grouping/groupingmanager.cpp b/searchcore/src/vespa/searchcore/grouping/groupingmanager.cpp
index 799d2663589..9ece37665cd 100644
--- a/searchcore/src/vespa/searchcore/grouping/groupingmanager.cpp
+++ b/searchcore/src/vespa/searchcore/grouping/groupingmanager.cpp
@@ -19,13 +19,6 @@ using vespalib::make_string_short::fmt;
 
 //-----------------------------------------------------------------------------
 
-GroupingManager::GroupingManager(GroupingContext & groupingContext)
-    : _groupingContext(groupingContext)
-{
-}
-
-GroupingManager::~GroupingManager() = default;
-
 using search::expression::ExpressionNode;
 using search::expression::AttributeNode;
 using search::expression::ConfigureStaticParams;
@@ -66,15 +59,17 @@ GroupingManager::init(const IAttributeContext &attrCtx)
 }
 
 void
-GroupingManager::groupInRelevanceOrder(const RankedHit *searchResults, uint32_t binSize)
+GroupingManager::groupInRelevanceOrder(uint32_t distributionKey, const RankedHit *searchResults, uint32_t binSize)
 {
     _groupingContext.groupInRelevanceOrder(searchResults, binSize);
+    _groupingContext.setDistributionKey(distributionKey);
 }
 
 void
-GroupingManager::groupUnordered(const RankedHit *searchResults, uint32_t binSize, const search::BitVector * overflow)
+GroupingManager::groupUnordered(uint32_t distributionKey, const RankedHit *searchResults, uint32_t binSize, const search::BitVector * overflow)
 {
     _groupingContext.groupUnordered(searchResults, binSize, overflow);
+    _groupingContext.setDistributionKey(distributionKey);
 }
 
 void
diff --git a/searchcore/src/vespa/searchcore/grouping/groupingmanager.h b/searchcore/src/vespa/searchcore/grouping/groupingmanager.h
index 48b84e8e4e4..b0856ae7189 100644
--- a/searchcore/src/vespa/searchcore/grouping/groupingmanager.h
+++ b/searchcore/src/vespa/searchcore/grouping/groupingmanager.h
@@ -20,7 +20,7 @@ class GroupingContext;
 class GroupingManager
 {
 private:
-    GroupingContext         &_groupingContext;
+    GroupingContext   &_groupingContext;
 public:
     GroupingManager(const GroupingManager &) = delete;
     GroupingManager &operator=(const GroupingManager &) = delete;
@@ -29,12 +29,9 @@ public:
      *
      * @param groupingContext Context to use for grouping
      **/
-    GroupingManager(GroupingContext & groupingContext);
-
-    /**
-     * Release resources
-     **/
-    ~GroupingManager();
+    GroupingManager(GroupingContext & groupingContext) noexcept
+        : _groupingContext(groupingContext)
+    {}
 
     /**
      * @return true if this manager is holding an empty grouping request.
@@ -56,7 +53,7 @@ public:
      * @param searchResults the result set in array form
      * @param binSize size of search result array
      **/
-    void groupInRelevanceOrder(const RankedHit *searchResults, uint32_t binSize);
+    void groupInRelevanceOrder(uint32_t distributionKey, const RankedHit *searchResults, uint32_t binSize);
 
     /**
      * Perform actual grouping on the given the results.
@@ -67,7 +64,7 @@ public:
      * @param binSize size of search result array
      * @param overflow The unranked hits.
      **/
-    void groupUnordered(const RankedHit *searchResults, uint32_t binSize, const BitVector * overflow);
+    void groupUnordered(uint32_t distributionKey, const RankedHit *searchResults, uint32_t binSize, const BitVector * overflow);
 
     /**
      * Merge another grouping context into the underlying context of
diff --git a/searchcore/src/vespa/searchcore/proton/matching/match_thread.cpp b/searchcore/src/vespa/searchcore/proton/matching/match_thread.cpp
index afe46d0fd36..211e67f1e2b 100644
--- a/searchcore/src/vespa/searchcore/proton/matching/match_thread.cpp
+++ b/searchcore/src/vespa/searchcore/proton/matching/match_thread.cpp
@@ -402,7 +402,7 @@ MatchThread::processResult(const Doom & doom, search::ResultSet::UP result, Resu
     if (doom.hard_doom()) return;
     if (hasGrouping) {
         search::grouping::GroupingManager man(*context.grouping);
-        man.groupUnordered(hits, numHits, bits);
+        man.groupUnordered(_distributionKey, hits, numHits, bits);
     }
     if (doom.hard_doom()) return;
     size_t sortLimit = hasGrouping ? numHits : context.result->maxSize();
@@ -410,7 +410,7 @@ MatchThread::processResult(const Doom & doom, search::ResultSet::UP result, Resu
     if (doom.hard_doom()) return;
     if (hasGrouping) {
         search::grouping::GroupingManager man(*context.grouping);
-        man.groupInRelevanceOrder(hits, numHits);
+        man.groupInRelevanceOrder(_distributionKey, hits, numHits);
     }
     if (doom.hard_doom()) return;
     fillPartialResult(context, totalHits, numHits, hits, bits);
@@ -421,9 +421,6 @@ MatchThread::processResult(const Doom & doom, search::ResultSet::UP result, Resu
     if (auto task = matchToolsFactory.createOnFirstPhaseTask()) {
         task->run(search::ResultSet::stealResult(std::move(*result)));
     }
-    if (hasGrouping) {
-        context.grouping->setDistributionKey(_distributionKey);
-    }
 }
 
 //-----------------------------------------------------------------------------