Add missing precharges in datashard init tx (#11516)

Author: snaury

ydb/core/tablet_flat/flat_database.cpp

@@ -40,7 +40,6 @@ bool TDatabase::TChangeCounter::operator<(const TChangeCounter& rhs) const {
 
 TDatabase::TDatabase(TDatabaseImpl *databaseImpl) noexcept
     : DatabaseImpl(databaseImpl ? databaseImpl : new TDatabaseImpl(0, new TScheme, nullptr))
-    , NoMoreReadsFlag(true)
 {
 
 }
@@ -59,7 +58,7 @@ TIntrusiveConstPtr<TRowScheme> TDatabase::GetRowScheme(ui32 table) const noexcep
 
 TAutoPtr<TTableIter> TDatabase::Iterate(ui32 table, TRawVals key, TTagsRef tags, ELookup mode) const noexcept
 {
-    Y_ABORT_UNLESS(!NoMoreReadsFlag, "Trying to read after reads prohibited, table %u", table);
+    CheckReadAllowed(table);
 
     const auto seekBy = [](TRawVals key, ELookup mode) {
         if (!key && mode != ELookup::ExactMatch) {
@@ -92,7 +91,7 @@ TAutoPtr<TTableIter> TDatabase::IterateExact(ui32 table, TRawVals key, TTagsRef
         const ITransactionMapPtr& visible,
         const ITransactionObserverPtr& observer) const noexcept
 {
-    Y_ABORT_UNLESS(!NoMoreReadsFlag, "Trying to read after reads prohibited, table %u", table);
+    CheckReadAllowed(table);
 
     auto iter = Require(table)->Iterate(key, tags, Env, ESeek::Exact, snapshot, visible, observer);
 
@@ -144,7 +143,7 @@ TAutoPtr<TTableIter> TDatabase::IterateRange(ui32 table, const TKeyRange& range,
         const ITransactionMapPtr& visible,
         const ITransactionObserverPtr& observer) const noexcept
 {
-    Y_ABORT_UNLESS(!NoMoreReadsFlag, "Trying to read after reads prohibited, table %u", table);
+    CheckReadAllowed(table);
 
     Y_DEBUG_ABORT_UNLESS(!IsAmbiguousRange(range, Require(table)->GetScheme()->Keys->Size()),
         "%s", IsAmbiguousRangeReason(range, Require(table)->GetScheme()->Keys->Size()));
@@ -171,7 +170,7 @@ TAutoPtr<TTableReverseIter> TDatabase::IterateRangeReverse(ui32 table, const TKe
         const ITransactionMapPtr& visible,
         const ITransactionObserverPtr& observer) const noexcept
 {
-    Y_ABORT_UNLESS(!NoMoreReadsFlag, "Trying to read after reads prohibited, table %u", table);
+    CheckReadAllowed(table);
 
     Y_DEBUG_ABORT_UNLESS(!IsAmbiguousRange(range, Require(table)->GetScheme()->Keys->Size()),
         "%s", IsAmbiguousRangeReason(range, Require(table)->GetScheme()->Keys->Size()));
@@ -226,7 +225,7 @@ EReady TDatabase::Select(ui32 table, TRawVals key, TTagsRef tags, TRowState &row
         const ITransactionObserverPtr& observer) const noexcept
 {
     TempIterators.clear();
-    Y_ABORT_UNLESS(!NoMoreReadsFlag, "Trying to read after reads prohibited, table %u", table);
+    CheckReadAllowed(table);
 
     auto prevSieved = stats.Sieved;
     auto prevWeeded = stats.Weeded;
@@ -268,7 +267,8 @@ void TDatabase::CalculateReadSize(TSizeEnv& env, ui32 table, TRawVals minKey, TR
                                   TTagsRef tags, ui64 flg, ui64 items, ui64 bytes,
                                   EDirection direction, TRowVersion snapshot)
 {
-    Y_ABORT_UNLESS(!NoMoreReadsFlag, "Trying to do precharge after reads prohibited, table %u", table);
+    CheckReadAllowed(table);
+
     TSelectStats stats;
     Require(table)->Precharge(minKey, maxKey, tags, &env, flg, items, bytes, direction, snapshot, stats);
 }
@@ -277,7 +277,8 @@ bool TDatabase::Precharge(ui32 table, TRawVals minKey, TRawVals maxKey,
                     TTagsRef tags, ui64 flg, ui64 items, ui64 bytes,
                     EDirection direction, TRowVersion snapshot)
 {
-    Y_ABORT_UNLESS(!NoMoreReadsFlag, "Trying to do precharge after reads prohibited, table %u", table);
+    CheckPrechargeAllowed(table, minKey, maxKey);
+
     TSelectStats stats;
     auto ready = Require(table)->Precharge(minKey, maxKey, tags, Env, flg, items, bytes, direction, snapshot, stats);
     Change->Stats.ChargeSieved += stats.Sieved;
@@ -428,6 +429,10 @@ void TDatabase::NoMoreReadsForTx() {
     NoMoreReadsFlag = true;
 }
 
+void TDatabase::NoMoreUnprechargedReadsForTx() {
+    NoMoreUnprechargedReadsFlag = true;
+}
+
 void TDatabase::Begin(TTxStamp stamp, IPages& env)
 {
     Y_ABORT_UNLESS(!Redo, "Transaction already in progress");
@@ -438,6 +443,8 @@ void TDatabase::Begin(TTxStamp stamp, IPages& env)
     Change = MakeHolder<TChange>(stamp, DatabaseImpl->Serial());
     Env = &env;
     NoMoreReadsFlag = false;
+    NoMoreUnprechargedReadsFlag = false;
+    PrechargedTables.clear();
 }
 
 void TDatabase::RollbackChanges()
@@ -802,6 +809,30 @@ TTable* TDatabase::RequireForUpdate(ui32 table) const noexcept
     return DatabaseImpl->GetForUpdate(table).Self.Get();
 }
 
+void TDatabase::CheckReadAllowed(ui32 table) const noexcept
+{
+    Y_ABORT_UNLESS(!NoMoreReadsFlag, "Trying to read after reads prohibited, table %u", table);
+    if (NoMoreUnprechargedReadsFlag) [[unlikely]] {
+        Y_DEBUG_ABORT_UNLESS(
+            std::find(PrechargedTables.begin(), PrechargedTables.end(), table) != PrechargedTables.end(),
+            "Trying to read a previously unprecharged table %u", table);
+    }
+}
+
+void TDatabase::CheckPrechargeAllowed(ui32 table, TRawVals minKey, TRawVals maxKey) const noexcept
+{
+    Y_ABORT_UNLESS(!NoMoreReadsFlag, "Trying to precharge after reads prohibited, table %u", table);
+    if (NoMoreUnprechargedReadsFlag) [[unlikely]] {
+        Y_DEBUG_ABORT_UNLESS(
+            std::find(PrechargedTables.begin(), PrechargedTables.end(), table) != PrechargedTables.end(),
+            "Trying to precharge a previously unprecharged table %u", table);
+    } else if (!minKey && !maxKey) {
+        // Note: tables are usually fully precharged only in init transcations
+        // This vector is small or empty most of the time and not worth using a hash table
+        PrechargedTables.push_back(table);
+    }
+}
+
 TGarbage TDatabase::RollUp(TTxStamp stamp, TArrayRef<const char> delta, TArrayRef<const char> redo,
                                 TMemGlobs annex)
 {

ydb/core/tablet_flat/flat_database.h

@@ -193,6 +193,13 @@ class TDatabase {
 
     void NoMoreReadsForTx();
 
+    /**
+     * Will debug assert when current transaction attempts to read tables that
+     * have not been precharged up to this point. Useful to detect missing
+     * precharges that avoid multiple transaction restarts.
+     */
+    void NoMoreUnprechargedReadsForTx();
+
     TAlter& Alter(); /* Begin DDL ALTER script */
 
     TEpoch TxSnapTable(ui32 table);
@@ -291,10 +298,14 @@ class TDatabase {
     TTable* Require(ui32 tableId) const noexcept;
     TTable* RequireForUpdate(ui32 tableId) const noexcept;
 
+    void CheckReadAllowed(ui32 table) const noexcept;
+    void CheckPrechargeAllowed(ui32 table, TRawVals minKey, TRawVals maxKey) const noexcept;
+
 private:
     const THolder<TDatabaseImpl> DatabaseImpl;
 
-    bool NoMoreReadsFlag;
+    bool NoMoreReadsFlag = false;
+    bool NoMoreUnprechargedReadsFlag = false;
     IPages* Env = nullptr;
     THolder<TChange> Change;
     TAutoPtr<TAlter> Alter_;
@@ -305,6 +316,7 @@ class TDatabase {
     TVector<TUpdateOp> ModifiedOps;
 
     mutable TDeque<TPartIter> TempIterators; // Keeps the last result of Select() valid
+    mutable TVector<ui32> PrechargedTables;
 
     TVector<std::function<void()>> OnCommit_;
     TVector<std::function<void()>> OnRollback_;

ydb/core/tx/datashard/datashard__init.cpp

@@ -156,6 +156,7 @@ bool TDataShard::TTxInit::ReadEverything(TTransactionContext &txc) {
         PRECHARGE_SYS_TABLE(Schema::PlanQueue);
         PRECHARGE_SYS_TABLE(Schema::DeadlineQueue);
         PRECHARGE_SYS_TABLE(Schema::SchemaOperations);
+        PRECHARGE_SYS_TABLE(Schema::ScanProgress);
         PRECHARGE_SYS_TABLE(Schema::SplitSrcSnapshots);
         PRECHARGE_SYS_TABLE(Schema::SplitDstReceivedSnapshots);
         PRECHARGE_SYS_TABLE(Schema::Snapshots);
@@ -173,6 +174,9 @@ bool TDataShard::TTxInit::ReadEverything(TTransactionContext &txc) {
         PRECHARGE_SYS_TABLE(Schema::Locks);
         PRECHARGE_SYS_TABLE(Schema::LockRanges);
         PRECHARGE_SYS_TABLE(Schema::LockConflicts);
+        PRECHARGE_SYS_TABLE(Schema::LockVolatileDependencies);
+        PRECHARGE_SYS_TABLE(Schema::LockChangeRecords);
+        PRECHARGE_SYS_TABLE(Schema::ChangeRecordCommits);
         PRECHARGE_SYS_TABLE(Schema::TxVolatileDetails);
         PRECHARGE_SYS_TABLE(Schema::TxVolatileParticipants);
         PRECHARGE_SYS_TABLE(Schema::CdcStreamScans);
@@ -184,6 +188,8 @@ bool TDataShard::TTxInit::ReadEverything(TTransactionContext &txc) {
 #undef PRECHARGE_SYS_TABLE
     }
 
+    txc.DB.NoMoreUnprechargedReadsForTx();
+
     // Reads from Sys table
     LOAD_SYS_UI64(db, Schema::Sys_State, Self->State);
     LOAD_SYS_UI64(db, Schema::Sys_LastLocalTid, Self->LastLocalTid);