Refactor Koladata pseudo-random generator to use a weak symbol for reproducibility

This patch replaces the mechanism used to de-randomize Koladata. Instead of relying on an environment variable, it now uses a weak symbol `_KoladataInternalPseudoRandomUint64`.

This allows the behaviour to be controlled by adding a build dependency or by using LD_PRELOAD. The primary motivation is to remove "special" mode logic from client applications. Specifically, the overloaded version can enforce single-thread access to the random number generator without imposing this logic on the default implementation.

PiperOrigin-RevId: 843637091
Change-Id: Ia2b8fb2348b0529138dadb8ce1f3d7cb41138c42

Author: apronchenkov

koladata/BUILD

@@ -88,7 +88,7 @@ cc_library(
         "//koladata/internal:data_bag",
         "//koladata/internal:data_item",
         "//koladata/internal:object_id",
-        "//koladata/internal:random",
+        "//koladata/internal:pseudo_random",
         "//koladata/internal:schema_attrs",
         "//koladata/internal:triples",
         "@com_google_absl//absl/base:core_headers",

koladata/data_bag.h

@@ -38,7 +38,7 @@
 #include "arolla/util/repr.h"
 #include "koladata/internal/data_bag.h"
 #include "koladata/internal/object_id.h"
-#include "koladata/internal/random.h"
+#include "koladata/internal/pseudo_random.h"
 
 namespace koladata {
 
@@ -225,7 +225,7 @@ class DataBag : public arolla::RefcountedBase {
         has_mutable_fallbacks_(false),
         // NOTE: consider lazy initialization of the fingerprint if it becomes
         // expensive to compute.
-        fingerprint_(internal::MaybeDeterministicRandomFingerprint()) {}
+        fingerprint_(internal::PseudoRandomFingerprint()) {}
 
   // Returns a mutable DataBag that wraps provided low-level DataBagImpl.
   static DataBagPtr FromImpl(internal::DataBagImplPtr impl);

koladata/expr/BUILD

@@ -129,7 +129,7 @@ cc_library(
     hdrs = ["non_determinism.h"],
     deps = [
         ":init",
-        "//koladata/internal:random",
+        "//koladata/internal:pseudo_random",
         "@com_google_absl//absl/base:no_destructor",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:string_view",

koladata/expr/non_determinism.cc

@@ -22,7 +22,7 @@
 #include "arolla/expr/expr.h"
 #include "arolla/expr/expr_node.h"
 #include "arolla/expr/registered_expr_operator.h"
-#include "koladata/internal/random.h"
+#include "koladata/internal/pseudo_random.h"
 
 namespace koladata::expr {
 
@@ -35,9 +35,9 @@ absl::StatusOr<arolla::expr::ExprNodePtr> GenNonDeterministicToken() {
   static const absl::NoDestructor op(
       std::make_shared<RegisteredOperator>("koda_internal.non_deterministic"));
   static const absl::NoDestructor leaf(Leaf(kNonDeterministicTokenLeafKey));
-  return MakeOpNode(*op,
-                    {*leaf, Literal(static_cast<int64_t>(
-                                internal::MaybeDeterministicRandomUint64()))});
+  return MakeOpNode(
+      *op,
+      {*leaf, Literal(static_cast<int64_t>(internal::PseudoRandomUint64()))});
 }
 
 }  // namespace koladata::expr

koladata/internal/BUILD

@@ -65,7 +65,7 @@ cc_library(
     srcs = ["object_id.cc"],
     hdrs = ["object_id.h"],
     deps = [
-        ":random",
+        ":pseudo_random",
         "//koladata/internal/op_utils:base62",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_set",
@@ -1089,26 +1089,22 @@ cc_library(
 )
 
 cc_library(
-    name = "random",
-    srcs = ["random.cc"],
-    hdrs = ["random.h"],
+    name = "pseudo_random",
+    srcs = ["pseudo_random.cc"],
+    hdrs = ["pseudo_random.h"],
     deps = [
         "@com_google_absl//absl/base:core_headers",
-        "@com_google_absl//absl/base:no_destructor",
         "@com_google_absl//absl/numeric:int128",
         "@com_google_absl//absl/random",
-        "@com_google_absl//absl/random:distributions",
-        "@com_google_absl//absl/strings:string_view",
-        "@com_google_absl//absl/synchronization",
         "@com_google_arolla//arolla/util",
     ],
 )
 
 cc_test(
-    name = "random_test",
-    srcs = ["random_test.cc"],
+    name = "pseudo_random_test",
+    srcs = ["pseudo_random_test.cc"],
     deps = [
-        ":random",
+        ":pseudo_random",
         "@com_google_googletest//:gtest_main",
     ],
 )

koladata/internal/object_id.cc

@@ -31,15 +31,15 @@
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
 #include "arolla/qtype/simple_qtype.h"
-#include "koladata/internal/random.h"
+#include "koladata/internal/pseudo_random.h"
 
 namespace koladata::internal {
 
 namespace {
 
 uint64_t AllocatorId() {
   static uint64_t kAllocatorId =
-      MaybeDeterministicRandomUint64() & ((1ull << 52) - 1);
+      PseudoRandomUint64() & ((1ull << 52) - 1);
   return kAllocatorId;
 }
 
@@ -105,9 +105,6 @@ AllocationId Allocate(size_t size) {
             id << id_bit_count);
       };
 
-  // TODO: b/464002636 — If we decicde to force single-threaded usage, when
-  // KOLADATA_DETERMINISTIC_SEED variable is set, we have to add an assertion
-  // here as well.
   thread_local std::array<std::pair<uint64_t, uint64_t>, kMaxOffsetBits + 1>
       thread_id_per_offset_;
 

koladata/internal/random.h koladata/internal/pseudo_random.cckoladata/internal/random.h renamed to koladata/internal/pseudo_random.cc

@@ -12,29 +12,28 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 //
-#ifndef KOLADATA_INTERNAL_RANDOM_H_
-#define KOLADATA_INTERNAL_RANDOM_H_
+#include "koladata/internal/pseudo_random.h"
 
 #include <cstdint>
 
+#include "absl/base/attributes.h"
+#include "absl/numeric/int128.h"
+#include "absl/random/random.h"
 #include "arolla/util/fingerprint.h"
 
+// Note: Implement the function as a weak symbol so that this implementation can
+// be overridden by a build dependency or via LD_PRELOAD.
+extern "C" ABSL_ATTRIBUTE_WEAK uint64_t KoladataInternalPseudoRandomUint64() {
+  static thread_local absl::BitGen bitgen;
+  return absl::Uniform<uint64_t>(bitgen);
+}
+
 namespace koladata::internal {
 
-// Returns a random uniformly distributed uint64.
-//
-// If KOLADATA_DETERMINISTIC_SEED environment variable is set (the changes of
-// the value in runtime are ignored), the random number generator will be
-// deterministic.
-uint64_t MaybeDeterministicRandomUint64();
-
-// Returns a random uniformly distributed fingerprint.
-//
-// If KOLADATA_DETERMINISTIC_SEED environment variable is set (the changes of
-// the value in runtime are ignored), the random number generator will be
-// deterministic.
-arolla::Fingerprint MaybeDeterministicRandomFingerprint();
+arolla::Fingerprint PseudoRandomFingerprint() {
+  return arolla::Fingerprint(
+      absl::MakeUint128(KoladataInternalPseudoRandomUint64(),
+                        KoladataInternalPseudoRandomUint64()));
+}
 
 }  // namespace koladata::internal
-
-#endif  // KOLADATA_INTERNAL_RANDOM_H_

koladata/internal/pseudo_random.h

@@ -0,0 +1,49 @@
+// Copyright 2025 Google LLC
+//
+// 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.
+//
+#ifndef KOLADATA_INTERNAL_PSEUDO_RANDOM_H_
+#define KOLADATA_INTERNAL_PSEUDO_RANDOM_H_
+
+#include <cstdint>
+
+#include "arolla/util/fingerprint.h"
+
+// Returns a uniformly distributed pseudo-random uint64.
+//
+// While not cryptographically secure, the default implementation generates
+// unique sequences of random numbers for each process.
+//
+// Note: This function supports overriding via LD_PRELOAD. It is declared as "C"
+// function because providing multiple definitions in C++ constitutes an ODR
+// violation.
+//
+// WARNING: Koladata relies on high-entropy randomness to prevent collisions
+// across processes. If you override this function to return a fixed sequence,
+// you must ensure the process is isolated; it must not interact with any
+// other process using the same sequence.
+extern "C" uint64_t KoladataInternalPseudoRandomUint64();
+
+namespace koladata::internal {
+
+// Returns a pseudo-random uniformly distributed uint64.
+inline uint64_t PseudoRandomUint64() {
+  return KoladataInternalPseudoRandomUint64();
+}
+
+// Returns a pseudo-random uniformly distributed fingerprint.
+arolla::Fingerprint PseudoRandomFingerprint();
+
+}  // namespace koladata::internal
+
+#endif  // KOLADATA_INTERNAL_PSEUDO_RANDOM_H_

koladata/internal/random_test.cc koladata/internal/pseudo_random_test.cckoladata/internal/random_test.cc renamed to koladata/internal/pseudo_random_test.cc

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 //
-#include "koladata/internal/random.h"
+#include "koladata/internal/pseudo_random.h"
 
 #include <cstdint>
 #include <thread>  // NOLINT(build/c++11): only used for tests.
@@ -21,48 +21,17 @@
 #include "gtest/gtest.h"
 
 namespace koladata::internal {
-
-uint64_t DeterministicRandomUint64();
-
 namespace {
 
-TEST(RandomTest, MaybeDeterministicRandomUint64) {
-  EXPECT_NE(MaybeDeterministicRandomUint64(), MaybeDeterministicRandomUint64());
-}
-
-TEST(RandomTest, MaybeDeterministicRandomFingerprint) {
-  EXPECT_NE(MaybeDeterministicRandomFingerprint(),
-            MaybeDeterministicRandomFingerprint());
+TEST(PseudoRandomTest, PseudoRandomUint64) {
+  EXPECT_NE(PseudoRandomUint64(), PseudoRandomUint64());
 }
 
-TEST(RandomTest, MaybeDeterministicRandomUint64_MultiThreaded) {
-  constexpr int kNumThreads = 10;
-  constexpr int kNumValues = 10;
-  std::vector<std::vector<uint64_t>> results(kNumThreads);
-
-  std::vector<std::thread> threads;
-  threads.reserve(kNumThreads);
-  for (int i = 0; i < kNumThreads; ++i) {
-    threads.emplace_back([&results, i] {
-      results[i].reserve(kNumValues);
-      for (int j = 0; j < kNumValues; ++j) {
-        results[i].push_back(MaybeDeterministicRandomUint64());
-      }
-    });
-  }
-
-  for (auto& thread : threads) {
-    thread.join();
-  }
-
-  for (int i = 0; i < kNumThreads; ++i) {
-    for (int j = i + 1; j < kNumThreads; ++j) {
-      EXPECT_NE(results[i], results[j]);
-    }
-  }
+TEST(PseudoRandomTest, PseudoRandomFingerprint) {
+  EXPECT_NE(PseudoRandomFingerprint(), PseudoRandomFingerprint());
 }
 
-TEST(RandomTest, DeterministicRandomUint64_MultiThreaded) {
+TEST(PseudoRandomTest, PseudoRandomUint64_MultiThreaded) {
   constexpr int kNumThreads = 10;
   constexpr int kNumValues = 10;
   std::vector<std::vector<uint64_t>> results(kNumThreads);
@@ -73,7 +42,7 @@ TEST(RandomTest, DeterministicRandomUint64_MultiThreaded) {
     threads.emplace_back([&results, i] {
       results[i].reserve(kNumValues);
       for (int j = 0; j < kNumValues; ++j) {
-        results[i].push_back(DeterministicRandomUint64());
+        results[i].push_back(PseudoRandomUint64());
       }
     });
   }