Add JustKnobs '64_bit_ids' to gate new db format

Summary:
# What
Put 64-bit db generation behind a JK for deployment safety.

# Stack
This is diff (3/4)

Reviewed By: iamirzhan

Differential Revision: D88260777

fbshipit-source-id: b23f8ffc67f590f208fd27b8f0e357864c9c88d2

Author: simonhollis

glean/rocksdb/database-impl.cpp

@@ -11,6 +11,10 @@
 
 #include "glean/rts/timer.h"
 
+#ifndef OSS
+#include "justknobs/JustKnobProxy.h"
+#endif
+
 namespace facebook {
 namespace glean {
 namespace rocks {
@@ -147,16 +151,36 @@ DatabaseImpl::DatabaseImpl(
   }
 
   // Initialize ownership format version
-  // For new DBs: write the current (64-bit) format version
-  // For existing DBs: if no version marker exists, assume 32-bit format
-  ownership_format_version = static_cast<uint32_t>(initAdminValue(
-      container_,
-      AdminId::OWNERSHIP_FORMAT_VERSION,
-      static_cast<uint64_t>(
-          container_.mode == Mode::Create ? OWNERSHIP_FORMAT_VERSION_CURRENT
-                                          : OWNERSHIP_FORMAT_VERSION_32BIT),
-      container_.mode == Mode::Create,
-      [] {}));
+  // For new DBs: Check JustKnob to decide whether to use 64-bit or 32-bit
+  // format. For existing DBs: if no version marker exists, assume 32-bit format
+  // Both formats are always readable regardless of gatekeeper value
+  if (container_.mode == Mode::Create) {
+    // New DB: check JustKnob to determine format
+#ifndef OSS
+    static facebook::jk::BooleanKnob use64BitOwnership(
+        "glean/ownership:64_bit_ids");
+    uint32_t format_to_use = use64BitOwnership()
+        ? OWNERSHIP_FORMAT_VERSION_64BIT
+        : OWNERSHIP_FORMAT_VERSION_32BIT;
+#else
+    // OSS builds always use 32-bit format for now
+    uint32_t format_to_use = OWNERSHIP_FORMAT_VERSION_32BIT;
+#endif
+    ownership_format_version = static_cast<uint32_t>(initAdminValue(
+        container_,
+        AdminId::OWNERSHIP_FORMAT_VERSION,
+        static_cast<uint64_t>(format_to_use),
+        true, // write
+        [] {}));
+  } else {
+    // Existing DB: read format version, default to 32-bit if not present
+    ownership_format_version = static_cast<uint32_t>(initAdminValue(
+        container_,
+        AdminId::OWNERSHIP_FORMAT_VERSION,
+        static_cast<uint64_t>(OWNERSHIP_FORMAT_VERSION_32BIT),
+        false, // don't write
+        [] {}));
+  }
 
   VLOG(1) << folly::sformat(
       "ownership_format_version: {} ({})",

glean/rocksdb/ownership.cpp

@@ -10,6 +10,7 @@
 #include <cpp/memory.h> // @manual
 #else
 #include "common/hs/util/cpp/memory.h"
+#include "justknobs/JustKnobProxy.h"
 #endif
 #include <glean/rts/ownership/uset.h>
 #include "glean/rocksdb/database-impl.h"
@@ -73,6 +74,17 @@ DatabaseImpl::loadOwnershipDerivedCounters() {
 
 namespace {
 
+// Helper to serialize UnitId/UsetId based on format version
+rocksdb::Slice
+toSliceForFormat(binary::Output& out, uint64_t id, uint32_t format_version) {
+  if (format_version == OWNERSHIP_FORMAT_VERSION_32BIT) {
+    out.fixed(static_cast<uint32_t>(id));
+  } else {
+    out.fixed(id);
+  }
+  return slice(out);
+}
+
 void putOwnerSet(
     ContainerImpl& container,
     rocksdb::WriteBatch& batch,
@@ -272,10 +284,12 @@ void DatabaseImpl::addOwnership(const std::vector<OwnershipSet>& ownership) {
       touched.push_back(unit_id);
     } else {
       unit_id = first_unit_id + ownership_unit_counters.size() + new_count;
+      // Write UnitId in format based on ownership_format_version
+      binary::Output unitIdOut;
       check(batch.Put(
           container_.family(Family::ownershipUnits),
           slice(set.unit),
-          toSlice(unit_id)));
+          toSliceForFormat(unitIdOut, unit_id, ownership_format_version)));
       EncodedNat key(unit_id);
       check(batch.Put(
           container_.family(Family::ownershipUnitIds),
@@ -907,6 +921,17 @@ std::optional<UsetId> DatabaseImpl::FactOwnerCache::getOwner(
 }
 
 void DatabaseImpl::FactOwnerCache::prepare(ContainerImpl& container) {
+  // Static function - need to check gatekeeper for format
+#ifndef OSS
+  static facebook::jk::BooleanKnob use64BitOwnership(
+      "glean/ownership:64_bit_ids");
+  uint32_t format_version = use64BitOwnership()
+      ? OWNERSHIP_FORMAT_VERSION_64BIT
+      : OWNERSHIP_FORMAT_VERSION_32BIT;
+#else
+  uint32_t format_version = OWNERSHIP_FORMAT_VERSION_32BIT;
+#endif
+
   auto t = makeAutoTimer("prepareFactOwnerCache");
 
   std::unique_ptr<rocksdb::Iterator> iter(container.db->NewIterator(
@@ -934,7 +959,18 @@ void DatabaseImpl::FactOwnerCache::prepare(ContainerImpl& container) {
       index.push_back(HAS_PAGE);
       binary::Output out;
       out.bytes(ids.data(), ids.size() * sizeof(uint16_t));
-      out.bytes(sets.data(), sets.size() * sizeof(UsetId));
+      // Write UsetIds in format based on format_version
+      if (format_version == OWNERSHIP_FORMAT_VERSION_32BIT) {
+        // Write 32-bit UsetIds
+        std::vector<uint32_t> sets32(sets.size());
+        for (size_t i = 0; i < sets.size(); i++) {
+          sets32[i] = static_cast<uint32_t>(sets[i]);
+        }
+        out.bytes(sets32.data(), sets32.size() * sizeof(uint32_t));
+      } else {
+        // Write 64-bit UsetIds
+        out.bytes(sets.data(), sets.size() * sizeof(UsetId));
+      }
       batch.Put(
           container.family(Family::factOwnerPages),
           toSlice(prefix),
@@ -980,19 +1016,39 @@ void DatabaseImpl::FactOwnerCache::prepare(ContainerImpl& container) {
   // write the last page
   writePage();
 
-  batch.Put(
-      container.family(Family::factOwnerPages),
-      INDEX_KEY,
-      slice(
-          folly::ByteRange(
-              reinterpret_cast<const uint8_t*>(index.data()),
-              index.size() * sizeof(UsetId))));
+  // Write index in format based on format_version
+  if (format_version == OWNERSHIP_FORMAT_VERSION_32BIT) {
+    // Write 32-bit index entries
+    std::vector<uint32_t> index32(index.size());
+    if (!index.empty()) {
+      for (size_t i = 0; i < index.size(); i++) {
+        index32[i] = static_cast<uint32_t>(index[i]);
+      }
+    }
+    batch.Put(
+        container.family(Family::factOwnerPages),
+        INDEX_KEY,
+        slice(
+            folly::ByteRange(
+                reinterpret_cast<const uint8_t*>(index32.data()),
+                index32.size() * sizeof(uint32_t))));
+  } else {
+    // Write 64-bit index entries
+    batch.Put(
+        container.family(Family::factOwnerPages),
+        INDEX_KEY,
+        slice(
+            folly::ByteRange(
+                reinterpret_cast<const uint8_t*>(index.data()),
+                index.size() * sizeof(UsetId))));
+  }
 
   t.logFormat(
-      "{} index entries, {} populated, {} orphans",
+      "{} index entries, {} populated, {} orphans (format: {})",
       index.size(),
       populated,
-      orphaned);
+      orphaned,
+      format_version == OWNERSHIP_FORMAT_VERSION_32BIT ? "32-bit" : "64-bit");
 
   // record the number of orphaned facts, this will be fetched by ownershipStats
   batch.Put(
@@ -1110,16 +1166,38 @@ void DatabaseImpl::addDefineOwnership(DefineOwnership& def) {
         pred.new_ids_.size() *
             sizeof(std::remove_reference<
                    decltype(pred.new_ids_)>::type::value_type));
-    val.bytes(
-        pred.owners_.data(),
-        pred.owners_.size() *
-            sizeof(std::remove_reference<
-                   decltype(pred.owners_)>::type::value_type));
-    val.bytes(
-        pred.new_owners_.data(),
-        pred.new_owners_.size() *
-            sizeof(std::remove_reference<
-                   decltype(pred.new_owners_)>::type::value_type));
+
+    // Write owners in format based on ownership_format_version
+    if (ownership_format_version == OWNERSHIP_FORMAT_VERSION_32BIT) {
+      // 32-bit UsetId format
+      std::vector<uint32_t> owners32(pred.owners_.size());
+      if (!pred.owners_.empty()) {
+        for (size_t i = 0; i < pred.owners_.size(); i++) {
+          owners32[i] = static_cast<uint32_t>(pred.owners_[i]);
+        }
+      }
+      val.bytes(owners32.data(), owners32.size() * sizeof(uint32_t));
+
+      std::vector<uint32_t> new_owners32(pred.new_owners_.size());
+      if (!pred.new_owners_.empty()) {
+        for (size_t i = 0; i < pred.new_owners_.size(); i++) {
+          new_owners32[i] = static_cast<uint32_t>(pred.new_owners_[i]);
+        }
+      }
+      val.bytes(new_owners32.data(), new_owners32.size() * sizeof(uint32_t));
+    } else {
+      // 64-bit UsetId format
+      val.bytes(
+          pred.owners_.data(),
+          pred.owners_.size() *
+              sizeof(std::remove_reference<
+                     decltype(pred.owners_)>::type::value_type));
+      val.bytes(
+          pred.new_owners_.data(),
+          pred.new_owners_.size() *
+              sizeof(std::remove_reference<
+                     decltype(pred.new_owners_)>::type::value_type));
+    }
 
     check(batch.Put(
         container_.family(Family::ownershipDerivedRaw),