Improve testshard KIKIMR-11082

6 files changed, 99 insertions, 30 deletions

diff --git a/ydb/core/test_tablet/load_actor_delete.cpp b/ydb/core/test_tablet/load_actor_delete.cpp
index 11683caa3a..e63e717b15 100644
--- a/ydb/core/test_tablet/load_actor_delete.cpp
+++ b/ydb/core/test_tablet/load_actor_delete.cpp
@@ -4,10 +4,9 @@ namespace NKikimr::NTestShard {
 
     void TLoadActor::IssueDelete() {
         std::vector<TString> options;
-        options.reserve(Keys.size());
-        for (const auto& [key, info] : Keys) {
-            if (info.ConfirmedState == info.PendingState && info.ConfirmedState == ::NTestShard::TStateServer::CONFIRMED &&
-                    !KeysBeingRead.contains(key)) {
+        options.reserve(ConfirmedKeys.size());
+        for (const TString& key : ConfirmedKeys) {
+            if (!KeysBeingRead.contains(key)) {
                 options.emplace_back(key);
             }
         }
@@ -30,14 +29,14 @@ namespace NKikimr::NTestShard {
 
             STLOG(PRI_INFO, TEST_SHARD, TS09, "deleting data", (TabletId, TabletId), (Key, key));
 
-            const auto [difIt, difInserted] = DeletesInFlight.try_emplace(record.GetCookie(), key);
-            Y_VERIFY(difInserted);
-            Y_VERIFY(difIt->second.KeysInQuery.size() == 1);
-
             const auto it = Keys.find(key);
             Y_VERIFY(it != Keys.end());
             RegisterTransition(*it, ::NTestShard::TStateServer::CONFIRMED, ::NTestShard::TStateServer::DELETE_PENDING, std::move(ev));
 
+            const auto [difIt, difInserted] = DeletesInFlight.try_emplace(record.GetCookie(), std::move(key));
+            Y_VERIFY(difInserted);
+            Y_VERIFY(difIt->second.KeysInQuery.size() == 1);
+
             BytesOfData -= it->second.Len;
             BytesProcessed += it->second.Len;
             ++KeysDeleted;
diff --git a/ydb/core/test_tablet/load_actor_impl.cpp b/ydb/core/test_tablet/load_actor_impl.cpp
index 84b696096e..693cf7adcc 100644
--- a/ydb/core/test_tablet/load_actor_impl.cpp
+++ b/ydb/core/test_tablet/load_actor_impl.cpp
@@ -10,6 +10,21 @@ namespace NKikimr::NTestShard {
         , Settings(settings)
     {}
 
+    TLoadActor::~TLoadActor() {
+        ClearKeys();
+    }
+
+    void TLoadActor::ClearKeys() {
+        for (auto& [key, info] : Keys) {
+            Y_VERIFY((info.ConfirmedState != ::NTestShard::TStateServer::CONFIRMED && info.ConfirmedKeyIndex == Max<size_t>()) ||
+                (info.ConfirmedState == ::NTestShard::TStateServer::CONFIRMED && info.ConfirmedKeyIndex != Max<size_t>() &&
+                 ConfirmedKeys[info.ConfirmedKeyIndex] == key));
+            info.ConfirmedKeyIndex = Max<size_t>();
+        }
+        Keys.clear();
+        ConfirmedKeys.clear();
+    }
+
     void TLoadActor::Bootstrap(const TActorId& parentId) {
         STLOG(PRI_DEBUG, TEST_SHARD, TS31, "TLoadActor::Bootstrap", (TabletId, TabletId));
         TabletActorId = parentId;
@@ -173,7 +188,7 @@ namespace NKikimr::NTestShard {
                 DeletesInFlight.erase(it);
             }
             if (const auto it = ReadsInFlight.find(record.GetCookie()); it != ReadsInFlight.end()) {
-                const auto& [key, offset, size, timestamp] = it->second;
+                const auto& [key, offset, size, timestamp, payloadInResponse] = it->second;
                 const auto jt = KeysBeingRead.find(key);
                 Y_VERIFY(jt != KeysBeingRead.end() && jt->second);
                 if (!--jt->second) {
@@ -195,7 +210,7 @@ namespace NKikimr::NTestShard {
             STLOG(PRI_INFO, TEST_SHARD, TS04, "TEvKeyValue::TEvResponse", (TabletId, TabletId), (Msg, makeResponse()));
             ProcessWriteResult(record.GetCookie(), record.GetWriteResult());
             ProcessDeleteResult(record.GetCookie(), record.GetDeleteRangeResult());
-            ProcessReadResult(record.GetCookie(), record.GetReadResult());
+            ProcessReadResult(record.GetCookie(), record.GetReadResult(), *ev->Get());
         }
         if (WritesInFlight.size() != Settings.GetMaxInFlight() && NextWriteTimestamp == TMonotonic::Max()) {
             NextWriteTimestamp = TMonotonic::Now() + GenerateRandomInterval(Settings.GetWritePeriods());
diff --git a/ydb/core/test_tablet/load_actor_impl.h b/ydb/core/test_tablet/load_actor_impl.h
index 089ea83145..6fc923ed13 100644
--- a/ydb/core/test_tablet/load_actor_impl.h
+++ b/ydb/core/test_tablet/load_actor_impl.h
@@ -22,10 +22,15 @@ namespace NKikimr::NTestShard {
             ::NTestShard::TStateServer::EEntityState ConfirmedState = ::NTestShard::TStateServer::ABSENT;
             ::NTestShard::TStateServer::EEntityState PendingState = ::NTestShard::TStateServer::ABSENT;
             std::unique_ptr<TEvKeyValue::TEvRequest> Request;
+            size_t ConfirmedKeyIndex = Max<size_t>();
 
             TKeyInfo(ui32 len)
                 : Len(len)
             {}
+
+            ~TKeyInfo() {
+                Y_VERIFY(ConfirmedKeyIndex == Max<size_t>());
+            }
         };
 
         enum {
@@ -51,6 +56,8 @@ namespace NKikimr::NTestShard {
 
         TLoadActor(ui64 tabletId, ui32 generation, const TActorId tablet,
             const NKikimrClient::TTestShardControlRequest::TCmdInitialize& settings);
+        ~TLoadActor();
+        void ClearKeys();
         void Bootstrap(const TActorId& parentId);
         void PassAway() override;
         void HandleWakeup();
@@ -83,6 +90,7 @@ namespace NKikimr::NTestShard {
         // Key state
 
         std::unordered_map<TString, TKeyInfo> Keys;
+        std::vector<TString> ConfirmedKeys;
 
         using TKey = std::unordered_map<TString, TKeyInfo>::value_type;
 
@@ -133,11 +141,12 @@ namespace NKikimr::NTestShard {
         void HandleWriteOnTime();
         void HandleDoSomeAction();
 
-        std::unordered_map<ui64, std::tuple<TString, ui32, ui32, TMonotonic>> ReadsInFlight;
+        std::unordered_map<ui64, std::tuple<TString, ui32, ui32, TMonotonic, bool>> ReadsInFlight;
         std::unordered_map<TString, ui32> KeysBeingRead;
 
         bool IssueRead();
-        void ProcessReadResult(ui64 cookie, const NProtoBuf::RepeatedPtrField<NKikimrClient::TKeyValueResponse::TReadResult>& results);
+        void ProcessReadResult(ui64 cookie, const NProtoBuf::RepeatedPtrField<NKikimrClient::TKeyValueResponse::TReadResult>& results,
+            TEvKeyValue::TEvResponse& event);
 
         ////////////////////////////////////////////////////////////////////////////////////////////////////////////////
         // KV tablet delete management code
diff --git a/ydb/core/test_tablet/load_actor_read_validate.cpp b/ydb/core/test_tablet/load_actor_read_validate.cpp
index 8effeaefe5..0d4ea88471 100644
--- a/ydb/core/test_tablet/load_actor_read_validate.cpp
+++ b/ydb/core/test_tablet/load_actor_read_validate.cpp
@@ -626,10 +626,16 @@ namespace NKikimr::NTestShard {
         Send(TabletActorId, new TTestShard::TEvSwitchMode(TTestShard::EMode::WRITE));
         ValidationActorId = {};
         BytesProcessed = 0;
+        ClearKeys();
         Keys = std::move(ev->Get()->Keys);
         BytesOfData = 0;
-        for (const auto& [key, info] : Keys) {
+        for (auto& [key, info] : Keys) {
             BytesOfData += info.Len;
+            Y_VERIFY(info.ConfirmedKeyIndex == Max<size_t>());
+            if (info.ConfirmedState == ::NTestShard::TStateServer::CONFIRMED) {
+                info.ConfirmedKeyIndex = ConfirmedKeys.size();
+                ConfirmedKeys.push_back(key);
+            }
         }
         Action();
 
@@ -640,23 +646,21 @@ namespace NKikimr::NTestShard {
     }
 
     bool TLoadActor::IssueRead() {
-        std::vector<TString> options;
-        for (const auto& [key, info] : Keys) {
-            if (info.ConfirmedState == info.PendingState && info.ConfirmedState == ::NTestShard::TStateServer::CONFIRMED) {
-                options.push_back(key);
-            }
-        }
-        if (options.empty()) {
+        if (ConfirmedKeys.empty()) {
             return false;
         }
-        const size_t index = RandomNumber(options.size());
-        const TString& key = options[index];
+
+        const size_t index = RandomNumber(ConfirmedKeys.size());
+        const TString& key = ConfirmedKeys[index];
         ui64 len, seed, id;
         StringSplitter(key).Split(',').CollectInto(&len, &seed, &id);
-        const ui32 offset = RandomNumber(len);
-        const ui32 size = 1 + RandomNumber(len - offset);
+        const ui32 offset = len ? RandomNumber(len) : 0;
+        const ui32 size = len ? 1 + RandomNumber(len - offset) : 0;
 
         auto request = CreateRequest();
+        if (RandomNumber(2u)) {
+            request->Record.SetUsePayloadInResponse(true);
+        }
         auto *cmdRead = request->Record.AddCmdRead();
         cmdRead->SetKey(key);
         cmdRead->SetOffset(offset);
@@ -664,18 +668,20 @@ namespace NKikimr::NTestShard {
 
         STLOG(PRI_INFO, TEST_SHARD, TS16, "reading key", (TabletId, TabletId), (Key, key), (Offset, offset), (Size, size));
 
-        ReadsInFlight.try_emplace(request->Record.GetCookie(), key, offset, size, TActivationContext::Monotonic());
+        ReadsInFlight.try_emplace(request->Record.GetCookie(), key, offset, size, TActivationContext::Monotonic(),
+            request->Record.GetUsePayloadInResponse());
         ++KeysBeingRead[key];
         Send(TabletActorId, request.release());
 
         return true;
     }
 
-    void TLoadActor::ProcessReadResult(ui64 cookie, const NProtoBuf::RepeatedPtrField<NKikimrClient::TKeyValueResponse::TReadResult>& results) {
+    void TLoadActor::ProcessReadResult(ui64 cookie, const NProtoBuf::RepeatedPtrField<NKikimrClient::TKeyValueResponse::TReadResult>& results,
+            TEvKeyValue::TEvResponse& event) {
         for (const auto& result : results) {
             auto node = ReadsInFlight.extract(cookie);
             Y_VERIFY(node);
-            auto& [key, offset, size, timestamp] = node.mapped();
+            auto& [key, offset, size, timestamp, payloadInResponse] = node.mapped();
             const TMonotonic now = TActivationContext::Monotonic();
 
             auto it = KeysBeingRead.find(key);
@@ -694,9 +700,17 @@ namespace NKikimr::NTestShard {
             Y_VERIFY(result.GetStatus() == NKikimrProto::OK || result.GetStatus() == NKikimrProto::ERROR);
 
             if (result.GetStatus() == NKikimrProto::OK) {
-                const TString value = result.GetValue();
+                TRope value;
+                if (payloadInResponse) {
+                    Y_VERIFY(result.GetDataCase() == NKikimrClient::TKeyValueResponse::TReadResult::kPayloadId);
+                    value = event.GetPayload(result.GetPayloadId());
+                } else {
+                    Y_VERIFY(result.GetDataCase() == NKikimrClient::TKeyValueResponse::TReadResult::kValue);
+                    value = TRope(result.GetValue());
+                }
 
-                Y_VERIFY_S(offset < len && size <= len - offset && value.size() == size && data.substr(offset, size) == value,
+                Y_VERIFY_S((offset < len || !len) && size <= len - offset && value.size() == size &&
+                        data.substr(offset, size) == value.ConvertToString(),
                     "TabletId# " << TabletId << " Key# " << key << " value mismatch"
                     << " value.size# " << value.size() << " offset# " << offset << " size# " << size);
 
diff --git a/ydb/core/test_tablet/load_actor_state.cpp b/ydb/core/test_tablet/load_actor_state.cpp
index a1922120dc..a574c767fa 100644
--- a/ydb/core/test_tablet/load_actor_state.cpp
+++ b/ydb/core/test_tablet/load_actor_state.cpp
@@ -15,6 +15,26 @@ namespace NKikimr::NTestShard {
         Y_VERIFY(from != ::NTestShard::TStateServer::DELETED);
         Y_VERIFY(to != ::NTestShard::TStateServer::ABSENT);
 
+        // unconfirm the key
+        if (from == ::NTestShard::TStateServer::CONFIRMED) {
+            Y_VERIFY(key.second.ConfirmedKeyIndex != Max<size_t>() && key.second.ConfirmedKeyIndex < ConfirmedKeys.size());
+            auto& cell = ConfirmedKeys[key.second.ConfirmedKeyIndex];
+            Y_VERIFY(cell == key.first);
+            std::swap(cell, ConfirmedKeys.back());
+            ConfirmedKeys.pop_back();
+            if (key.second.ConfirmedKeyIndex != ConfirmedKeys.size()) {
+                const auto it = Keys.find(cell);
+                Y_VERIFY(it != Keys.end());
+                Y_VERIFY(it->second.ConfirmedKeyIndex == ConfirmedKeys.size());
+                it->second.ConfirmedKeyIndex = key.second.ConfirmedKeyIndex;
+            }
+            key.second.ConfirmedKeyIndex = Max<size_t>();
+        } else if (to == ::NTestShard::TStateServer::CONFIRMED) {
+            Y_VERIFY(key.second.ConfirmedKeyIndex == Max<size_t>());
+            key.second.ConfirmedKeyIndex = ConfirmedKeys.size();
+            ConfirmedKeys.push_back(key.first);
+        }
+
         if (!Settings.HasStorageServerHost()) {
             if (from == ::NTestShard::TStateServer::WRITE_PENDING && to == ::NTestShard::TStateServer::CONFIRMED) {
                 BytesOfData += key.second.Len;
@@ -23,6 +43,9 @@ namespace NKikimr::NTestShard {
                 Keys.erase(key.first);
             } else {
                 key.second.ConfirmedState = key.second.PendingState = to;
+                Y_VERIFY((key.second.ConfirmedState != ::NTestShard::TStateServer::CONFIRMED && key.second.ConfirmedKeyIndex == Max<size_t>()) ||
+                    (key.second.ConfirmedState == ::NTestShard::TStateServer::CONFIRMED && key.second.ConfirmedKeyIndex != Max<size_t>() &&
+                     ConfirmedKeys[key.second.ConfirmedKeyIndex] == key.first));
             }
             if (ev) {
                 Send(TabletActorId, ev.release());
@@ -93,7 +116,12 @@ namespace NKikimr::NTestShard {
             Send(TabletActorId, r.release());
         }
         if (key.second.ConfirmedState == ::NTestShard::TStateServer::DELETED) {
+            Y_VERIFY(key.second.ConfirmedKeyIndex == Max<size_t>());
             Keys.erase(key.first);
+        } else {
+            Y_VERIFY((key.second.ConfirmedState != ::NTestShard::TStateServer::CONFIRMED && key.second.ConfirmedKeyIndex == Max<size_t>()) ||
+                (key.second.ConfirmedState == ::NTestShard::TStateServer::CONFIRMED && key.second.ConfirmedKeyIndex != Max<size_t>() &&
+                 ConfirmedKeys[key.second.ConfirmedKeyIndex] == key.first));
         }
 
         // perform some action if possible
diff --git a/ydb/core/test_tablet/load_actor_write.cpp b/ydb/core/test_tablet/load_actor_write.cpp
index 5242b7dba5..0e9ba33911 100644
--- a/ydb/core/test_tablet/load_actor_write.cpp
+++ b/ydb/core/test_tablet/load_actor_write.cpp
@@ -22,7 +22,11 @@ namespace NKikimr::NTestShard {
         auto& r = ev->Record;
         auto *write = r.AddCmdWrite();
         write->SetKey(key);
-        write->SetValue(value);
+        if (RandomNumber(2u)) {
+            write->SetPayloadId(ev->AddPayload(TRope(value)));
+        } else {
+            write->SetValue(value);
+        }
         if (isInline) {
             write->SetStorageChannel(NKikimrClient::TKeyValueRequest::INLINE);
         }