Fixes for set size limit

Summary: Followup from D68108359. Thanks to Simon for the comments. Reviewed By: simonmar Differential Revision: D68164934 fbshipit-source-id: c05ee855518cf500a4c68121bd26e54c6b2a15e3
facebookincubator · Jan 15, 2025 · de1f90b · de1f90b
1 parent 3aec7fa
commit de1f90b
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 30 deletions.
diff --git a/glean/rts/set.cpp b/glean/rts/set.cpp
@@ -33,11 +33,10 @@ void SetOps::insertOutputSet(SetOps::SetToken token, binary::Output* out) {
     s.insert(out->moveToFbString());
     set_sizes[token] += size;
   } else {
-    throw std::overflow_error(folly::sformat(
-        "Set size limit exceeded for standard set. Set token: {}. Max size: {}. Size: {}",
-        token,
+    rts::error(
+        "Set size limit exceeded for standard set. Max size: {}. Size: {}",
         max_set_size,
-        set_sizes[token] + size));
+        set_sizes[token] + size);
   }
 }
 
@@ -56,19 +55,16 @@ void SetOps::freeSet(SetOps::SetToken token) {
 
 SetOps::SetToken SetOps::newWordSet() {
   wordsets.emplace_back(WordSet());
-  wordset_sizes.emplace_back(0);
   return wordsets.size() - 1;
 }
 
 void SetOps::insertWordSet(SetOps::SetToken token, uint64_t value) {
   auto& set = wordsets[token];
-  if (set.size() + sizeof(uint64_t) <= max_set_size) {
+  if ((set.size() + 1) * sizeof(uint64_t) <= max_set_size) {
     set.insert(value);
-    wordset_sizes[token] += sizeof(uint64_t);
   } else {
     rts::error(
-        "Set size limit exceeded for nat set. Set token: {}. Max size: {}. Size: {}",
-        token,
+        "Set size limit exceeded for nat set. Max size: {}. Size: {}",
         max_set_size,
         set.size() + sizeof(uint64_t));
   }
@@ -79,15 +75,13 @@ void SetOps::insertBytesWordSet(
     const unsigned char* start,
     const unsigned char* end) {
   auto& set = wordsets[token];
-  if (set.size() + (end - start) <= max_set_size) {
+  if (set.size() * sizeof(uint64_t) + (end - start) <= max_set_size) {
     for (const unsigned char* p = start; p < end; p++) {
       set.insert(*p);
     }
-    wordset_sizes[token] += end - start;
   } else {
     rts::error(
-        "Set size limit exceeded for byte set. Set token: {}. Max size: {}. Size: {}",
-        token,
+        "Set size limit exceeded for byte set. Max size: {}. Size: {}",
         max_set_size,
         set.size() + (end - start));
   }
@@ -111,7 +105,6 @@ void SetOps::byteSetToByteArray(SetOps::SetToken token, binary::Output* out) {
 
 void SetOps::freeWordSet(SetOps::SetToken token) {
   wordsets.erase(wordsets.begin() + token);
-  wordset_sizes.erase(wordset_sizes.begin() + token);
 }
 
 } // namespace rts

diff --git a/glean/rts/set.h b/glean/rts/set.h
@@ -61,7 +61,6 @@ struct SetOps {
   std::vector<BytestringSet> sets = {};
   std::vector<WordSet> wordsets = {};
   std::vector<uint64_t> set_sizes = {};
-  std::vector<uint64_t> wordset_sizes = {};
   uint64_t max_set_size = UINT64_MAX;
 };
 

diff --git a/glean/test/tests/Angle/SetTest.hs b/glean/test/tests/Angle/SetTest.hs
@@ -9,6 +9,8 @@
 {-# LANGUAGE QuasiQuotes #-}
 {-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE AllowAmbiguousTypes #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 module Angle.SetTest (main) where
 
 import Control.Exception hiding (assert)
@@ -18,7 +20,7 @@ import Data.Default
 import Data.Text
 
 import Glean.Angle.Parser
-import Glean.Angle.Types hiding (Nat)
+import Glean.Angle.Types hiding (Nat, Type)
 import Glean.Init
 import Glean.Query.Typecheck
 import Glean.Query.Flatten
@@ -27,6 +29,7 @@ import Glean.Schema.Resolve
 import Glean.Database.Schema.Types
 import Glean.Database.Test
 import Glean.Database.Types
+import Glean.Typed.Binary
 import Glean.Types
 
 import qualified Data.HashMap.Strict as HashMap
@@ -132,18 +135,29 @@ setSemanticsTest = TestList
   ]
 
 setLimitTest :: Test
-setLimitTest =
-  TestLabel "Fail when exceeding limit" $
-    dbTestCaseSettings [setMaxSetSize 8] $ \env repo -> do
-      r <- try (runQuery_ env repo $ angleData @(Set Nat) [s| all (1|2) |])
-        :: IO (Either FFIError [Set Nat])
-      case r of
-        Left ffiExc -> do
-          -- Remove the stack trace and just check the actual error
-          let runtimeMsg = Prelude.take 71 (ffiErrorMessage ffiExc)
-          assertEqual "Exception " runtimeMsg errMsg
-        Right _ -> assert False
+setLimitTest = TestList
+  [ TestLabel  "Fail when exceeding limit for nat set" $
+      testQuery @(Set Nat) [s| all (1|2) |]
+        "Set size limit exceeded for nat set. Max size: 8. Size: 9"
+  , TestLabel "Fail when exceeding limit for word set" $
+      testQuery @(Set Byte)
+        [s| all ( 1 : byte | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 ) |]
+        "Set size limit exceeded for byte set. Max size: 8. Size: 9"
+  , TestLabel "Fail when exceeding limit for word set" $
+      testQuery @(Set Text) [s| all ("foo"|"bar") |]
+        "Set size limit exceeded for standard set. Max size: 8. Size: 10"
+  ]
 
-errMsg :: String
-errMsg =
-  "Set size limit exceeded for nat set. Set token: 0. Max size: 8. Size: 9"
+testQuery :: forall ty . (Show ty, Type ty) => Text -> String -> Test
+testQuery query errMsg =
+  dbTestCaseSettings [setMaxSetSize 8] $ \env repo -> do
+    r <- try (runQuery_ env repo $ angleData query)
+          :: IO (Either FFIError [ty])
+    case r of
+      Left ffiExc -> do
+        let runtimeMsg =
+              Prelude.take (Prelude.length errMsg) (ffiErrorMessage ffiExc)
+        assertEqual "Exception" runtimeMsg errMsg
+      Right v -> do
+        print v
+        assert False