Check the upper bound of sets before we store them

Summary:
The crashes we've seen all involve a set element being out of range,
and the current check catches this when the sets are read back from
the DB. Let's check the upper bound of sets before we store them, so
we can catch any bugs closer to the source.

Reviewed By: phlalx

Differential Revision: D66704648

fbshipit-source-id: 25936f649f9d1994655ef1e7442e2e00a65c88ef

glean/rocksdb/ownership.cpp

@@ -386,7 +386,9 @@ void DatabaseImpl::storeOwnership(ComputedOwnership& ownership) {
     auto t = makeAutoTimer("storeOwnership(sets)");
     rocksdb::WriteBatch batch;
 
-    auto serialized = ownership.sets_.toEliasFano();
+    auto upper = ownership.sets_.getFirstId() + ownership.sets_.size();
+
+    auto serialized = ownership.sets_.toEliasFano(upper);
     uint32_t id = ownership.sets_.getFirstId();
     CHECK_GE(id, next_uset_id);
 
@@ -399,7 +401,7 @@ void DatabaseImpl::storeOwnership(ComputedOwnership& ownership) {
       id++;
     }
 
-    next_uset_id = ownership.sets_.getFirstId() + ownership.sets_.size();
+    next_uset_id = upper;
     check(batch.Put(
         container_.family(Family::admin),
         toSlice(AdminId::NEXT_UNIT_ID),
@@ -949,7 +951,7 @@ void DatabaseImpl::addDefineOwnership(DefineOwnership& def) {
         usets_->promote(p);
         assert(next_uset_id == p->id);
         next_uset_id++;
-        auto ownerset = p->toEliasFano();
+        auto ownerset = p->toEliasFano(next_uset_id);
         putOwnerSet(container_, batch, p->id, ownerset.op, ownerset.set);
         ownerset.set.free();
         numNewSets++;

glean/rts/ownership/setu32.cpp

@@ -440,11 +440,12 @@ SetU32::merge(SetU32& result, const SetU32& left, const SetU32& right) {
   }
 }
 
-SetU32::MutableEliasFanoList SetU32::toEliasFano() {
-  auto upperBound = this->upper();
+SetU32::MutableEliasFanoList SetU32::toEliasFano(uint32_t upperBound) const {
+  auto max = this->upper();
+  // upperBound is a check only, to catch garbage before we store it
+  CHECK_LT(max, upperBound);
   size_t size = this->size();
-  folly::compression::EliasFanoEncoder<uint32_t, uint32_t> encoder(
-      size, upperBound);
+  folly::compression::EliasFanoEncoder<uint32_t, uint32_t> encoder(size, max);
 
   VLOG(5) << "upper=" << upperBound << ", size=" << size;
   foreach([&](uint32_t elt) { encoder.add(elt); });

glean/rts/ownership/setu32.h

@@ -525,7 +525,7 @@ class SetU32 {
   using MutableEliasFanoList =
       folly::compression::MutableEliasFanoCompressedList;
   using EliasFanoList = folly::compression::EliasFanoCompressedList;
-  MutableEliasFanoList toEliasFano();
+  MutableEliasFanoList toEliasFano(uint32_t upper) const;
   static SetU32 fromEliasFano(const EliasFanoList& list);
 
   static void dump(SetU32&);

glean/rts/ownership/uset.cpp

@@ -12,16 +12,17 @@ namespace facebook {
 namespace glean {
 namespace rts {
 
-SetExpr<Usets::MutableEliasFanoList> Uset::toEliasFano() {
-  return {exp.op, exp.set.toEliasFano()};
+SetExpr<Usets::MutableEliasFanoList> Uset::toEliasFano(UsetId max) const {
+  return {exp.op, exp.set.toEliasFano(max)};
 }
 
-std::vector<SetExpr<Usets::MutableEliasFanoList>> Usets::toEliasFano() {
+std::vector<SetExpr<Usets::MutableEliasFanoList>> Usets::toEliasFano(
+    UsetId max) {
   std::vector<SetExpr<Usets::MutableEliasFanoList>> sets(stats.promoted);
   for (auto uset : usets) {
     if (uset->promoted()) {
       VLOG(5) << "exporting: " << uset->id;
-      sets[uset->id - firstId] = uset->toEliasFano();
+      sets[uset->id - firstId] = uset->toEliasFano(max);
     }
   }
   return sets;

glean/rts/ownership/uset.h

@@ -117,7 +117,7 @@ struct Uset {
   };
 
   using MutableEliasFanoList = SetU32::MutableEliasFanoList;
-  SetExpr<MutableEliasFanoList> toEliasFano();
+  SetExpr<MutableEliasFanoList> toEliasFano(UsetId max) const;
 };
 
 /**
@@ -280,7 +280,7 @@ struct Usets {
   }
 
   using MutableEliasFanoList = SetU32::MutableEliasFanoList;
-  std::vector<SetExpr<MutableEliasFanoList>> toEliasFano();
+  std::vector<SetExpr<MutableEliasFanoList>> toEliasFano(UsetId max);
 
  private:
   folly::F14FastSet<Uset*, Uset::Hash, Uset::Eq> usets;

glean/rts/tests/OwnershipTest.cpp

@@ -287,7 +287,7 @@ TEST(OwnershipTest, SliceTest) {
 
   auto usets = buildExampleSets(units);
   auto firstUsetId = usets.getFirstId();
-  auto sets = usets.toEliasFano();
+  auto sets = usets.toEliasFano(usets.getFirstId() + usets.size());
   uint32_t numSets = usets.statistics().promoted;
 
   // One fact with each different set
@@ -314,7 +314,7 @@ struct SetSerializationTest : testing::Test {};
 RC_GTEST_PROP(SetSerializationTest, testSerialization, ()) {
   const auto set = *rc::gen::nonEmpty(rc::gen::arbitrary<std::set<uint32_t>>());
   SetU32 a = SetU32::from(set);
-  MutableOwnerSet b = a.toEliasFano();
+  MutableOwnerSet b = a.toEliasFano(a.upper() + 1);
   binary::Output o;
   serializeEliasFano(o, b);
   auto size = o.size();