Fix possible segfault in gracejoin during memory preallocation

commit_hash:17e8d9ea9ee4e9ae0722424efaf1cc4d03a802a5

2 files changed, 67 insertions, 6 deletions

diff --git a/yql/essentials/minikql/comp_nodes/mkql_grace_join_imp.cpp b/yql/essentials/minikql/comp_nodes/mkql_grace_join_imp.cpp
index 8c9b8a6a22..16a94f457c 100644
--- a/yql/essentials/minikql/comp_nodes/mkql_grace_join_imp.cpp
+++ b/yql/essentials/minikql/comp_nodes/mkql_grace_join_imp.cpp
@@ -363,12 +363,12 @@ bool TTable::TryToPreallocateMemoryForJoin(TTable & t1, TTable & t2, EJoinKind /
             bucketForPreallocation.JoinSlots.reserve(nSlots*slotSize);
         } catch (TMemoryLimitExceededException) {
             for (ui64 i = 0; i < bucket; ++i) {
-                GraceJoin::TTableBucket * b1 = &JoinTable1->TableBuckets[i];
-                b1->JoinSlots.resize(0);
-                b1->JoinSlots.shrink_to_fit();
-                GraceJoin::TTableBucket * b2 = &JoinTable2->TableBuckets[i];
-                b2->JoinSlots.resize(0);
-                b2->JoinSlots.shrink_to_fit();
+                auto& b1 = t1.TableBuckets[i];
+                b1.JoinSlots.resize(0);
+                b1.JoinSlots.shrink_to_fit();
+                auto& b2 = t2.TableBuckets[i];
+                b2.JoinSlots.resize(0);
+                b2.JoinSlots.shrink_to_fit();
             }
             return false;
         }
diff --git a/yql/essentials/minikql/comp_nodes/ut/mkql_grace_join_ut.cpp b/yql/essentials/minikql/comp_nodes/ut/mkql_grace_join_ut.cpp
index 529250fb1d..e8fbd575a9 100644
--- a/yql/essentials/minikql/comp_nodes/ut/mkql_grace_join_ut.cpp
+++ b/yql/essentials/minikql/comp_nodes/ut/mkql_grace_join_ut.cpp
@@ -194,6 +194,67 @@ Y_UNIT_TEST_SUITE(TMiniKQLGraceJoinImpTest) {
     constexpr ui64 SmallTableTuples = 150000;
     constexpr ui64 BigTupleSize = 40;
 
+    Y_UNIT_TEST_TWIN(TestTryToPreallocateMemoryForJoin, EXCEPTION) {
+        TSetup<false> setup;
+        ui64 tuple[11] = {0,1,2,3,4,5,6,7,8,9,10};
+        ui32 strSizes[2] = {4, 4};
+        char * strVals[] = {(char *)"aaaaa", (char *)"bbbb"};
+
+        char * bigStrVal[] = {(char *)"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
+                             (char *)"bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"};
+        ui32 bigStrSize[2] = {151, 151};
+
+
+        GraceJoin::TTable bigTable(1,1,1,1);
+        GraceJoin::TTable smallTable(1,1,1,1);
+        GraceJoin::TTable joinTable(1,1,1,1);
+
+        const ui64 TupleSize = 1024;
+
+        ui64 bigTuple[TupleSize];
+
+        std::mt19937_64 rng; // deterministic PRNG
+
+        std::uniform_int_distribution<ui64> dist(0, 10000 - 1);
+
+        for (ui64 i = 0; i < TupleSize; i++) {
+            bigTuple[i] = dist(rng);
+        }
+
+
+        std::uniform_int_distribution<ui64> smallDist(0, SmallTableTuples - 1);
+
+        smallTable.AddTuple(tuple, bigStrVal, bigStrSize);
+
+        for ( ui64 i = 0; i < SmallTableTuples + 1; i++) {
+            tuple[1] = smallDist(rng);
+            tuple[2] = tuple[1];
+            smallTable.AddTuple(tuple, strVals, strSizes);
+        }
+
+
+        for ( ui64 i = 0; i < BigTableTuples; i++) {
+            tuple[1] = smallDist(rng);
+            tuple[2] = tuple[1];
+            bigTable.AddTuple(tuple, strVals, strSizes);
+        }
+
+        ui64 allocationsCount = 0;
+        if (EXCEPTION) {
+            TlsAllocState->SetLimit(1);
+            TlsAllocState->SetIncreaseMemoryLimitCallback([&allocationsCount](ui64, ui64 required) {
+                // Preallocate memory for some buckets before fail
+                if (allocationsCount++ > 5) {
+                    throw TMemoryLimitExceededException();
+                }
+                TlsAllocState->SetLimit(required);
+            });
+        }
+
+        bool preallocationResult = joinTable.TryToPreallocateMemoryForJoin(smallTable, bigTable, EJoinKind::Inner, true, true);
+        UNIT_ASSERT_EQUAL(preallocationResult, !EXCEPTION);
+    }
+
     Y_UNIT_TEST_LLVM(TestImp1) {
             TSetup<LLVM> setup;
             ui64 tuple[11] = {0,1,2,3,4,5,6,7,8,9,10};