Extend use of bucket ownership hash.

7 files changed, 86 insertions, 96 deletions

diff --git a/storage/src/tests/distributor/bucketdbupdatertest.cpp b/storage/src/tests/distributor/bucketdbupdatertest.cpp
index 0a3a001222f..22d9c945262 100644
--- a/storage/src/tests/distributor/bucketdbupdatertest.cpp
+++ b/storage/src/tests/distributor/bucketdbupdatertest.cpp
@@ -1778,7 +1778,7 @@ TEST_F(BucketDBUpdaterTest, no_db_resurrection_for_bucket_not_owned_in_current_s
         uint32_t expectedMsgs = messageCount(2), dummyBucketsToReturn = 1;
         ASSERT_NO_FATAL_FAILURE(setAndEnableClusterState(stateAfter, expectedMsgs, dummyBucketsToReturn));
     }
-    EXPECT_FALSE(getDistributorBucketSpace().owns_bucket_in_current_state(bucket));
+    EXPECT_FALSE(getDistributorBucketSpace().get_bucket_ownership_flags(bucket).owned_in_current_state());
 
     ASSERT_NO_FATAL_FAILURE(sendFakeReplyForSingleBucketRequest(*rbi));
 
@@ -1804,8 +1804,8 @@ TEST_F(BucketDBUpdaterTest, no_db_resurrection_for_bucket_not_owned_in_pending_s
     lib::ClusterState stateAfter("distributor:3 storage:3");
     // Set, but _don't_ enable cluster state. We want it to be pending.
     setSystemState(stateAfter);
-    EXPECT_TRUE(getDistributorBucketSpace().owns_bucket_in_current_state(bucket));
-    EXPECT_FALSE(getDistributorBucketSpace().check_ownership_in_pending_state(bucket).isOwned());
+    EXPECT_TRUE(getDistributorBucketSpace().get_bucket_ownership_flags(bucket).owned_in_current_state());
+    EXPECT_FALSE(getDistributorBucketSpace().get_bucket_ownership_flags(bucket).owned_in_pending_state());
 
     ASSERT_NO_FATAL_FAILURE(sendFakeReplyForSingleBucketRequest(*rbi));
 
@@ -1925,7 +1925,7 @@ TEST_F(BucketDBUpdaterTest, changed_distribution_config_does_not_elide_bucket_db
     setDistribution(getDistConfig6Nodes2Groups());
 
     getBucketDatabase().forEach(*func_processor([&](const auto& e) {
-        EXPECT_TRUE(getDistributorBucketSpace().check_ownership_in_pending_state(e.getBucketId()).isOwned());
+        EXPECT_TRUE(getDistributorBucketSpace().get_bucket_ownership_flags(e.getBucketId()).owned_in_pending_state());
     }));
 }
 
diff --git a/storage/src/tests/distributor/distributor_bucket_space_test.cpp b/storage/src/tests/distributor/distributor_bucket_space_test.cpp
index 3880dd3d768..e206308618f 100644
--- a/storage/src/tests/distributor/distributor_bucket_space_test.cpp
+++ b/storage/src/tests/distributor/distributor_bucket_space_test.cpp
@@ -55,8 +55,6 @@ DistributorBucketSpaceTest::count_distributor_buckets()
     for (uint32_t i = 0; i < (1u << distribution_bits); ++i) {
         BucketId bucket(distribution_bits, i);
         bool owned = bucket_space.check_ownership_in_pending_and_current_state(bucket).isOwned();
-        bool check_owned = bucket_space.check_ownership_in_pending_and_current_state_fallback(bucket).isOwned();
-        EXPECT_EQ(check_owned, owned);
         if (owned) {
             ++owned_buckets;
         }
@@ -100,8 +98,6 @@ DistributorBucketSpaceTest::count_deep_split_distributor_buckets()
     for (uint32_t i = 0; i < 100; ++i) {
         BucketId bucket(42u, i * (1ul << 32) + bias);
         bool owned = bucket_space.check_ownership_in_pending_and_current_state(bucket).isOwned();
-        bool check_owned = bucket_space.check_ownership_in_pending_and_current_state_fallback(bucket).isOwned();
-        EXPECT_EQ(check_owned, owned);
         if (owned) {
             ++owned_buckets;
         }
diff --git a/storage/src/tests/distributor/distributortest.cpp b/storage/src/tests/distributor/distributortest.cpp
index 25221b9cc7d..72115da8de6 100644
--- a/storage/src/tests/distributor/distributortest.cpp
+++ b/storage/src/tests/distributor/distributortest.cpp
@@ -557,7 +557,7 @@ TEST_F(DistributorTest, no_db_resurrection_for_bucket_not_owned_in_pending_state
     getBucketDBUpdater().onSetSystemState(stateCmd);
 
     document::BucketId nonOwnedBucket(16, 3);
-    EXPECT_FALSE(getDistributorBucketSpace().check_ownership_in_pending_state(nonOwnedBucket).isOwned());
+    EXPECT_FALSE(getDistributorBucketSpace().get_bucket_ownership_flags(nonOwnedBucket).owned_in_pending_state());
     EXPECT_FALSE(getDistributorBucketSpace().check_ownership_in_pending_and_current_state(nonOwnedBucket).isOwned());
 
     std::vector<BucketCopy> copies;
diff --git a/storage/src/vespa/storage/distributor/bucket_ownership_flags.h b/storage/src/vespa/storage/distributor/bucket_ownership_flags.h
new file mode 100644
index 00000000000..c9003552f12
--- /dev/null
+++ b/storage/src/vespa/storage/distributor/bucket_ownership_flags.h
@@ -0,0 +1,29 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+#pragma once
+
+namespace storage::distributor {
+
+/*
+ * Compact bucket ownership representation. Default value is not owned
+ * by current state or pending state.  Bucket is always considered
+ * owned in pending state if there is no pending state.
+ */
+class BucketOwnershipFlags {
+    uint8_t _flags;
+    
+    static constexpr uint8_t owned_in_current_state_flag = 0x1;
+    static constexpr uint8_t owned_in_pending_state_flag = 0x2;
+    
+public:
+    BucketOwnershipFlags() noexcept
+        : _flags(0)
+    {
+    }
+
+    bool owned_in_current_state() const noexcept { return ((_flags & owned_in_current_state_flag) != 0); }
+    bool owned_in_pending_state() const noexcept {  return ((_flags & owned_in_pending_state_flag) != 0); }
+    void set_owned_in_current_state() noexcept { _flags |= owned_in_current_state_flag; }
+    void set_owned_in_pending_state() noexcept { _flags |= owned_in_pending_state_flag; }
+};
+
+}
diff --git a/storage/src/vespa/storage/distributor/distributor_bucket_space.cpp b/storage/src/vespa/storage/distributor/distributor_bucket_space.cpp
index 82bbeadb9a5..f9122650311 100644
--- a/storage/src/vespa/storage/distributor/distributor_bucket_space.cpp
+++ b/storage/src/vespa/storage/distributor/distributor_bucket_space.cpp
@@ -118,55 +118,6 @@ DistributorBucketSpace::owns_bucket_in_state(
     return owns_bucket_in_state(*_distribution, clusterState, bucket);
 }
 
-bool
-DistributorBucketSpace::owns_bucket_in_current_state(document::BucketId bucket) const
-{
-    return owns_bucket_in_state(*_distribution, *_clusterState, bucket);
-}
-
-BucketOwnership
-DistributorBucketSpace::check_ownership_in_pending_state(document::BucketId bucket) const
-{
-    if (_pending_cluster_state) {
-        if (!owns_bucket_in_state(*_pending_cluster_state, bucket)) {
-            return BucketOwnership::createNotOwnedInState(*_pending_cluster_state);
-        }
-    }
-    return BucketOwnership::createOwned();
-}
-BucketOwnership
-DistributorBucketSpace::check_ownership_in_pending_and_given_state(
-        const lib::Distribution& distribution,
-        const lib::ClusterState& clusterState,
-        document::BucketId bucket) const
-{
-    try {
-        BucketOwnership pendingRes(
-                check_ownership_in_pending_state(bucket));
-        if (!pendingRes.isOwned()) {
-            return pendingRes;
-        }
-        uint16_t distributor = distribution.getIdealDistributorNode(
-                clusterState, bucket);
-
-        if (_node_index == distributor) {
-            return BucketOwnership::createOwned();
-        } else {
-            return BucketOwnership::createNotOwnedInState(clusterState);
-        }
-    } catch (lib::TooFewBucketBitsInUseException& e) {
-        return BucketOwnership::createNotOwnedInState(clusterState);
-    } catch (lib::NoDistributorsAvailableException& e) {
-        return BucketOwnership::createNotOwnedInState(clusterState);
-    }
-}
-
-BucketOwnership
-DistributorBucketSpace::check_ownership_in_pending_and_current_state_fallback(document::BucketId bucket) const
-{
-    return check_ownership_in_pending_and_given_state(*_distribution, *_clusterState, bucket);
-}
-
 std::vector<uint16_t>
 DistributorBucketSpace::get_ideal_nodes(document::BucketId bucket) const
 {
@@ -185,21 +136,46 @@ DistributorBucketSpace::get_ideal_nodes(document::BucketId bucket) const
     return insres.first->second;
 }
 
-BucketOwnership
-DistributorBucketSpace::check_ownership_in_pending_and_current_state(document::BucketId bucket) const
+BucketOwnershipFlags
+DistributorBucketSpace::get_bucket_ownership_flags(document::BucketId bucket) const
 {
     if (bucket.getUsedBits() < _distribution_bits) {
-        // Cannot map to super bucket ==> cannot cache result
-        return check_ownership_in_pending_and_current_state_fallback(bucket);
+        BucketOwnershipFlags flags;
+        if (!_pending_cluster_state) {
+            flags.set_owned_in_pending_state();
+        }
+        return flags;
     }
     document::BucketId super_bucket(_distribution_bits, bucket.getId());
     auto itr = _ownerships.find(super_bucket);
     if (itr != _ownerships.end()) {
         return itr->second;
     }
-    auto insres = _ownerships.insert(std::make_pair(super_bucket, check_ownership_in_pending_and_current_state_fallback(super_bucket)));
+    BucketOwnershipFlags flags;
+    if (!_pending_cluster_state || owns_bucket_in_state(*_distribution, *_pending_cluster_state, super_bucket)) {
+        flags.set_owned_in_pending_state();
+    }
+    if (owns_bucket_in_state(*_distribution, *_clusterState, super_bucket)) {
+        flags.set_owned_in_current_state();
+    }
+    auto insres = _ownerships.insert(std::make_pair(super_bucket, flags));
     assert(insres.second);
     return insres.first->second;
 }
 
+BucketOwnership
+DistributorBucketSpace::check_ownership_in_pending_and_current_state(document::BucketId bucket) const
+{
+    auto flags = get_bucket_ownership_flags(bucket);
+    if (!flags.owned_in_pending_state()) {
+        assert(_pending_cluster_state);
+        return BucketOwnership::createNotOwnedInState(*_pending_cluster_state);
+    }
+    if (flags.owned_in_current_state()) {
+        return BucketOwnership::createOwned();
+    } else {
+        return BucketOwnership::createNotOwnedInState(*_clusterState);
+    }
+}
+
 }
diff --git a/storage/src/vespa/storage/distributor/distributor_bucket_space.h b/storage/src/vespa/storage/distributor/distributor_bucket_space.h
index 96536b26da3..b3722ed9c91 100644
--- a/storage/src/vespa/storage/distributor/distributor_bucket_space.h
+++ b/storage/src/vespa/storage/distributor/distributor_bucket_space.h
@@ -2,6 +2,7 @@
 #pragma once
 
 #include "bucketownership.h"
+#include "bucket_ownership_flags.h"
 #include <vespa/document/bucket/bucketid.h>
 #include <vespa/vespalib/stllike/hash_map.h>
 #include <memory>
@@ -37,11 +38,12 @@ class DistributorBucketSpace {
     uint16_t                                 _distribution_bits;
     std::shared_ptr<const lib::ClusterState> _pending_cluster_state;
     std::vector<bool>                        _available_nodes;
-    mutable vespalib::hash_map<document::BucketId, BucketOwnership, document::BucketId::hash>       _ownerships;
+    mutable vespalib::hash_map<document::BucketId, BucketOwnershipFlags, document::BucketId::hash>  _ownerships;
     mutable vespalib::hash_map<document::BucketId, std::vector<uint16_t>, document::BucketId::hash> _ideal_nodes;
 
     void clear();
     void enumerate_available_nodes();
+    bool owns_bucket_in_state(const lib::Distribution& distribution, const lib::ClusterState& cluster_state, document::BucketId bucket) const;
 public:
     explicit DistributorBucketSpace();
     explicit DistributorBucketSpace(uint16_t node_index);
@@ -77,42 +79,29 @@ public:
     }
 
     void set_pending_cluster_state(std::shared_ptr<const lib::ClusterState> pending_cluster_state);
+    const lib::ClusterState& get_pending_cluster_state() const noexcept { return *_pending_cluster_state; }
 
     std::vector<uint16_t> get_ideal_nodes_fallback(document::BucketId bucket) const;
 
-    bool owns_bucket_in_state(const lib::Distribution& distribution, const lib::ClusterState& cluster_state, document::BucketId bucket) const;
-
     /**
      * Returns true if this distributor owns the given bucket in the
      * given cluster and current distribution config.
+     * Only used by unit tests.
      */
     bool owns_bucket_in_state(const lib::ClusterState& clusterState, document::BucketId bucket) const;
 
-    /**
-     * Returns true if this distributor owns the given bucket with the current
-     * cluster state and distribution config.
-     */
-    bool owns_bucket_in_current_state(document::BucketId bucket) const;
-
-    /**
-     * If there is a pending state, returns ownership state of bucket in it.
-     * Otherwise always returns "is owned", i.e. it must also be checked in the current state.
-     */
-    BucketOwnership check_ownership_in_pending_state(document::BucketId bucket) const;
-    /**
-     * Returns the ownership status of a bucket as decided with the given
-     * distribution and cluster state -and- that of the pending cluster
-     * state and distribution (if any pending exists).
-     */
-    BucketOwnership check_ownership_in_pending_and_given_state(const lib::Distribution& distribution,
-                                                               const lib::ClusterState& clusterState,
-                                                               document::BucketId bucket) const;
-    BucketOwnership check_ownership_in_pending_and_current_state_fallback(document::BucketId bucket) const;
     const std::vector<bool>& get_available_nodes() const { return _available_nodes; }
     /**
      * Returns the ideal nodes for the given bucket.
      */
     std::vector<uint16_t> get_ideal_nodes(document::BucketId bucket) const;
+
+    /*
+     * Return bucket ownership flags for the given bucket. Bucket is always
+     * considered owned in pending state if there is no pending state.
+     */
+    BucketOwnershipFlags get_bucket_ownership_flags(document::BucketId bucket) const;
+
     /**
      * Returns the ownership status of a bucket as decided with the current
      * distribution and cluster state -and- that of the pending cluster
diff --git a/storage/src/vespa/storage/distributor/externaloperationhandler.cpp b/storage/src/vespa/storage/distributor/externaloperationhandler.cpp
index cf00718b1fd..aebea14b1f1 100644
--- a/storage/src/vespa/storage/distributor/externaloperationhandler.cpp
+++ b/storage/src/vespa/storage/distributor/externaloperationhandler.cpp
@@ -181,7 +181,8 @@ ExternalOperationHandler::checkTimestampMutationPreconditions(api::StorageComman
                                                               PersistenceOperationMetricSet& persistenceMetrics)
 {
     auto &bucket_space(_op_ctx.bucket_space_repo().get(cmd.getBucket().getBucketSpace()));
-    if (!bucket_space.owns_bucket_in_current_state(bucketId)) {
+    auto bucket_ownership_flags = bucket_space.get_bucket_ownership_flags(bucketId);
+    if (!bucket_ownership_flags.owned_in_current_state()) {
         document::Bucket bucket(cmd.getBucket().getBucketSpace(), bucketId);
         LOG(debug, "Distributor manager received %s, bucket %s with wrong distribution",
             cmd.toString().c_str(), bucket.toString().c_str());
@@ -190,12 +191,11 @@ ExternalOperationHandler::checkTimestampMutationPreconditions(api::StorageComman
         return false;
     }
 
-    auto pending = bucket_space.check_ownership_in_pending_state(bucketId);
-    if (!pending.isOwned()) {
+    if (!bucket_ownership_flags.owned_in_pending_state()) {
         // We return BUSY here instead of WrongDistributionReply to avoid clients potentially
         // ping-ponging between cluster state versions during a state transition.
-        auto& current_state = _op_ctx.bucket_space_repo().get(document::FixedBucketSpaces::default_space()).getClusterState();
-        auto& pending_state = pending.getNonOwnedState();
+        auto& current_state = bucket_space.getClusterState();
+        auto& pending_state = bucket_space.get_pending_cluster_state();
         bounce_with_busy_during_state_transition(cmd, current_state, pending_state);
         return false;
     }
@@ -238,7 +238,8 @@ void ExternalOperationHandler::bounce_or_invoke_read_only_op(
         Func func)
 {
     auto &bucket_space(_op_ctx.bucket_space_repo().get(bucket.getBucketSpace()));
-    if (!bucket_space.owns_bucket_in_current_state(bucket.getBucketId())) {
+    auto bucket_ownership_flags = bucket_space.get_bucket_ownership_flags(bucket.getBucketId());
+    if (!bucket_ownership_flags.owned_in_current_state()) {
         LOG(debug, "Distributor manager received %s, bucket %s with wrong distribution",
             cmd.toString().c_str(), bucket.toString().c_str());
         bounce_with_wrong_distribution(cmd);
@@ -246,15 +247,14 @@ void ExternalOperationHandler::bounce_or_invoke_read_only_op(
         return;
     }
 
-    auto pending = bucket_space.check_ownership_in_pending_state(bucket.getBucketId());
-    if (pending.isOwned()) {
+    if (bucket_ownership_flags.owned_in_pending_state()) {
         func(_op_ctx.bucket_space_repo());
     } else {
         if (_op_ctx.distributor_config().allowStaleReadsDuringClusterStateTransitions()) {
             func(_op_ctx.read_only_bucket_space_repo());
         } else {
-            auto& current_state = _op_ctx.bucket_space_repo().get(document::FixedBucketSpaces::default_space()).getClusterState();
-            auto& pending_state = pending.getNonOwnedState();
+            auto& current_state = bucket_space.getClusterState();
+            auto& pending_state = bucket_space.get_pending_cluster_state();
             bounce_with_busy_during_state_transition(cmd, current_state, pending_state);
         }
     }