fix race for unsafe variable in tests

3 files changed, 12 insertions, 9 deletions

diff --git a/ydb/core/testlib/controllers/abstract.h b/ydb/core/testlib/controllers/abstract.h
index 523dc7451f6..b0d30f67bba 100644
--- a/ydb/core/testlib/controllers/abstract.h
+++ b/ydb/core/testlib/controllers/abstract.h
@@ -1,6 +1,7 @@
 #pragma once
 #include <ydb/library/accessor/accessor.h>
 #include <util/generic/singleton.h>
+#include <util/generic/refcount.h>
 #include <memory>
 
 namespace NKikimr::NOlap::NIndexedReader {
@@ -11,7 +12,7 @@ namespace NKikimr::NYDBTest {
 
 class ICSController {
 private:
-    YDB_READONLY(ui32, OnSortingPolicyCounter, 0);
+    YDB_READONLY(TAtomicCounter, OnSortingPolicyCounter, 0);
 protected:
     virtual bool DoOnSortingPolicy(std::shared_ptr<NOlap::NIndexedReader::IOrderPolicy> /*policy*/) {
         return true;
@@ -20,7 +21,7 @@ public:
     using TPtr = std::shared_ptr<ICSController>;
     virtual ~ICSController() = default;
     bool OnSortingPolicy(std::shared_ptr<NOlap::NIndexedReader::IOrderPolicy> policy) {
-        ++OnSortingPolicyCounter;
+        OnSortingPolicyCounter.Inc();
         return DoOnSortingPolicy(policy);
     }
 };
diff --git a/ydb/core/tx/columnshard/testlib/controller.cpp b/ydb/core/tx/columnshard/testlib/controller.cpp
index 16e722803b4..2ed500cae7f 100644
--- a/ydb/core/tx/columnshard/testlib/controller.cpp
+++ b/ydb/core/tx/columnshard/testlib/controller.cpp
@@ -6,13 +6,17 @@ namespace NKikimr::NYDBTest::NColumnShard {
 
 bool TController::DoOnSortingPolicy(std::shared_ptr<NOlap::NIndexedReader::IOrderPolicy> policy) {
     if (dynamic_cast<const NOlap::NIndexedReader::TPKSortingWithLimit*>(policy.get())) {
-        ++SortingWithLimit;
+        SortingWithLimit.Inc();
     } else if (dynamic_cast<const NOlap::NIndexedReader::TAnySorting*>(policy.get())) {
-        ++AnySorting;
+        AnySorting.Inc();
     } else {
         Y_VERIFY(false);
     }
     return true;
 }
 
+bool TController::HasPKSortingOnly() const {
+    return SortingWithLimit.Val() && !AnySorting.Val();
+}
+
 }
diff --git a/ydb/core/tx/columnshard/testlib/controller.h b/ydb/core/tx/columnshard/testlib/controller.h
index edb4f7a7118..f1ea14b5a4d 100644
--- a/ydb/core/tx/columnshard/testlib/controller.h
+++ b/ydb/core/tx/columnshard/testlib/controller.h
@@ -5,14 +5,12 @@ namespace NKikimr::NYDBTest::NColumnShard {
 
 class TController: public ICSController {
 private:
-    YDB_READONLY(ui32, SortingWithLimit, 0);
-    YDB_READONLY(ui32, AnySorting, 0);
+    YDB_READONLY(TAtomicCounter, SortingWithLimit, 0);
+    YDB_READONLY(TAtomicCounter, AnySorting, 0);
 protected:
     virtual bool DoOnSortingPolicy(std::shared_ptr<NOlap::NIndexedReader::IOrderPolicy> policy);
 public:
-    bool HasPKSortingOnly() const {
-        return SortingWithLimit && !AnySorting;
-    }
+    bool HasPKSortingOnly() const;
 };
 
 }