Try to preserve in-memory state (locks in particular) across restarts

Author: snaury

ydb/core/protos/counters_datashard.proto

@@ -527,4 +527,5 @@ enum ETxTypes {
  TXTYPE_PLAN_PREDICTED_TXS = 81 [(TxTypeOpts) = {Name: "TxPlanPredictedTxs"}];
  TXTYPE_WRITE = 82 [(TxTypeOpts) = {Name: "TxWrite"}];
  TXTYPE_REMOVE_SCHEMA_SNAPSHOTS = 83 [(TxTypeOpts) = {Name: "TxRemoveSchemaSnapshots"}];
+ TXTYPE_INIT_RESTORED = 84 [(TxTypeOpts) = {Name: "TxInitRestored"}];
 }

ydb/core/protos/datashard_config.proto

@@ -25,4 +25,5 @@ message TDataShardConfig {
  optional bool DisabledOnSchemeShard = 19 [default = false];
  optional uint64 IncrementalRestoreReadAheadLo = 20 [default = 524288];
  optional uint64 IncrementalRestoreReadAheadHi = 21 [default = 1048576];
+ optional uint64 InMemoryStateMigrationTimeoutMs = 24 [default = 1000];
 }

ydb/core/protos/feature_flags.proto

@@ -180,4 +180,6 @@ message TFeatureFlags {
  optional bool EnableReleaseNodeNameOnGracefulShutdown = 155 [default = false];
  optional bool ForceDistconfDisable = 156 [default = false];
  optional bool EnableScaleRecommender = 157 [default = false];
+ optional bool EnableDataShardInMemoryStateMigration = 158 [default = true];
+ optional bool EnableDataShardInMemoryStateMigrationAcrossGenerations = 159 [default = false];
 }

ydb/core/protos/tx_datashard.proto

@@ -2144,3 +2144,118 @@ message TSerializedEvent {
  // TEventSerializationInfo::IsExtendedFormat flag
  optional bool IsExtendedFormat = 2;
 }
+
+// In-memory variables that are important to restore between shard generations
+message TInMemoryVars {
+ optional NKikimrProto.TRowVersion ImmediateWriteEdge = 1;
+ optional NKikimrProto.TRowVersion ImmediateWriteEdgeReplied = 2;
+ optional NKikimrProto.TRowVersion UnprotectedReadEdge = 3;
+}
+
+// In-memory locks
+message TInMemoryLockInfo {
+ optional uint64 LockId = 1;
+ optional uint32 LockNodeId = 2;
+ optional uint32 Generation = 3;
+ optional uint32 Counter = 4;
+ optional uint64 CreateTs = 5;
+ optional uint64 Flags = 6;
+ optional NKikimrProto.TRowVersion BreakVersion = 7;
+ repeated NKikimrProto.TPathID ReadTables = 8;
+ repeated NKikimrProto.TPathID WriteTables = 9;
+}
+
+// In-memory ranges associated with the lock
+// Note: whole shard locks will have whole table read/write ranges
+message TInMemoryLockRange {
+ optional uint64 LockId = 1;
+ optional NKikimrProto.TPathID TableId = 2;
+ optional uint64 Flags = 3;
+ optional bytes Data = 4;
+}
+
+// In-memory conflicts between locks
+message TInMemoryLockConflict {
+ // When this lock is committed
+ optional uint64 LockId = 1;
+ // It must break another lock with the specified id
+ optional uint64 ConflictId = 2;
+}
+
+// In-memory volatile dependencies
+// This specifies that to commit this lock we must use a volatile transaction
+// that must have a dependency on the specified TxId.
+message TInMemoryLockVolatileDependency {
+ optional uint64 LockId = 1;
+ optional uint64 TxId = 2;
+}
+
+// In-memory prepared volatile transactions, i.e. those we have accepted for
+// execution and may be planned and executed later. This includes planned and
+// acknowledged transactions that haven't executed yet (including those that
+// have not been committed yet?)
+message TInMemoryPreparedVolatileTx {
+ optional uint64 TxId = 1;
+ optional NActorsProto.TActorId Source = 2;
+ optional uint64 Cookie = 3;
+ optional uint64 Type = 4;
+ optional bytes Body = 5;
+ optional uint64 Flags = 6;
+ // This is Min/Max step of a prepared transactions, mostly used for deadline cleanup
+ optional uint64 MinStep = 7;
+ optional uint64 MaxStep = 8;
+ // Specified when the transaction has been planned already (including predictions)
+ optional uint64 PredictedStep = 9;
+ optional uint64 Step = 10;
+}
+
+// In-memory waiting volatile transactions, i.e. those that have not been
+// decided yet (at least persistently), and for which we need to send a
+// result when they eventually commit or abort. This is matched against
+// transactions that are restored from disk.
+message TInMemoryWaitingVolatileTx {
+ optional uint64 TxId = 1;
+ optional NActorsProto.TActorId Source = 2;
+ optional uint64 Cookie = 3;
+ optional uint64 Type = 4;
+}
+
+// Partial in-memory state of the datashard
+message TInMemoryState {
+ optional TInMemoryVars Vars = 1;
+ repeated TInMemoryLockInfo Locks = 2;
+ repeated TInMemoryLockRange LockRanges = 3;
+ repeated TInMemoryLockConflict LockConflicts = 4;
+ repeated TInMemoryLockVolatileDependency LockVolatileDependencies = 5;
+ repeated TInMemoryPreparedVolatileTx PreparedVolatileTxs = 6;
+ repeated TInMemoryWaitingVolatileTx WaitingVolatileTxs = 7;
+}
+
+// Internal continuation token, may change between versions
+message TInMemoryStateContinuationToken {
+ // Index to continue from
+ optional uint64 NextIndex = 1;
+}
+
+// Sent to the state actor by new generations
+message TEvInMemoryStateRequest {
+ // Generation of the new tablet
+ optional uint32 Generation = 1;
+
+ // Continuation token to get the next response
+ optional bytes ContinuationToken = 2;
+}
+
+// Sent by the state actor to new generations
+message TEvInMemoryStateResponse {
+ // Partial serialized TInMemoryState chunk
+ optional uint32 SerializedStatePayloadIndex = 1;
+
+ // Optional checkpoint offsets within the chunk. The protobuf message may
+ // be parsed cleanly up to these offsets without the need to concatenate
+ // subsequent chunks.
+ repeated uint32 SerializedStateCheckpoints = 2 [packed = true];
+
+ // Will be set for incomplete replies
+ optional bytes ContinuationToken = 3;
+}

