Atomically update feature flags at runtime (#10281)

Author: snaury

ydb/core/base/appdata.cpp

@@ -25,6 +25,7 @@
 #include <ydb/core/protos/netclassifier.pb.h>
 #include <ydb/core/protos/datashard_config.pb.h>
 #include <ydb/core/protos/shared_cache.pb.h>
+#include <ydb/core/protos/feature_flags.pb.h>
 #include <ydb/library/pdisk_io/aio.h>
 
 #include <ydb/library/actors/interconnect/poller_tcp.h>
@@ -125,6 +126,14 @@ TAppData::TAppData(
 TAppData::~TAppData()
 {}
 
+void TAppData::InitFeatureFlags(const NKikimrConfig::TFeatureFlags& flags) {
+    Impl->FeatureFlags = flags;
+}
+
+void TAppData::UpdateRuntimeFlags(const NKikimrConfig::TFeatureFlags& flags) {
+    Impl->FeatureFlags.CopyRuntimeFrom(flags);
+}
+
 TIntrusivePtr<IRandomProvider> TAppData::RandomProvider = CreateDefaultRandomProvider();
 TIntrusivePtr<ITimeProvider> TAppData::TimeProvider = CreateDefaultTimeProvider();
 

ydb/core/base/appdata_fwd.h

@@ -64,6 +64,7 @@ namespace NKikimrConfig {
     class TGraphConfig;
     class TMetadataCacheConfig;
     class TMemoryControllerConfig;
+    class TFeatureFlags;
 }
 
 namespace NKikimrReplication {
@@ -271,6 +272,9 @@ struct TAppData {
             TProgramShouldContinue *kikimrShouldContinue);
 
     ~TAppData();
+
+    void InitFeatureFlags(const NKikimrConfig::TFeatureFlags& flags);
+    void UpdateRuntimeFlags(const NKikimrConfig::TFeatureFlags& flags);
 };
 
 inline TAppData* AppData(NActors::TActorSystem* actorSystem) {

ydb/core/base/events.h

@@ -182,6 +182,7 @@ struct TKikimrEvents : TEvents {
         ES_MEMORY = 4259,
         ES_GROUPED_ALLOCATIONS_MANAGER = 4260,
         ES_INCREMENTAL_RESTORE_SCAN = 4261,
+        ES_FEATURE_FLAGS = 4262,
     };
 };
 

ydb/core/base/feature_flags.h

@@ -1,13 +1,14 @@
 #pragma once
-
 #include "defs.h"
-
-#include <ydb/core/protos/feature_flags.pb.h>
+#include "runtime_feature_flags.h"
 
 namespace NKikimr {
 
-class TFeatureFlags: public NKikimrConfig::TFeatureFlags {
-    using TBase = NKikimrConfig::TFeatureFlags;
+/**
+ * Thin wrapper around TRuntimeFeatureFlags, for compatibility with existing code
+ */
+class TFeatureFlags : public TRuntimeFeatureFlags {
+    using TBase = TRuntimeFeatureFlags;
 
 public:
     using TBase::TBase;

ydb/core/base/feature_flags_service.cpp

@@ -0,0 +1,9 @@
+#include "feature_flags_service.h"
+
+namespace NKikimr {
+
+    TActorId MakeFeatureFlagsServiceID() {
+        return TActorId(ui32(0), "featureflags");
+    }
+
+} // namespace NKikimr

ydb/core/base/feature_flags_service.h

@@ -0,0 +1,37 @@
+#pragma once
+#include "defs.h"
+#include "events.h"
+
+namespace NKikimr {
+
+    struct TEvFeatureFlags {
+        enum EEv {
+            EvSubscribe = EventSpaceBegin(TKikimrEvents::ES_FEATURE_FLAGS),
+            EvUnsubscribe,
+            EvChanged,
+            EvEnd
+        };
+
+        static_assert(EvEnd < EventSpaceEnd(TKikimrEvents::ES_FEATURE_FLAGS),
+            "expect EvEnd < EventSpaceEnd(TKikimrEvents::ES_FEATURE_FLAGS)");
+
+        struct TEvSubscribe : public TEventLocal<TEvSubscribe, EvSubscribe> {
+            TEvSubscribe() = default;
+        };
+
+        struct TEvUnsubscribe : public TEventLocal<TEvUnsubscribe, EvUnsubscribe> {
+            TEvUnsubscribe() = default;
+        };
+
+        struct TEvChanged : public TEventLocal<TEvChanged, EvChanged> {
+            TEvChanged() = default;
+        };
+    };
+
+    /**
+     * Returns the service id that may be used to subscribe to feature flags
+     * change notifications.
+     */
+    TActorId MakeFeatureFlagsServiceID();
+
+} // namespace NKikimr

ydb/core/base/generated/codegen/__main__.py

@@ -0,0 +1,84 @@
+import os
+import sys
+from ydb.core.protos.feature_flags_pb2 import TFeatureFlags, RequireRestart
+
+from jinja2 import Environment, FileSystemLoader
+from google.protobuf.descriptor import FieldDescriptor
+
+
+class Slot(object):
+    def __init__(self, name, index):
+        self.name = name
+        self.index = index
+        self.fields = []
+        self.default_value = 0
+        self.runtime_flags_mask = 0
+
+
+class Field(object):
+    def __init__(self, name, slot, has_mask, value_mask, full_mask, default_value, is_runtime):
+        self.name = name
+        self.slot = slot
+        self.has_mask = has_mask
+        self.value_mask = value_mask
+        self.full_mask = full_mask
+        self.default_value = default_value
+        self.is_runtime = is_runtime
+
+
+class EnvironmentByDir(dict):
+    def __missing__(self, template_dir):
+        env = Environment(loader=FileSystemLoader(template_dir))
+        self[template_dir] = env
+        return env
+
+
+def main():
+    slots = []
+    fields = []
+    current_bits = 64
+    for field in TFeatureFlags.DESCRIPTOR.fields:
+        if field.type == FieldDescriptor.TYPE_BOOL:
+            if current_bits + 2 > 64:
+                index = len(slots)
+                slots.append(Slot(name=f'slot{index}', index=index))
+                current_bits = 0
+            shift = current_bits
+            current_bits += 2
+            slot = slots[-1]
+            has_mask = 1 << shift
+            value_mask = 2 << shift
+            full_mask = 3 << shift
+            default_value = 0
+            if field.default_value:
+                default_value = 2 << shift
+            is_runtime = not field.GetOptions().Extensions[RequireRestart]
+            fields.append(Field(
+                name=field.name,
+                slot=slot,
+                has_mask=has_mask,
+                value_mask=value_mask,
+                full_mask=full_mask,
+                default_value=default_value,
+                is_runtime=is_runtime,
+            ))
+            slot.fields.append(fields[-1])
+            slot.default_value |= default_value
+            if is_runtime:
+                slot.runtime_flags_mask |= full_mask
+
+    template_files = sys.argv[1::2]
+    output_files = sys.argv[2::2]
+    env_by_dir = EnvironmentByDir()
+    for (template_file, output_file) in zip(template_files, output_files):
+        (template_dir, template_name) = os.path.split(template_file)
+        env = env_by_dir[template_dir]
+        template = env.get_template(template_name)
+        result = template.render(generator=__file__, slots=slots, fields=fields)
+        with open(output_file, 'w') as f:
+            f.write(result)
+        print(f'Generated {output_file} from {template_name}')
+
+
+if __name__ == '__main__':
+    main()

ydb/core/base/generated/codegen/ya.make

@@ -0,0 +1,11 @@
+PY3_PROGRAM()
+
+PY_SRCS(__main__.py)
+
+PEERDIR(
+    contrib/python/MarkupSafe
+    contrib/python/Jinja2
+    ydb/core/protos
+)
+
+END()

ydb/core/base/generated/runtime_feature_flags.cpp.in

@@ -0,0 +1,127 @@
+// Auto-generated by {{ generator }}, do not modify.
+#include <ydb/core/base/generated/runtime_feature_flags.h>
+#include <ydb/core/protos/feature_flags.pb.h>
+
+namespace NKikimr {
+
+{% for slot in slots %}
+std::tuple<ui64, ui64> TRuntimeFeatureFlags::BitsFromProto_{{ slot.name }}(const NKikimrConfig::TFeatureFlags& flags) {
+    ui64 bits = 0;
+    ui64 mask = 0;
+{%- for field in slot.fields %}
+    if (flags.Has{{ field.name }}()) {
+        bits |= flags.Get{{ field.name }}() ? {{ field.full_mask }}ULL : {{ field.has_mask }}ULL;
+        mask |= {{ field.full_mask }}ULL;
+    }
+{%- endfor %}
+    return { bits, mask };
+}
+
+ui64 TRuntimeFeatureFlags::FromProto_{{ slot.name }}(const NKikimrConfig::TFeatureFlags& flags, ui64 value) {
+    auto [bits, mask] = BitsFromProto_{{ slot.name }}(flags);
+    return (value & ~mask) | bits;
+}
+
+void TRuntimeFeatureFlags::Update_{{ slot.name }}(ui64 bits, ui64 mask) {
+    ui64 slot = {{ slot.name }}_.load(std::memory_order_relaxed);
+    for (;;) {
+        ui64 updated = (slot & ~mask) | bits;
+        // We perform a CAS so unrelated updates by concurrent threads are not lost
+        if ({{ slot.name }}_.compare_exchange_weak(slot, updated, std::memory_order_relaxed)) {
+            break;
+        }
+    }
+}
+{% endfor %}
+
+TRuntimeFeatureFlags::TRuntimeFeatureFlags()
+{%- for slot in slots %}
+    {{ ',' if slot.index else ':' }} {{ slot.name }}_{ {{ slot.default_value }}ULL }
+{%- endfor %}
+{}
+
+TRuntimeFeatureFlags::TRuntimeFeatureFlags(const TRuntimeFeatureFlags& flags)
+{%- for slot in slots %}
+    {{ ',' if slot.index else ':' }} {{ slot.name }}_{ flags.{{ slot.name }}_.load(std::memory_order_relaxed) }
+{%- endfor %}
+{}
+
+TRuntimeFeatureFlags& TRuntimeFeatureFlags::operator=(const TRuntimeFeatureFlags& flags) {
+{%- for slot in slots %}
+    {{ slot.name }}_.store(
+        flags.{{ slot.name }}_.load(std::memory_order_relaxed),
+        std::memory_order_relaxed);
+{%- endfor %}
+    return *this;
+}
+
+TRuntimeFeatureFlags::TRuntimeFeatureFlags(const NKikimrConfig::TFeatureFlags& flags)
+{%- for slot in slots %}
+    {{ ',' if slot.index else ':' }} {{ slot.name }}_{ FromProto_{{ slot.name }}(flags, {{ slot.default_value }}ULL) }
+{%- endfor %}
+{}
+
+TRuntimeFeatureFlags& TRuntimeFeatureFlags::operator=(const NKikimrConfig::TFeatureFlags& flags) {
+    CopyFrom(flags);
+    return *this;
+}
+
+void TRuntimeFeatureFlags::CopyFrom(const NKikimrConfig::TFeatureFlags& flags) {
+{%- for slot in slots %}
+    ui64 {{ slot.name }} = FromProto_{{ slot.name }}(flags, {{ slot.default_value }}ULL);
+{%- endfor %}
+{%- for slot in slots %}
+    {{ slot.name }}_.store({{ slot.name }}, std::memory_order_relaxed);
+{%- endfor %}
+}
+
+void TRuntimeFeatureFlags::MergeFrom(const NKikimrConfig::TFeatureFlags& flags) {
+{%- for slot in slots %}
+    auto [{{ slot.name }}_bits, {{ slot.name }}_mask] = BitsFromProto_{{ slot.name }}(flags);
+{%- endfor %}
+{%- for slot in slots %}
+    Update_{{ slot.name }}({{ slot.name }}_bits, {{ slot.name }}_mask);
+{%- endfor %}
+}
+
+void TRuntimeFeatureFlags::CopyRuntimeFrom(const NKikimrConfig::TFeatureFlags& flags) {
+{%- for slot in slots %}
+    ui64 {{ slot.name }} = FromProto_{{ slot.name }}(flags, {{ slot.default_value }}ULL);
+{%- endfor %}
+{%- for slot in slots %}
+    Update_{{ slot.name }}({{ slot.name }} & {{ slot.runtime_flags_mask }}ULL, {{ slot.runtime_flags_mask }}ULL);
+{%- endfor %}
+}
+
+TRuntimeFeatureFlags::operator NKikimrConfig::TFeatureFlags() const {
+    NKikimrConfig::TFeatureFlags flags;
+{%- for slot in slots %}
+    ui64 {{ slot.name }} = {{ slot.name }}_.load(std::memory_order_relaxed);
+{%- endfor %}
+{%- for field in fields %}
+    if ({{ field.slot.name }} & {{ field.has_mask }}ULL) {
+        flags.Set{{ field.name }}(bool({{ field.slot.name }} & {{ field.value_mask }}ULL));
+    }
+{%- endfor %}
+    return flags;
+}
+
+{% for field in fields %}
+bool TRuntimeFeatureFlags::Has{{ field.name }}() const {
+    return {{ field.slot.name }}_.load(std::memory_order_relaxed) & {{ field.has_mask }}ULL;
+}
+
+bool TRuntimeFeatureFlags::Get{{ field.name }}() const {
+    return {{ field.slot.name }}_.load(std::memory_order_relaxed) & {{ field.value_mask }}ULL;
+}
+
+void TRuntimeFeatureFlags::Set{{ field.name }}(bool value) {
+    Update_{{ field.slot.name }}(value ? {{ field.full_mask }}ULL : {{ field.has_mask }}ULL, {{ field.full_mask}}ULL);
+}
+
+void TRuntimeFeatureFlags::Clear{{ field.name }}() {
+    Update_{{ field.slot.name }}({{ field.default_value }}ULL, {{ field.full_mask }}ULL);
+}
+{% endfor %}
+
+} // namespace NKikimr