cleanup(pubsublite): have routing policy return std::uint32_t (#8675)

Author: 18suresha

google/cloud/pubsublite/internal/default_routing_policy.cc

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "google/cloud/pubsublite/internal/default_routing_policy.h"
+#include "google/cloud/internal/sha256_hash.h"
 #include "google/cloud/version.h"
 #include <unordered_map>
 
@@ -21,10 +22,13 @@ namespace cloud {
 namespace pubsublite_internal {
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
 
+// will always be 32 as specified in sha256_hash.h
+std::uint32_t constexpr kNumBytesSha256 = 32;
+
 // Uses the identity that `(a*b) % m == ((a % m) * (b % m)) % m`
-std::uint64_t ModPow(std::uint64_t val, std::uint64_t pow, std::uint32_t mod) {
+std::uint64_t ModPow(std::uint64_t val, std::uint32_t pow, std::uint32_t mod) {
   std::uint64_t result = 1;
-  for (std::uint32_t i = 0; i < pow; ++i) {
+  for (std::uint32_t i = 0; i != pow; ++i) {
     result *= (val % mod);
     result %= mod;
   }
@@ -33,29 +37,29 @@ std::uint64_t ModPow(std::uint64_t val, std::uint64_t pow, std::uint32_t mod) {
 
 // Uses the identity that `(a*b) % m == ((a % m) * (b % m)) % m`
 // Uses the identity that `(a+b) % m == ((a % m) + (b % m)) % m`
-std::uint64_t GetMod(std::array<uint8_t, 32> big_endian, std::uint32_t mod) {
+std::uint32_t GetMod(std::array<uint8_t, kNumBytesSha256> big_endian,
+                     std::uint32_t mod) {
   std::uint64_t result = 0;
-  for (std::uint64_t i = 0; i < big_endian.size(); ++i) {
-    std::uint64_t val_mod = big_endian[i] % mod;
+  for (std::uint32_t i = 0; i != kNumBytesSha256; ++i) {
+    std::uint32_t val_mod = big_endian[i] % mod;
 
-    std::uint64_t pow = big_endian.size() - (i + 1);
-    std::uint64_t pow_mod = ModPow(
-        // 2^8
-        static_cast<std::uint64_t>(1 << 8), pow, mod);
+    std::uint32_t pow = kNumBytesSha256 - (i + 1);
+    std::uint64_t pow_mod = ModPow(256, pow, mod);
 
     result += (val_mod * pow_mod) % mod;
     result %= mod;
   }
-  return result;
+  // within bounds because `mod`ed by a std::uint32_t value
+  return static_cast<std::uint32_t>(result);
 }
 
-std::uint64_t DefaultRoutingPolicy::Route(std::uint32_t num_partitions) {
+Partition DefaultRoutingPolicy::Route(Partition num_partitions) {
   // atomic operation
-  return counter_++ % num_partitions;
+  return static_cast<std::uint32_t>(counter_++ % num_partitions);
 }
 
-std::uint64_t DefaultRoutingPolicy::Route(std::string const& message_key,
-                                          std::uint32_t num_partitions) {
+Partition DefaultRoutingPolicy::Route(std::string const& message_key,
+                                      Partition num_partitions) {
   return GetMod(google::cloud::internal::Sha256Hash(message_key),
                 num_partitions);
 }

google/cloud/pubsublite/internal/default_routing_policy.h

@@ -16,7 +16,6 @@
 #define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_PUBSUBLITE_INTERNAL_DEFAULT_ROUTING_POLICY_H
 
 #include "google/cloud/pubsublite/internal/routing_policy.h"
-#include "google/cloud/internal/sha256_hash.h"
 #include "google/cloud/version.h"
 #include <atomic>
 #include <cstddef>
@@ -33,7 +32,7 @@ GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
 std::uint64_t ModPow(std::uint64_t val, std::uint64_t pow, std::uint32_t mod);
 
 // returns <integer value of `big_endian`> % `mod` while accounting for overflow
-std::uint64_t GetMod(std::array<uint8_t, 32> big_endian, std::uint32_t mod);
+std::uint32_t GetMod(std::array<uint8_t, 32> big_endian, std::uint32_t mod);
 
 /**
  * Implements the same routing policy as all the other Pub/Sub Lite client
@@ -49,9 +48,9 @@ std::uint64_t GetMod(std::array<uint8_t, 32> big_endian, std::uint32_t mod);
  */
 class DefaultRoutingPolicy : public RoutingPolicy {
  public:
-  std::uint64_t Route(std::uint32_t num_partitions) override;
-  std::uint64_t Route(std::string const& message_key,
-                      std::uint32_t num_partitions) override;
+  Partition Route(Partition num_partitions) override;
+  Partition Route(std::string const& message_key,
+                  Partition num_partitions) override;
 
  private:
   std::atomic<std::int64_t> counter_{0};

google/cloud/pubsublite/internal/default_routing_policy_test.cc

@@ -14,6 +14,7 @@
 
 #include "google/cloud/pubsublite/internal/default_routing_policy.h"
 #include <gmock/gmock.h>
+#include <array>
 #include <unordered_map>
 
 namespace google {
@@ -50,9 +51,9 @@ TEST(DefaultRoutingPolicyTest, RouteWithKey) {
 TEST(DefaultRoutingPolicyTest, RouteWithoutKey) {
   unsigned int num_partitions = 29;
   DefaultRoutingPolicy rp;
-  uint64_t initial_partition = rp.Route(num_partitions);
+  std::uint32_t initial_partition = rp.Route(num_partitions);
   for (unsigned int i = 0; i < num_partitions; ++i) {
-    uint64_t next_partition = rp.Route(num_partitions);
+    std::uint32_t next_partition = rp.Route(num_partitions);
     EXPECT_EQ((initial_partition + 1) % num_partitions,
               next_partition % num_partitions);
     initial_partition = next_partition;
@@ -70,6 +71,7 @@ TEST(TestGetMod, MaxValue) {
   EXPECT_EQ(GetMod(arr, 100), 35);
   EXPECT_EQ(GetMod(arr, 10023), 5397);
   EXPECT_EQ(GetMod(arr, UINT8_MAX), 0);
+  EXPECT_EQ(GetMod(arr, UINT32_MAX - 1), 255);
 }
 
 TEST(TestGetMod, OneLessThanMaxValue) {
@@ -82,6 +84,7 @@ TEST(TestGetMod, OneLessThanMaxValue) {
   EXPECT_EQ(GetMod(arr, 100), 34);
   EXPECT_EQ(GetMod(arr, 10023), 5396);
   EXPECT_EQ(GetMod(arr, UINT8_MAX), 254);
+  EXPECT_EQ(GetMod(arr, UINT32_MAX - 1), 254);
 }
 
 TEST(TestGetMod, Zeros) {
@@ -93,6 +96,7 @@ TEST(TestGetMod, Zeros) {
   EXPECT_EQ(GetMod(arr, 100), 0);
   EXPECT_EQ(GetMod(arr, 10023), 0);
   EXPECT_EQ(GetMod(arr, UINT8_MAX), 0);
+  EXPECT_EQ(GetMod(arr, UINT32_MAX - 1), 0);
 }
 
 TEST(TestGetMod, ArbitraryValue) {
@@ -105,6 +109,7 @@ TEST(TestGetMod, ArbitraryValue) {
   EXPECT_EQ(GetMod(arr, 10023), 3346);
   EXPECT_EQ(GetMod(arr, 109000), 60390);
   EXPECT_EQ(GetMod(arr, UINT8_MAX), 100);
+  EXPECT_EQ(GetMod(arr, UINT32_MAX - 1), 1136793478);
 }
 
 TEST(TestGetMod, ArbitraryValue1) {
@@ -115,6 +120,7 @@ TEST(TestGetMod, ArbitraryValue1) {
   EXPECT_EQ(GetMod(arr, 102301), 93535);
   EXPECT_EQ(GetMod(arr, 23), 13);
   EXPECT_EQ(GetMod(arr, UINT8_MAX), 37);
+  EXPECT_EQ(GetMod(arr, UINT32_MAX - 1), 3416191692);
 }
 
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END

google/cloud/pubsublite/internal/routing_policy.h

@@ -22,6 +22,11 @@ namespace cloud {
 namespace pubsublite_internal {
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
 
+/**
+ * A partition should be an integer in the range [0, UINT32_MAX).
+ */
+using Partition = std::uint32_t;
+
 /**
  * Interface for a Pub/Sub Lite routing policy that determines the partition
  * that a message should be sent to, either depending on the message's key or
@@ -30,9 +35,9 @@ GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
 class RoutingPolicy {
  public:
   virtual ~RoutingPolicy() = default;
-  virtual std::uint64_t Route(std::uint32_t num_partitions) = 0;
-  virtual std::uint64_t Route(std::string const& message_key,
-                              std::uint32_t num_partitions) = 0;
+  virtual Partition Route(Partition num_partitions) = 0;
+  virtual Partition Route(std::string const& message_key,
+                          Partition num_partitions) = 0;
 };
 
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END