ydb/core/tx/datashard/datashard.cpp

@@ -1,6 +1,7 @@
 #include "datashard_impl.h"
 #include "datashard_txs.h"
 #include "datashard_locks_db.h"
+#include "memory_state_migration.h"
 #include "probes.h"
 
 #include <ydb/core/base/interconnect_channels.h>
@@ -174,6 +175,15 @@ TDataShard::TDataShard(const TActorId &tablet, TTabletStorageInfo *info)
  RegisterDataShardProbes();
 }
 
+TDataShard::~TDataShard() {
+ if (InMemoryRestoreActor) {
+ InMemoryRestoreActor->OnTabletDestroyed();
+ }
+ if (InMemoryStateActor) {
+ InMemoryStateActor->OnTabletDestroyed();
+ }
+}
+
 void TDataShard::OnDetach(const TActorContext &ctx) {
  Cleanup(ctx);
  LOG_INFO_S(ctx, NKikimrServices::TX_DATASHARD, "OnDetach: " << TabletID());
@@ -265,6 +275,13 @@ void TDataShard::Cleanup(const TActorContext& ctx) {
 }
 
 void TDataShard::Die(const TActorContext& ctx) {
+ if (InMemoryRestoreActor) {
+ InMemoryRestoreActor->OnTabletDead();
+ }
+ if (InMemoryStateActor) {
+ InMemoryStateActor->OnTabletDead();
+ }
+
  NTabletPipe::CloseAndForgetClient(SelfId(), SchemeShardPipe);
  NTabletPipe::CloseAndForgetClient(SelfId(), StateReportPipe);
  NTabletPipe::CloseAndForgetClient(SelfId(), DbStatsReportPipe);
@@ -1620,6 +1637,29 @@ void TDataShard::PersistSys(NIceDb::TNiceDb& db, ui64 key, bool value) const {
  db.Table<Schema::Sys>().Key(key).Update(NIceDb::TUpdate<Schema::Sys::Uint64>(value ? 1 : 0));
 }
 
+void TDataShard::PersistSys(NIceDb::TNiceDb& db, ui64 key, const TActorId& value) const {
+ char buf[sizeof(ui64) * 2];
+ WriteUnaligned<ui64>(buf, value.RawX1());
+ WriteUnaligned<ui64>(buf + sizeof(ui64), value.RawX2());
+ db.Table<Schema::Sys>().Key(key).Update(NIceDb::TUpdate<Schema::Sys::Bytes>(TString(buf, sizeof(ui64) * 2)));
+}
+
+bool TDataShard::SysGetActorId(NIceDb::TNiceDb& db, ui64 key, TActorId& value) {
+ auto rowset = db.Table<Schema::Sys>().Key(key).Select<Schema::Sys::Bytes>();
+ if (!rowset.IsReady()) {
+ return false;
+ }
+ if (rowset.IsValid()) {
+ TString buf = rowset.GetValue<Schema::Sys::Bytes>();
+ Y_ABORT_UNLESS(buf.size() == sizeof(ui64) * 2, "Unexpected TActorId value size %" PRISZT, buf.size());
+ const char* data = buf.data();
+ ui64 x1 = ReadUnaligned<ui64>(data);
+ ui64 x2 = ReadUnaligned<ui64>(data + sizeof(ui64));
+ value = TActorId(x1, x2);
+ }
+ return true;
+}
+
 void TDataShard::PersistUserTable(NIceDb::TNiceDb& db, ui64 tableId, const TUserTable& tableInfo) {
  db.Table<Schema::UserTables>().Key(tableId).Update(
  NIceDb::TUpdate<Schema::UserTables::LocalTid>(tableInfo.LocalTid),
@@ -2484,11 +2524,18 @@ void TDataShard::SendImmediateWriteResult(
  const ui64 step = version.Step;
  const ui64 observedStep = GetMaxObservedStep();
  if (step <= observedStep) {
- SnapshotManager.PromoteImmediateWriteEdgeReplied(version);
- if (!sessionId) {
- Send(target, event, 0, cookie, span.GetTraceId());
+ // We avoid sending replies that would have promoted the replied edge
+ // when it's frozen. This prevents us replying and causing the next
+ // generation to potentially keep reading stale data.
+ if (Y_LIKELY(!InMemoryVarsFrozen) || version <= SnapshotManager.GetImmediateWriteEdgeReplied()) {
+ SnapshotManager.PromoteImmediateWriteEdgeReplied(version);
+ if (!sessionId) {
+ Send(target, event, 0, cookie, span.GetTraceId());
+ } else {
+ SendViaSession(sessionId, target, SelfId(), event, 0, cookie, span.GetTraceId());
+ }
  } else {
- SendViaSession(sessionId, target, SelfId(), event, 0, cookie, span.GetTraceId());
+ span.EndError("Dropped");
  }
  return;
  }
@@ -2627,13 +2674,20 @@ void TDataShard::SendAfterMediatorStepActivate(ui64 mediatorStep, const TActorCo
  }
 
  if (step <= mediatorStep) {
- if (it->second.Span) {
- it->second.Span.Attribute("ActivateStep", std::to_string(mediatorStep));
+ // We avoid sending replies that would have promoted the replied edge
+ // when it's frozen. This prevents us replying and causing the next
+ // generation to potentially keep reading stale data.
+ if (Y_LIKELY(!InMemoryVarsFrozen) || it->first <= SnapshotManager.GetImmediateWriteEdgeReplied()) {
+ if (it->second.Span) {
+ it->second.Span.Attribute("ActivateStep", std::to_string(mediatorStep));
+ }
+ SnapshotManager.PromoteImmediateWriteEdgeReplied(it->first);
+ SendViaSession(it->second.SessionId,
+ it->second.Target, SelfId(), it->second.Event.Release(),
+ 0, it->second.Cookie, it->second.Span.GetTraceId());
+ } else {
+ it->second.Span.EndError("Dropped");
  }
- SnapshotManager.PromoteImmediateWriteEdgeReplied(it->first);
- SendViaSession(it->second.SessionId,
- it->second.Target, SelfId(), it->second.Event.Release(),
- 0, it->second.Cookie, it->second.Span.GetTraceId());
  it = MediatorDelayedReplies.erase(it);
  continue;
  }
@@ -2693,6 +2747,12 @@ bool TDataShard::NeedMediatorStateRestored() const {
  return false;
  }
 
+ if (InMemoryVarsRestored) {
+ // We have migrated in-memory vars from previous generations
+ // We don't need mediator state for correctness
+ return false;
+ }
+
  switch (State) {
  case TShardState::Ready:
  case TShardState::Readonly:

ydb/core/tx/datashard/datashard.h

@@ -342,6 +342,10 @@ namespace TEvDataShard {
  EvReadScanStarted,
  EvReadScanFinished,
 
+ // Used to transfer in-memory state between generations
+ EvInMemoryStateRequest,
+ EvInMemoryStateResponse,
+
  EvEnd
  };
 
@@ -1769,6 +1773,29 @@ namespace TEvDataShard {
  Record.SetErrorDescription(error);
  }
  };
+
+ struct TEvInMemoryStateRequest
+ : public TEventPB<TEvInMemoryStateRequest,
+ NKikimrTxDataShard::TEvInMemoryStateRequest,
+ EvInMemoryStateRequest>
+ {
+ TEvInMemoryStateRequest() = default;
+
+ explicit TEvInMemoryStateRequest(ui32 generation, const TString& continuationToken = {}) {
+ Record.SetGeneration(generation);
+ if (!continuationToken.empty()) {
+ Record.SetContinuationToken(continuationToken);
+ }
+ }
+ };
+
+ struct TEvInMemoryStateResponse
+ : public TEventPB<TEvInMemoryStateResponse,
+ NKikimrTxDataShard::TEvInMemoryStateResponse,
+ EvInMemoryStateResponse>
+ {
+ TEvInMemoryStateResponse() = default;
+ };
 };
 
 IActor* CreateDataShard(const TActorId &tablet, TTabletStorageInfo *info);