[core] Add a util for size literals (ray-project#48638)

As titled, I think having `MB` explicitly called out is more readable than `1024 * 1024` or `1<<20` Intended use case: ray-project#48262 (comment) Signed-off-by: dentiny <dentinyhao@gmail.com>
ShuN6211 · Nov 12, 2024 · 4f6a419 · 4f6a419
1 parent 6e5584e
commit 4f6a419
Show file tree

Hide file tree

Showing 6 changed files with 160 additions and 19 deletions.
diff --git a/BUILD.bazel b/BUILD.bazel
@@ -620,6 +620,7 @@ ray_cc_library(
     deps = [
         ":reporter_rpc",
         ":stats_metric",
+        "//src/ray/util:size_literals",
         "@com_github_grpc_grpc//:grpc_opencensus_plugin",
     ],
 )

diff --git a/src/ray/stats/metric_defs.cc b/src/ray/stats/metric_defs.cc
@@ -14,9 +14,11 @@
 
 #include "ray/stats/metric_defs.h"
 
-namespace ray {
+#include "ray/util/size_literals.h"
 
-namespace stats {
+using namespace ray::literals;
+
+namespace ray::stats {
 
 /// The definitions of metrics that you can use everywhere.
 ///
@@ -114,23 +116,19 @@ DEFINE_stats(
     (),
     ray::stats::GAUGE);
 
-double operator""_MB(unsigned long long int x) {
-  return static_cast<double>(1024L * 1024L * x);
-}
-
 DEFINE_stats(object_store_dist,
              "The distribution of object size in bytes",
              ("Source"),
-             ({32_MB,
-               64_MB,
-               128_MB,
-               256_MB,
-               512_MB,
-               1024_MB,
-               2048_MB,
-               4096_MB,
-               8192_MB,
-               16384_MB}),
+             ({32_MiB,
+               64_MiB,
+               128_MiB,
+               256_MiB,
+               512_MiB,
+               1024_MiB,
+               2048_MiB,
+               4096_MiB,
+               8192_MiB,
+               16384_MiB}),
              ray::stats::HISTOGRAM);
 
 /// Placement group metrics from the GCS.
@@ -387,6 +385,4 @@ DEFINE_stats(
     (),
     ray::stats::GAUGE);
 
-}  // namespace stats
-
-}  // namespace ray
+}  // namespace ray::stats
diff --git a/src/ray/util/BUILD b/src/ray/util/BUILD
@@ -43,6 +43,12 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "size_literals",
+    hdrs = ["size_literals.h"],
+    visibility = ["//visibility:public"],
+)
+
 cc_library(
     name = "thread_checker",
     hdrs = ["thread_checker.h"],

diff --git a/src/ray/util/size_literals.h b/src/ray/util/size_literals.h
@@ -0,0 +1,95 @@
+// Copyright 2024 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// These user-defined literals makes sizes.
+
+#pragma once
+
+#include <cstdint>
+
+namespace ray::literals {
+
+// Size literal for integer values.
+constexpr unsigned long long operator""_B(unsigned long long sz) { return sz; }
+
+constexpr unsigned long long operator""_KiB(unsigned long long sz) { return sz * 1024_B; }
+constexpr unsigned long long operator""_KB(unsigned long long sz) { return sz * 1000_B; }
+
+constexpr unsigned long long operator""_MiB(unsigned long long sz) {
+  return sz * 1024_KiB;
+}
+constexpr unsigned long long operator""_MB(unsigned long long sz) { return sz * 1000_KB; }
+
+constexpr unsigned long long operator""_GiB(unsigned long long sz) {
+  return sz * 1024_MiB;
+}
+constexpr unsigned long long operator""_GB(unsigned long long sz) { return sz * 1000_MB; }
+
+constexpr unsigned long long operator""_TiB(unsigned long long sz) {
+  return sz * 1024_GiB;
+}
+constexpr unsigned long long operator""_TB(unsigned long long sz) { return sz * 1000_GB; }
+
+constexpr unsigned long long operator""_PiB(unsigned long long sz) {
+  return sz * 1024_TiB;
+}
+constexpr unsigned long long operator""_PB(unsigned long long sz) { return sz * 1000_TB; }
+
+// Size literals for floating point values.
+constexpr unsigned long long operator""_KiB(long double sz) {
+  const long double res = sz * 1_KiB;
+  return static_cast<unsigned long long>(res);
+}
+constexpr unsigned long long operator""_KB(long double sz) {
+  const long double res = sz * 1_KB;
+  return static_cast<unsigned long long>(res);
+}
+
+constexpr unsigned long long operator""_MiB(long double sz) {
+  const long double res = sz * 1_MiB;
+  return static_cast<unsigned long long>(res);
+}
+constexpr unsigned long long operator""_MB(long double sz) {
+  const long double res = sz * 1_MB;
+  return static_cast<unsigned long long>(res);
+}
+
+constexpr unsigned long long operator""_GiB(long double sz) {
+  const long double res = sz * 1_GiB;
+  return static_cast<unsigned long long>(res);
+}
+constexpr unsigned long long operator""_GB(long double sz) {
+  const long double res = sz * 1_GB;
+  return static_cast<unsigned long long>(res);
+}
+
+constexpr unsigned long long operator""_TiB(long double sz) {
+  const long double res = sz * 1_TiB;
+  return static_cast<unsigned long long>(res);
+}
+constexpr unsigned long long operator""_TB(long double sz) {
+  const long double res = sz * 1_TB;
+  return static_cast<unsigned long long>(res);
+}
+
+constexpr unsigned long long operator""_PiB(long double sz) {
+  const long double res = sz * 1_PiB;
+  return static_cast<unsigned long long>(res);
+}
+constexpr unsigned long long operator""_PB(long double sz) {
+  const long double res = sz * 1_PB;
+  return static_cast<unsigned long long>(res);
+}
+
+}  // namespace ray::literals
diff --git a/src/ray/util/tests/BUILD b/src/ray/util/tests/BUILD
@@ -182,3 +182,15 @@ cc_test(
         "//src/ray/protobuf:gcs_cc_proto",
     ],
 )
+
+cc_test(
+    name = "size_literals_test",
+    srcs = ["size_literals_test.cc"],
+    size = "small",
+    copts = COPTS,
+    tags = ["team:core"],
+    deps = [
+        "//src/ray/util:size_literals",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
diff --git a/src/ray/util/tests/size_literals_test.cc b/src/ray/util/tests/size_literals_test.cc
@@ -0,0 +1,31 @@
+// Copyright 2024 The Ray Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "src/ray/util/size_literals.h"
+
+#include <gtest/gtest.h>
+
+namespace ray::literals {
+
+namespace {
+
+TEST(SizeLiteralsTest, BasicTest) {
+  static_assert(2_MiB == 2 * 1024 * 1024);
+  static_assert(2.5_KB == 2500);
+  static_assert(4_GB == 4'000'000'000);
+}
+
+}  // namespace
+
+}  // namespace ray::literals