Make it possible to change in-memory setting for tables

snaury · snaury · commit f6685aede56b · 2024-11-28T14:28:45.000Z
diff --git a/ydb/core/tablet_flat/flat_dbase_apply.cpp b/ydb/core/tablet_flat/flat_dbase_apply.cpp
@@ -97,7 +97,14 @@ bool TSchemeModifier::Apply(const TAlterRecord &delta)
         ui32 large = delta.HasLarge() ? delta.GetLarge() : family.Large;
 
         Y_ABORT_UNLESS(ui32(cache) <= 2, "Invalid pages cache policy value");
-        changes |= ChangeTableSetting(table, family.Cache, cache);
+        if (family.Cache != cache) {
+            PreserveTable(table);
+            family.Cache = cache;
+            changes |= true;
+            if (cache == ECache::Ever) {
+                AddNeedToCacheTable(table);
+            }
+        }
         changes |= ChangeTableSetting(table, family.Codec, codec);
         changes |= ChangeTableSetting(table, family.Small, small);
         changes |= ChangeTableSetting(table, family.Large, large);
@@ -405,6 +412,15 @@ bool TSchemeModifier::SetCompactionPolicy(ui32 tid, const NKikimrCompaction::TCo
     return true;
 }
 
+void TSchemeModifier::AddNeedToCacheTable(ui32 tid) noexcept
+{
+    if (Scheme.NeedToCacheTables.insert(tid).second) {
+        if (RollbackState) {
+            RollbackState->NeedToCacheTables.insert(tid);
+        }
+    }
+}
+
 void TSchemeModifier::PreserveTable(ui32 tid) noexcept
 {
     if (RollbackState && !RollbackState->Tables.contains(tid)) {
diff --git a/ydb/core/tablet_flat/flat_dbase_apply.h b/ydb/core/tablet_flat/flat_dbase_apply.h
@@ -83,6 +83,8 @@ namespace NTable {
             return false;
         }
 
+        void AddNeedToCacheTable(ui32 tid) noexcept;
+
         void PreserveTable(ui32 tid) noexcept;
         void PreserveExecutor() noexcept;
         void PreserveRedo() noexcept;
diff --git a/ydb/core/tablet_flat/flat_dbase_naked.h b/ydb/core/tablet_flat/flat_dbase_naked.h
@@ -361,6 +361,7 @@ namespace NTable {
             SchemeRollbackState.Tables.clear();
             SchemeRollbackState.Executor.reset();
             SchemeRollbackState.Redo.reset();
+            SchemeRollbackState.NeedToCacheTables.clear();
         }
 
         void RollbackTransaction() noexcept
@@ -401,6 +402,10 @@ namespace NTable {
             // Note: we assume schema rollback is very rare,
             // so it doesn't have to be efficient
             TScheme& scheme = *Scheme;
+            for (ui32 tid : SchemeRollbackState.NeedToCacheTables) {
+                scheme.NeedToCacheTables.erase(tid);
+            }
+            SchemeRollbackState.NeedToCacheTables.clear();
             if (SchemeRollbackState.Redo) {
                 scheme.Redo = *SchemeRollbackState.Redo;
                 SchemeRollbackState.Redo.reset();
diff --git a/ydb/core/tablet_flat/flat_dbase_scheme.h b/ydb/core/tablet_flat/flat_dbase_scheme.h
@@ -171,6 +171,9 @@ class TScheme {
     THashMap<TString, ui32> TableNames;
     TExecutorInfo Executor;
     TRedo Redo;
+
+    // A set of tables that have their in-memory state enabled
+    mutable THashSet<ui32> NeedToCacheTables;
 };
 
 /**
@@ -186,6 +189,8 @@ struct TSchemeRollbackState {
     std::optional<TScheme::TExecutorInfo> Executor;
     // Previous redo settings if modified
     std::optional<TScheme::TRedo> Redo;
+    // Newly added NeedToCacheTables entries
+    THashSet<ui32> NeedToCacheTables;
 };
 
 /**
diff --git a/ydb/core/tablet_flat/flat_executor.cpp b/ydb/core/tablet_flat/flat_executor.cpp
@@ -415,6 +415,8 @@ void TExecutor::ActivateFollower(const TActorContext &ctx) {
     RecreatePageCollectionsCache();
     ReflectSchemeSettings();
 
+    RequestInMemPagesForDatabase();
+
     Become(&TThis::StateFollower);
     Stats->IsActive = true;
     Stats->FollowerId = FollowerId;
@@ -664,16 +666,28 @@ void TExecutor::TranslateCacheTouchesToSharedCache() {
 }
 
 void TExecutor::RequestInMemPagesForDatabase() {
+    const auto& scheme = Scheme();
+    for (auto& pr : Scheme().Tables) {
+        scheme.NeedToCacheTables.insert(pr.first);
+    }
+    RequestInMemPagesForNeedToCacheTables();
+}
+
+void TExecutor::RequestInMemPagesForNeedToCacheTables() {
     const auto &scheme = Scheme();
-    for (auto &sxpair : scheme.Tables) {
-        auto stickyColumns = GetStickyColumns(sxpair.first);
+    for (ui32 tid : scheme.NeedToCacheTables) {
+        if (!scheme.Tables.contains(tid)) {
+            continue;
+        }
+        auto stickyColumns = GetStickyColumns(tid);
         if (stickyColumns) {
-            auto subset = Database->Subset(sxpair.first, NTable::TEpoch::Max(), { } , { });
+            auto subset = Database->Subset(tid, NTable::TEpoch::Max(), { } , { });
 
             for (auto &partView: subset->Flatten)
-                RequestInMemPagesForPartStore(sxpair.first, partView, stickyColumns);
+                RequestInMemPagesForPartStore(tid, partView, stickyColumns);
         }
     }
+    scheme.NeedToCacheTables.clear();
 }
 
 TExecutorCaches TExecutor::CleanupState() {
@@ -984,6 +998,7 @@ void TExecutor::ApplyFollowerUpdate(THolder<TEvTablet::TFUpdateBody> update) {
         if (schemeUpdate) {
             ReadResourceProfile();
             ReflectSchemeSettings();
+            RequestInMemPagesForNeedToCacheTables();
             Owner->OnFollowerSchemaUpdated();
         }
 
@@ -2081,6 +2096,7 @@ void TExecutor::CommitTransactionLog(TAutoPtr<TSeat> seat, TPageCollectionTxEnv
 
             ReadResourceProfile();
             ReflectSchemeSettings();
+            RequestInMemPagesForNeedToCacheTables();
 
             // For every table that changed strategy we need to generate a
             // special part switch that notifies bootlogic about new strategy
diff --git a/ydb/core/tablet_flat/flat_executor.h b/ydb/core/tablet_flat/flat_executor.h
@@ -530,6 +530,7 @@ class TExecutor
 
     void TranslateCacheTouchesToSharedCache();
     void RequestInMemPagesForDatabase();
+    void RequestInMemPagesForNeedToCacheTables();
     void RequestInMemPagesForPartStore(ui32 tableId, const NTable::TPartView &partView, const THashSet<NTable::TTag> &stickyColumns);
     THashSet<NTable::TTag> GetStickyColumns(ui32 tableId);
     void RequestFromSharedCache(TAutoPtr<NPageCollection::TFetch> fetch,
diff --git a/ydb/core/ydb_convert/column_families.h b/ydb/core/ydb_convert/column_families.h
@@ -183,13 +183,15 @@ namespace NKikimr {
                 case Ydb::FeatureFlag::STATUS_UNSPECIFIED:
                     break;
                 case Ydb::FeatureFlag::ENABLED:
-                    *code = Ydb::StatusIds::BAD_REQUEST;
-                    *error = TStringBuilder()
-                        << "Setting keep_in_memory to ENABLED is not supported in column family '"
-                        << familySettings.name() << "'";
-                    return false;
+                    if (!AppData()->FeatureFlags.GetEnablePublicApiKeepInMemory()) {
+                        *code = Ydb::StatusIds::BAD_REQUEST;
+                        *error = "Setting keep_in_memory to ENABLED is not allowed";
+                        return false;
+                    }
+                    family->SetColumnCache(NKikimrSchemeOp::ColumnCacheEver);
+                    break;
                 case Ydb::FeatureFlag::DISABLED:
-                    family->ClearColumnCache();
+                    family->SetColumnCache(NKikimrSchemeOp::ColumnCacheNone);
                     break;
                 default:
                     *code = Ydb::StatusIds::BAD_REQUEST;
diff --git a/ydb/public/sdk/cpp/client/ydb_table/table.cpp b/ydb/public/sdk/cpp/client/ydb_table/table.cpp
@@ -1115,6 +1115,11 @@ TColumnFamilyBuilder& TColumnFamilyBuilder::SetCompression(EColumnFamilyCompress
     return *this;
 }
 
+TColumnFamilyBuilder& TColumnFamilyBuilder::SetKeepInMemory(bool enabled) {
+    Impl_->Proto.set_keep_in_memory(enabled ? Ydb::FeatureFlag::ENABLED : Ydb::FeatureFlag::DISABLED);
+    return *this;
+}
+
 TColumnFamilyDescription TColumnFamilyBuilder::Build() const {
     return TColumnFamilyDescription(Impl_->Proto);
 }
diff --git a/ydb/public/sdk/cpp/client/ydb_table/table.h b/ydb/public/sdk/cpp/client/ydb_table/table.h
@@ -806,6 +806,7 @@ class TColumnFamilyBuilder {
 
     TColumnFamilyBuilder& SetData(const TString& media);
     TColumnFamilyBuilder& SetCompression(EColumnFamilyCompression compression);
+    TColumnFamilyBuilder& SetKeepInMemory(bool enabled);
 
     TColumnFamilyDescription Build() const;
 
@@ -868,6 +869,11 @@ class TTableColumnFamilyBuilder {
         return *this;
     }
 
+    TTableColumnFamilyBuilder& SetKeepInMemory(bool enabled) {
+        Builder_.SetKeepInMemory(enabled);
+        return *this;
+    }
+
     TTableBuilder& EndColumnFamily();
 
 private:
@@ -1486,6 +1492,11 @@ class TAlterColumnFamilyBuilder {
         return *this;
     }
 
+    TAlterColumnFamilyBuilder& SetKeepInMemory(bool enabled) {
+        Builder_.SetKeepInMemory(enabled);
+        return *this;
+    }
+
     TAlterTableSettings& EndAddColumnFamily();
     TAlterTableSettings& EndAlterColumnFamily();
 
diff --git a/ydb/services/ydb/ydb_ut.cpp b/ydb/services/ydb/ydb_ut.cpp
@@ -824,6 +824,87 @@ Y_UNIT_TEST_SUITE(TGRpcNewClient) {
         client.CreateSession().Apply(createSessionHandler).Wait();
         UNIT_ASSERT(done);
     }
+
+    Y_UNIT_TEST(InMemoryTables) {
+        TKikimrWithGrpcAndRootSchemaNoSystemViews server;
+        server.Server_->GetRuntime()->GetAppData().FeatureFlags.SetEnablePublicApiKeepInMemory(true);
+
+        ui16 grpc = server.GetPort();
+        TString location = TStringBuilder() << "localhost:" << grpc;
+
+        auto connection = NYdb::TDriver(
+            TDriverConfig()
+                .SetEndpoint(location));
+
+        auto client = NYdb::NTable::TTableClient(connection);
+        auto createSessionResult = client.CreateSession().ExtractValueSync();
+        UNIT_ASSERT(!createSessionResult.IsTransportError());
+        auto session = createSessionResult.GetSession();
+
+        auto createTableResult = session.CreateTable("/Root/Table", client.GetTableBuilder()
+            .AddNullableColumn("Key", EPrimitiveType::Int32)
+            .AddNullableColumn("Value", EPrimitiveType::String)
+            .SetPrimaryKeyColumn("Key")
+            // Note: only needed because this test doesn't initial table profiles
+            .BeginStorageSettings()
+                .SetTabletCommitLog0("ssd")
+                .SetTabletCommitLog1("ssd")
+            .EndStorageSettings()
+            .BeginColumnFamily("default")
+                .SetData("ssd")
+                .SetKeepInMemory(true)
+            .EndColumnFamily()
+            .Build()).ExtractValueSync();
+        UNIT_ASSERT_C(createTableResult.IsSuccess(), (NYdb::TStatus&)createTableResult);
+
+        {
+            auto describeTableResult = session.DescribeTable("/Root/Table").ExtractValueSync();
+            UNIT_ASSERT_C(describeTableResult.IsSuccess(), (NYdb::TStatus&)describeTableResult);
+            auto desc = describeTableResult.GetTableDescription();
+            auto families = desc.GetColumnFamilies();
+            UNIT_ASSERT_VALUES_EQUAL(families.size(), 1u);
+            auto family = families.at(0);
+            UNIT_ASSERT_VALUES_EQUAL(family.GetKeepInMemory(), true);
+        }
+
+        {
+            auto alterTableResult = session.AlterTable("/Root/Table", NYdb::NTable::TAlterTableSettings()
+                .BeginAlterColumnFamily("default")
+                    .SetKeepInMemory(false)
+                .EndAlterColumnFamily()).ExtractValueSync();
+            UNIT_ASSERT_C(alterTableResult.IsSuccess(), (NYdb::TStatus&)alterTableResult);
+        }
+
+        {
+            auto describeTableResult = session.DescribeTable("/Root/Table").ExtractValueSync();
+            UNIT_ASSERT_C(describeTableResult.IsSuccess(), (NYdb::TStatus&)describeTableResult);
+            auto desc = describeTableResult.GetTableDescription();
+            auto families = desc.GetColumnFamilies();
+            UNIT_ASSERT_VALUES_EQUAL(families.size(), 1u);
+            auto family = families.at(0);
+            // Note: server cannot currently distinguish between implicitly
+            // unset and explicitly disabled, so it returns the former.
+            UNIT_ASSERT_VALUES_EQUAL(family.GetKeepInMemory(), Nothing());
+        }
+
+        {
+            auto alterTableResult = session.AlterTable("/Root/Table", NYdb::NTable::TAlterTableSettings()
+                .BeginAlterColumnFamily("default")
+                    .SetKeepInMemory(true)
+                .EndAlterColumnFamily()).ExtractValueSync();
+            UNIT_ASSERT_C(alterTableResult.IsSuccess(), (NYdb::TStatus&)alterTableResult);
+        }
+
+        {
+            auto describeTableResult = session.DescribeTable("/Root/Table").ExtractValueSync();
+            UNIT_ASSERT_C(describeTableResult.IsSuccess(), (NYdb::TStatus&)describeTableResult);
+            auto desc = describeTableResult.GetTableDescription();
+            auto families = desc.GetColumnFamilies();
+            UNIT_ASSERT_VALUES_EQUAL(families.size(), 1u);
+            auto family = families.at(0);
+            UNIT_ASSERT_VALUES_EQUAL(family.GetKeepInMemory(), true);
+        }
+    }
 }
 
 static TString CreateSession(std::shared_ptr<grpc::Channel> channel) {

Original file line number	Diff line number	Diff line change
`@@ -83,6 +83,8 @@ namespace NTable {`
`83`	`83`	`return false;`
`84`	`84`	`}`
`85`	`85`
	`86`	`+ void AddNeedToCacheTable(ui32 tid) noexcept;`
	`87`	`+`
`86`	`88`	`void PreserveTable(ui32 tid) noexcept;`
`87`	`89`	`void PreserveExecutor() noexcept;`
`88`	`90`	`void PreserveRedo() noexcept;`
Original file line number	Diff line number	Diff line change
`@@ -1115,6 +1115,11 @@ TColumnFamilyBuilder& TColumnFamilyBuilder::SetCompression(EColumnFamilyCompress`
`1115`	`1115`	`return *this;`
`1116`	`1116`	`}`
`1117`	`1117`
	`1118`	`+TColumnFamilyBuilder& TColumnFamilyBuilder::SetKeepInMemory(bool enabled) {`
	`1119`	`+ Impl_->Proto.set_keep_in_memory(enabled ? Ydb::FeatureFlag::ENABLED : Ydb::FeatureFlag::DISABLED);`
	`1120`	`+ return *this;`
	`1121`	`+}`
	`1122`	`+`
`1118`	`1123`	`TColumnFamilyDescription TColumnFamilyBuilder::Build() const {`
`1119`	`1124`	`return TColumnFamilyDescription(Impl_->Proto);`
`1120`	`1125`	`}`