fix(meta): fence datanode processes replaced by a newer process (#9878)

* fix(meta): fence datanode processes replaced by a newer process

Two datanode processes can run with the same node id, and metasrv
grants leases and delivers instructions by node id only. Every datanode
heartbeat already carries node_epoch, the time the process started.

The metasrv leader records the highest accepted epoch per node id under
__datanode_epoch/{node_id}, a prefix of the leader cache, and raises it
with compare-and-put. The heartbeat session checks each datanode request
with a non-zero epoch before the handler chain: a lower epoch gets a
reply without running the chain, so no lease, keep-alive, stats or
mailbox reply. An epoch more than 5 minutes ahead of the leader's clock
is not recorded and is fenced too.

Datanode pushers carry the epoch of their process: a lower epoch cannot
replace a higher one, and a stream removes only its own pusher.

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* fix(meta): keep the leader cache load and pusher updates coherent

A write made while the leader cache is disabled, as during the load at
leader start, went to the store without the cache noticing, so the load
could install a value read before the write. Such writes now create new
cache versions around the store call, and the load drops what it loaded
for a prefix if a write happened meanwhile. This also covers node
addresses, which heartbeats write during the load.

A stream's pusher that takes its first epoch is updated in place instead
of being replaced, so messages waiting on it are not cancelled, and the
connection gauge counts only a pusher that fills an empty slot.

The epoch check moves into the heartbeat service.

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* fix(meta): keep the leader cache coherent with writes while it is disabled

A write that started while the cache was disabled could commit after the
load had read the old value and the cache had been enabled, leaving the
old value cached. Writes now maintain the cache whether or not it is
enabled, as enabled writes already did: new versions around the store
call, then the key is updated or invalidated.

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* test(meta): pin the epoch 0 exemption and pusher order, and update docs

The datanode test without an epoch now runs after another process has
recorded one, and the pusher test checks that a stream's first accepted
request receives what its handlers broadcast. The leader cache load no
longer claims that writes must not run while it loads, and the meta-srv
guide states that epoch 0 is not checked.

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* fix(meta): register a datanode pusher only after its first request passes

A datanode stream registered an epoch-0 pusher on its handshake, so an
older process's reconnecting stream could take an empty slot, or a newer
process's pusher before its first heartbeat, for about a round trip. A
datanode stream now registers its pusher on its first request that
passes the epoch check (epoch-0 requests after the handshake included),
still before the handlers run.

Pusher registration is logged only when the insert happens. Tests cover
a fenced stream removing its own pusher, a handshake not taking an empty
slot, and a KV error failing the request; two steps that could not fail
are removed.

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* fix(meta): check a datanode epoch again before replying

A request could pass the epoch check, wait in the handler chain while a
newer process of the node was recorded, and still return a region lease.
After the handlers run, the session checks the epoch again against the
cached record and, if a newer process has been recorded, replies as to a
fenced request and removes the stream's own pusher. This costs one more
read of the cached record per checked datanode request.

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* fix(meta): handle a datanode's requests one at a time

A request of an older process could pass the epoch check, wait in the
handlers while a newer process was admitted, and then keep the node
alive with its own address. Checking the epoch again after the handlers
only dropped the reply.

The heartbeat session now takes an in-memory lock for the node id
before checking a request that carries an epoch, and releases it after
the handlers, before the reply is sent. A newer process is therefore
admitted only after the handlers of an older request finish. The check
after the handlers is removed.

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

---------

Signed-off-by: jeremyhi <fengjiachun@gmail.com>
This commit is contained in:
jeremyhi
2026-10-26 09:02:40 +00:00
committed by GitHub
parent c96cdd0026
commit 6f5c064f54
7 changed files with 841 additions and 71 deletions
+10 -2
View File
@@ -65,6 +65,9 @@
//! 13. Topic name to region map key `__topic_region/{topic_name}/{region_id}`
//! - Mapping {topic_name} to {region_id}
//!
//! 14. Datanode epoch key: `__datanode_epoch/{datanode_id}`
//! - The highest `node_epoch` the metasrv leader has accepted from the datanode
//!
//! All keys have related managers. The managers take care of the serialization and deserialization
//! of keys and values, and the interaction with the underlying KV store backend.
//!
@@ -98,6 +101,7 @@
//! {partition_id}
pub mod catalog_name;
pub mod datanode_epoch;
pub mod datanode_table;
pub mod flow;
pub mod node_address;
@@ -149,6 +153,7 @@ use topic_region::{TopicRegionKey, TopicRegionManager};
use view_info::{ViewInfoKey, ViewInfoManager, ViewInfoValue};
use self::catalog_name::{CatalogManager, CatalogNameKey, CatalogNameValue};
use self::datanode_epoch::DatanodeEpochValue;
use self::datanode_table::RegionInfo;
use self::flow::flow_info::FlowInfoValue;
use self::flow::flow_name::FlowNameValue;
@@ -179,6 +184,7 @@ pub const PAUSE_PROCEDURE_KEY: &str = "__switches/pause_procedure";
pub const RECOVERY_MODE_KEY: &str = "__switches/recovery";
pub const DATANODE_TABLE_KEY_PREFIX: &str = "__dn_table";
pub const DATANODE_EPOCH_KEY_PREFIX: &str = "__datanode_epoch";
pub const TABLE_INFO_KEY_PREFIX: &str = "__table_info";
pub const VIEW_INFO_KEY_PREFIX: &str = "__view_info";
pub const TABLE_NAME_KEY_PREFIX: &str = "__table_name";
@@ -198,12 +204,13 @@ pub const ELECTION_KEY: &str = "__metasrv_election";
pub const CANDIDATES_ROOT: &str = "__metasrv_election_candidates/";
/// The keys with these prefixes will be loaded into the cache when the leader starts.
pub const CACHE_KEY_PREFIXES: [&str; 5] = [
pub const CACHE_KEY_PREFIXES: [&str; 6] = [
TABLE_NAME_KEY_PREFIX,
CATALOG_NAME_KEY_PREFIX,
SCHEMA_NAME_KEY_PREFIX,
TABLE_ROUTE_PREFIX,
NODE_ADDRESS_PREFIX,
DATANODE_EPOCH_KEY_PREFIX,
];
/// A set of regions with the same role.
@@ -1886,7 +1893,8 @@ impl_metadata_value! {
SchemaNameValue,
FlowStateValue,
PoisonValue,
TopicRegionValue
TopicRegionValue,
DatanodeEpochValue
}
impl_optional_metadata_value! {
+66
View File
@@ -0,0 +1,66 @@
// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use std::fmt::Display;
use serde::{Deserialize, Serialize};
use crate::DatanodeId;
use crate::key::DATANODE_EPOCH_KEY_PREFIX;
/// The key of the highest `node_epoch` the metasrv leader has accepted from a datanode.
///
/// The layout: `__datanode_epoch/{node_id}`
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct DatanodeEpochKey {
pub node_id: DatanodeId,
}
impl DatanodeEpochKey {
pub fn new(node_id: DatanodeId) -> Self {
Self { node_id }
}
pub fn to_bytes(&self) -> Vec<u8> {
self.to_string().into_bytes()
}
}
impl Display for DatanodeEpochKey {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}/{}", DATANODE_EPOCH_KEY_PREFIX, self.node_id)
}
}
/// The value of [DatanodeEpochKey]. It only grows.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct DatanodeEpochValue {
pub node_epoch: u64,
}
#[cfg(test)]
mod tests {
use super::*;
use crate::key::MetadataValue;
#[test]
fn test_frozen_datanode_epoch_record() {
let value =
include_bytes!("../../../../../tests/compatibility/meta/datanode_epoch/value.json");
assert_eq!(DatanodeEpochKey::new(7).to_bytes(), b"__datanode_epoch/7");
let decoded = DatanodeEpochValue::try_from_raw_value(value).unwrap();
assert_eq!(decoded.node_epoch, 1761091200000);
assert_eq!(decoded.try_as_raw_value().unwrap(), value);
}
}
+8
View File
@@ -43,6 +43,14 @@ abstraction, election, key encoding, and the DDL manager live in `common-meta`;
`handler/utils.rs` for per-node address persistence. Datanode address
updates retain the statistics handler's per-node ordering. Address epochs are
confirmed only after a successful save; cache invalidation remains best effort.
- **Datanode process fence**: before the handler chain, the leader's heartbeat
session checks a datanode request's non-zero `node_epoch` against the highest
one recorded under `__datanode_epoch`, read and raised through the leader
cache. A lower epoch gets a reply without running the chain; requests that
carry an epoch take a per-node lock through the check and the chain; epoch 0
(handshakes, older datanodes) is not checked. A datanode stream registers
its pusher on its first request that passes, not on the handshake; pushers
carry the epoch, and a lower one cannot replace a higher one.
- **DDL**: `service/procedure.rs` (`ddl`) is leader-only and hands off to
`common-meta`'s `DdlManager`, executed via `common-procedure`.
- **Region migration**: `procedure/region_migration/manager.rs` plus per-step
+84 -11
View File
@@ -51,7 +51,7 @@ use response_header_handler::ResponseHeaderHandler;
use snafu::{OptionExt, ResultExt};
use store_api::storage::RegionId;
use tokio::sync::mpsc::Sender;
use tokio::sync::{Notify, RwLock, oneshot, watch};
use tokio::sync::{Notify, OwnedMutexGuard, RwLock, oneshot, watch};
use crate::error::{self, DeserializeFromJsonSnafu, Result, UnexpectedInstructionReplySnafu};
use crate::handler::collect_topic_stats_handler::CollectTopicStatsHandler;
@@ -198,6 +198,8 @@ pub struct Pusher {
deregister_signal_receiver: DeregisterSignalReceiver,
res_header: ResponseHeader,
/// The `node_epoch` of the datanode process that registered the pusher, or 0.
node_epoch: u64,
}
impl Drop for Pusher {
@@ -220,9 +222,15 @@ impl Pusher {
deregister_signal_sender,
deregister_signal_receiver,
res_header,
node_epoch: 0,
}
}
pub(crate) fn with_node_epoch(mut self, node_epoch: u64) -> Self {
self.node_epoch = node_epoch;
self
}
#[inline]
pub fn header(&self) -> ResponseHeader {
self.res_header.clone()
@@ -327,12 +335,38 @@ impl Pushers {
Ok(())
}
pub(crate) async fn insert(&self, pusher_id: PusherId, pusher: Pusher) -> Option<Pusher> {
self.0.write().await.insert(pusher_id, pusher)
/// Inserts the pusher unless the slot holds one with a higher node epoch, so an older
/// datanode process cannot take over a newer one's pusher. A pusher of the same stream
/// only takes the new epoch, keeping the messages waiting on it. Returns `None` if refused,
/// otherwise whether the slot was empty.
pub(crate) async fn insert(&self, pusher_id: PusherId, pusher: Pusher) -> Option<bool> {
let mut pushers = self.0.write().await;
match pushers.get_mut(&pusher_id) {
None => {
pushers.insert(pusher_id, pusher);
return Some(true);
}
Some(current) if current.node_epoch > pusher.node_epoch => return None,
Some(current) if current.sender.same_channel(&pusher.sender) => {
current.node_epoch = pusher.node_epoch;
}
Some(current) => *current = pusher,
}
Some(false)
}
async fn remove(&self, pusher_id: PusherId) -> Option<Pusher> {
self.0.write().await.remove(&pusher_id)
/// Removes the pusher if it pushes to `sender`'s channel, keeping one that another stream
/// has registered since.
async fn remove(
&self,
pusher_id: PusherId,
sender: &Sender<std::result::Result<HeartbeatResponse, tonic::Status>>,
) -> bool {
let mut pushers = self.0.write().await;
let own = pushers
.get(&pusher_id)
.is_some_and(|pusher| pusher.sender.same_channel(sender));
own && pushers.remove(&pusher_id).is_some()
}
}
@@ -357,24 +391,43 @@ pub type HeartbeatHandlerGroupRef = Arc<HeartbeatHandlerGroup>;
pub struct HeartbeatHandlerGroup {
handlers: Vec<NameCachedHandler>,
pushers: Pushers,
/// One lock per datanode id seen, kept for the life of the process.
datanode_locks: DashMap<u64, Arc<tokio::sync::Mutex<()>>>,
}
impl HeartbeatHandlerGroup {
/// Registers the heartbeat response [`Pusher`] with the given key to the group.
/// Registers the heartbeat response [`Pusher`] with the given key to the group, unless the
/// group holds one with a higher node epoch.
pub async fn register_pusher(&self, pusher_id: PusherId, pusher: Pusher) {
METRIC_META_HEARTBEAT_CONNECTION_NUM.inc();
let Some(vacant) = self.pushers.insert(pusher_id, pusher).await else {
return;
};
info!("Pusher register: {}", pusher_id);
let _ = self.pushers.insert(pusher_id, pusher).await;
if vacant {
METRIC_META_HEARTBEAT_CONNECTION_NUM.inc();
}
}
/// Deregisters the heartbeat response [`Pusher`] with the given key from the group.
pub async fn deregister_push(&self, pusher_id: PusherId) {
if self.pushers.remove(pusher_id).await.is_some() {
/// Deregisters the heartbeat response [`Pusher`] with the given key from the group if
/// `sender` registered it.
pub async fn deregister_push(
&self,
pusher_id: PusherId,
sender: &Sender<std::result::Result<HeartbeatResponse, tonic::Status>>,
) {
if self.pushers.remove(pusher_id, sender).await {
info!("Pusher unregister: {}", pusher_id);
METRIC_META_HEARTBEAT_CONNECTION_NUM.dec();
}
}
/// Locks the datanode with the given id until the guard is dropped, so its requests are
/// checked and handled one at a time.
pub(crate) async fn lock_datanode(&self, node_id: u64) -> OwnedMutexGuard<()> {
let lock = self.datanode_locks.entry(node_id).or_default().clone();
lock.lock_owned().await
}
#[cfg(test)]
/// Returns whether the group contains the heartbeat response [`Pusher`] with the given key.
pub async fn contains_pusher(&self, pusher_id: &PusherId) -> bool {
@@ -783,6 +836,7 @@ impl HeartbeatHandlerGroupBuilder {
Ok(HeartbeatHandlerGroup {
handlers: self.handlers,
pushers: self.pushers,
datanode_locks: DashMap::new(),
})
}
@@ -1477,6 +1531,25 @@ mod tests {
assert!(healthy_rx.try_recv().is_ok());
}
#[tokio::test]
async fn test_pusher_insert_reports_an_empty_slot() {
let pushers = Pushers::default();
let id = PusherId::new(Role::Datanode, 1);
let (tx1, _rx1) = mpsc::channel(1);
let (tx2, _rx2) = mpsc::channel(1);
assert_eq!(
pushers.insert(id, Pusher::new(tx1.clone())).await,
Some(true)
);
// Taking an epoch and replacing another stream's pusher fill no slot.
let in_place = Pusher::new(tx1).with_node_epoch(2);
assert_eq!(pushers.insert(id, in_place).await, Some(false));
let lower = Pusher::new(tx2.clone()).with_node_epoch(1);
assert_eq!(pushers.insert(id, lower).await, None);
let other = Pusher::new(tx2).with_node_epoch(2);
assert_eq!(pushers.insert(id, other).await, Some(false));
}
#[tokio::test]
async fn test_pushers_broadcast_does_not_fence_the_registry() {
let pushers = Pushers::default();
+604 -29
View File
@@ -12,16 +12,23 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::cmp::Ordering;
use std::io::ErrorKind;
use std::sync::Arc;
use std::sync::atomic::AtomicU64;
use std::time::Duration;
use api::v1::meta::{
AskLeaderRequest, AskLeaderResponse, HeartbeatRequest, HeartbeatResponse, Peer, RequestHeader,
ResponseHeader, Role, heartbeat_server,
AskLeaderRequest, AskLeaderResponse, Error, HeartbeatRequest, HeartbeatResponse,
PROTOCOL_VERSION, Peer, RequestHeader, ResponseHeader, Role, heartbeat_server,
};
use common_meta::election::LeaderChangeMessage;
use common_meta::key::MetadataValue;
use common_meta::key::datanode_epoch::{DatanodeEpochKey, DatanodeEpochValue};
use common_meta::kv_backend::KvBackend;
use common_meta::rpc::store::CompareAndPutRequest;
use common_telemetry::{debug, error, info, warn};
use common_time::util::current_time_millis;
use futures::StreamExt;
use once_cell::sync::OnceCell;
use snafu::{OptionExt, ResultExt};
@@ -39,6 +46,9 @@ use crate::service::{GrpcResult, GrpcStream};
type HeartbeatResponseResult = std::result::Result<HeartbeatResponse, Status>;
/// How far ahead of the leader's clock a datanode epoch can be recorded.
const MAX_NODE_EPOCH_AHEAD: Duration = Duration::from_secs(5 * 60);
#[async_trait::async_trait]
trait HeartbeatRequestStream {
async fn next(&mut self) -> Option<std::result::Result<HeartbeatRequest, Status>>;
@@ -109,6 +119,8 @@ struct HeartbeatSession<R, L> {
handler_group: Arc<HeartbeatHandlerGroup>,
ctx: Context,
sender_id: PusherId,
/// The node epoch the stream's pusher was registered with, if any.
pusher_epoch: Option<u64>,
}
impl<R, L> HeartbeatSession<R, L>
@@ -144,7 +156,12 @@ where
return None;
};
let sender_id = register_pusher(&handler_group, header, tx.clone()).await;
let sender_id = PusherId::new(header.role(), get_node_id(header));
// A datanode stream registers its pusher on its first request that passes the check.
if sender_id.role != Role::Datanode {
let pusher = Pusher::new(tx.clone());
handler_group.register_pusher(sender_id, pusher).await;
}
let mut session = Self {
requests,
tx,
@@ -152,6 +169,7 @@ where
handler_group,
ctx,
sender_id,
pusher_epoch: None,
};
if session.handle_request(req, true).await {
@@ -219,10 +237,23 @@ where
.with_label_values(&[sender_id.as_str()])
.inc();
let res = self
.handler_group
.handle(req, self.ctx.clone().with_handshake(is_handshake))
.await
// A datanode's requests that carry an epoch are checked and handled one at a time, so
// an older process's request cannot run its handlers after a newer process is admitted.
let datanode_lock = if self.sender_id.role == Role::Datanode && req.node_epoch != 0 {
Some(self.handler_group.lock_datanode(self.sender_id.id).await)
} else {
None
};
let res = match self.check_node_epoch(&req, is_handshake).await {
Some(res) => res,
None => {
self.handler_group
.handle(req, self.ctx.clone().with_handshake(is_handshake))
.await
}
};
drop(datanode_lock);
let res = res
.inspect_err(
|e| warn!(e; "Failed to handle heartbeat request, sender: {}", self.sender_id),
)
@@ -257,6 +288,48 @@ where
true
}
/// Checks the node epoch of a datanode request, and registers the stream's pusher once
/// the request passes. Returns the response to send instead of running the handlers, if any.
async fn check_node_epoch(
&mut self,
req: &HeartbeatRequest,
is_handshake: bool,
) -> Option<Result<HeartbeatResponse>> {
let node_epoch = req.node_epoch;
if self.sender_id.role != Role::Datanode || (node_epoch == 0 && is_handshake) {
return None;
}
// Datanodes that do not report their epoch are not checked.
if node_epoch != 0 {
if let Some(election) = &self.ctx.election
&& !election.is_leader()
{
return Some(Ok(header_only(Some(Error::is_not_leader()))));
}
let kv_backend = &*self.ctx.leader_cached_kv_backend;
match accept_node_epoch(kv_backend, self.sender_id.id, node_epoch).await {
Ok(true) => {}
// No error, so the datanode does not reconnect.
Ok(false) => {
self.handler_group
.deregister_push(self.sender_id, &self.tx)
.await;
return Some(Ok(header_only(None)));
}
Err(e) => return Some(Err(e)),
}
}
// Before the handlers run, so the stream receives what they broadcast.
if self.pusher_epoch.is_none_or(|epoch| node_epoch > epoch) {
self.pusher_epoch = Some(node_epoch);
let pusher = Pusher::new(self.tx.clone()).with_node_epoch(node_epoch);
self.handler_group
.register_pusher(self.sender_id, pusher)
.await;
}
None
}
async fn send_not_leader_error(&mut self) {
let _ = self
.tx
@@ -279,7 +352,9 @@ where
async fn cleanup(&self) {
info!("Heartbeat stream closed, sender: {}", self.sender_id);
let _ = self.handler_group.deregister_push(self.sender_id).await;
self.handler_group
.deregister_push(self.sender_id, &self.tx)
.await;
}
}
@@ -400,17 +475,62 @@ fn get_node_id(header: &RequestHeader) -> u64 {
}
}
async fn register_pusher(
handler_group: &HeartbeatHandlerGroup,
header: &RequestHeader,
sender: Sender<std::result::Result<HeartbeatResponse, tonic::Status>>,
) -> PusherId {
let role = header.role();
let id = get_node_id(header);
let pusher_id = PusherId::new(role, id);
let pusher = Pusher::new(sender);
handler_group.register_pusher(pusher_id, pusher).await;
pusher_id
/// Returns whether `node_epoch` is not lower than the highest epoch recorded for the datanode,
/// raising the record when it is higher. An epoch too far ahead of this metasrv's clock is
/// refused, since recording it would fence every later process.
async fn accept_node_epoch(
kv_backend: &dyn KvBackend<Error = common_meta::error::Error>,
node_id: u64,
node_epoch: u64,
) -> Result<bool> {
let key = DatanodeEpochKey::new(node_id).to_bytes();
let mut current = kv_backend.get(&key).await.context(error::KvBackendSnafu)?;
let limit = current_time_millis().max(0) as u64 + MAX_NODE_EPOCH_AHEAD.as_millis() as u64;
loop {
let (expect, highest) = match current {
Some(kv) => {
let value = DatanodeEpochValue::try_from_raw_value(&kv.value)
.context(error::KvBackendSnafu)?;
(kv.value, value.node_epoch)
}
None => (vec![], 0),
};
match node_epoch.cmp(&highest) {
Ordering::Equal => return Ok(true),
Ordering::Greater if node_epoch <= limit => {}
_ => {
warn!("Fenced datanode {node_id} with epoch {node_epoch}, recorded {highest}");
return Ok(false);
}
}
let value = DatanodeEpochValue { node_epoch }
.try_as_raw_value()
.context(error::KvBackendSnafu)?;
let req = CompareAndPutRequest {
key: key.clone(),
expect,
value,
};
let resp = kv_backend
.compare_and_put(req)
.await
.context(error::KvBackendSnafu)?;
if resp.success {
info!("Recorded epoch {node_epoch} of datanode {node_id}");
return Ok(true);
}
current = resp.prev_kv;
}
}
fn header_only(error: Option<Error>) -> HeartbeatResponse {
HeartbeatResponse {
header: Some(ResponseHeader {
protocol_version: PROTOCOL_VERSION,
error,
}),
..Default::default()
}
}
#[cfg(test)]
@@ -596,9 +716,9 @@ mod tests {
async fn test_heartbeat_session_init_registers_sender() {
let (tx, mut rx) = mpsc::channel(8);
let handler_group = test_handler_group();
let sender_id = sender_id(Role::Datanode, 42);
let sender_id = sender_id(Role::Flownode, 42);
let requests =
MockHeartbeatRequestStream::new(vec![Ok(heartbeat_request(Role::Datanode, 42))]);
MockHeartbeatRequestStream::new(vec![Ok(heartbeat_request(Role::Flownode, 42))]);
let session = init_session(
requests,
@@ -619,9 +739,9 @@ mod tests {
async fn test_heartbeat_session_run_deregisters_sender_on_stream_close() {
let (tx, mut rx) = mpsc::channel(8);
let handler_group = test_handler_group();
let sender_id = sender_id(Role::Datanode, 42);
let sender_id = sender_id(Role::Flownode, 42);
let requests =
MockHeartbeatRequestStream::new(vec![Ok(heartbeat_request(Role::Datanode, 42))]);
MockHeartbeatRequestStream::new(vec![Ok(heartbeat_request(Role::Flownode, 42))]);
let session = init_session(
requests,
tx,
@@ -641,9 +761,9 @@ mod tests {
async fn test_heartbeat_session_run_forwards_stream_error_after_init() {
let (tx, mut rx) = mpsc::channel(8);
let handler_group = test_handler_group();
let sender_id = sender_id(Role::Datanode, 42);
let sender_id = sender_id(Role::Flownode, 42);
let requests = MockHeartbeatRequestStream::new(vec![
Ok(heartbeat_request(Role::Datanode, 42)),
Ok(heartbeat_request(Role::Flownode, 42)),
Err(Status::unavailable("temporary")),
]);
let session = init_session(
@@ -668,9 +788,9 @@ mod tests {
async fn test_heartbeat_session_leader_step_down_sends_aborted_and_deregisters() {
let (tx, mut rx) = mpsc::channel(8);
let handler_group = test_handler_group();
let sender_id = sender_id(Role::Datanode, 42);
let sender_id = sender_id(Role::Flownode, 42);
let requests = MockHeartbeatRequestStream::pending_after(vec![Ok(heartbeat_request(
Role::Datanode,
Role::Flownode,
42,
))]);
let session = init_session(
@@ -694,9 +814,9 @@ mod tests {
async fn test_heartbeat_session_leader_watcher_closed_sends_unavailable_and_deregisters() {
let (tx, mut rx) = mpsc::channel(8);
let handler_group = test_handler_group();
let sender_id = sender_id(Role::Datanode, 42);
let sender_id = sender_id(Role::Flownode, 42);
let requests = MockHeartbeatRequestStream::pending_after(vec![Ok(heartbeat_request(
Role::Datanode,
Role::Flownode,
42,
))]);
let session = init_session(
@@ -768,4 +888,459 @@ mod tests {
assert_eq!(i, get_node_id(&header));
}
}
mod node_epoch {
use std::sync::Arc;
use std::time::Duration;
use api::v1::meta::{HeartbeatRequest, HeartbeatResponse, MailboxMessage, Peer, Role};
use common_meta::error::UnexpectedSnafu;
use common_meta::key::datanode_epoch::{DatanodeEpochKey, DatanodeEpochValue};
use common_meta::key::{MetadataValue, TableMetadataManager};
use common_meta::kv_backend::memory::MemoryKvBackend;
use common_meta::kv_backend::test_util::MockKvBackendBuilder;
use common_meta::kv_backend::{KvBackend, KvBackendRef, ResettableKvBackend};
use common_meta::region_keeper::MemoryRegionKeeper;
use common_meta::rpc::store::{PutRequest, RangeResponse};
use common_meta::sequence::SequenceBuilder;
use common_time::util::current_time_millis;
use tokio::sync::{Notify, mpsc};
use super::{MockHeartbeatRequestStream, MockLeaderStepDown, heartbeat_request};
use crate::discovery::utils::find_datanode_lease_value;
use crate::handler::region_lease_handler::RegionLeaseHandler;
use crate::handler::test_utils::TestEnv;
use crate::handler::{
HandleControl, HeartbeatAccumulator, HeartbeatHandler, HeartbeatHandlerGroup,
HeartbeatHandlerGroupBuilder, HeartbeatMailbox, Pusher, PusherId, Pushers,
};
use crate::metasrv::Context;
use crate::service::heartbeat::{
HeartbeatResponseResult, HeartbeatSession, accept_node_epoch,
};
use crate::service::mailbox::{BroadcastChannel, Channel, MailboxRef};
use crate::service::store::cached_kv::LeaderCachedKvBackend;
const NODE_ID: u64 = 1;
/// A datanode stream that has sent its handshake, which carries no epoch.
struct Stream {
session: HeartbeatSession<MockHeartbeatRequestStream, MockLeaderStepDown>,
rx: mpsc::Receiver<HeartbeatResponseResult>,
}
impl Stream {
async fn open(group: &Arc<HeartbeatHandlerGroup>, ctx: &Context) -> Self {
let (tx, mut rx) = mpsc::channel(8);
let handshake = heartbeat_request(Role::Datanode, NODE_ID);
let session = HeartbeatSession::init(
MockHeartbeatRequestStream::new(vec![Ok(handshake)]),
tx,
None,
group.clone(),
ctx.clone(),
)
.await
.unwrap();
rx.recv().await.unwrap().unwrap();
Self { session, rx }
}
async fn send(&mut self, req: HeartbeatRequest) -> HeartbeatResponse {
assert!(self.session.handle_request(req, false).await);
self.rx.recv().await.unwrap().unwrap()
}
/// Returns whether the stream has received a pushed message.
fn received_push(&mut self) -> bool {
self.rx.try_recv().is_ok()
}
}
fn group_builder() -> HeartbeatHandlerGroupBuilder {
let lease_handler = RegionLeaseHandler::new(
10,
Arc::new(TableMetadataManager::new(Arc::new(MemoryKvBackend::new()))),
Arc::new(MemoryRegionKeeper::default()),
None,
);
let mut builder = HeartbeatHandlerGroupBuilder::new(Pushers::default())
.with_region_lease_handler(Some(lease_handler))
.add_default_handlers();
builder
.add_handler_after("ResponseHeaderHandler", BroadcastHandler)
.unwrap();
builder
}
fn new_group() -> Arc<HeartbeatHandlerGroup> {
Arc::new(group_builder().build().unwrap())
}
/// Holds a request from the address "pause": notifies `entered`, then waits for `release`.
struct PauseHandler {
entered: Arc<Notify>,
release: Arc<Notify>,
}
#[async_trait::async_trait]
impl HeartbeatHandler for PauseHandler {
fn is_acceptable(&self, role: Role) -> bool {
role == Role::Datanode
}
async fn handle(
&self,
req: &HeartbeatRequest,
_: &mut Context,
_: &mut HeartbeatAccumulator,
) -> crate::error::Result<HandleControl> {
if req.peer.as_ref().is_some_and(|peer| peer.addr == "pause") {
self.entered.notify_one();
self.release.notified().await;
}
Ok(HandleControl::Continue)
}
}
/// Broadcasts to datanodes while handling a request from the address "broadcast".
struct BroadcastHandler;
#[async_trait::async_trait]
impl HeartbeatHandler for BroadcastHandler {
fn is_acceptable(&self, role: Role) -> bool {
role == Role::Datanode
}
async fn handle(
&self,
req: &HeartbeatRequest,
ctx: &mut Context,
_: &mut HeartbeatAccumulator,
) -> crate::error::Result<HandleControl> {
if req
.peer
.as_ref()
.is_some_and(|peer| peer.addr == "broadcast")
{
let msg = MailboxMessage::default();
ctx.mailbox
.broadcast(&BroadcastChannel::Datanode, &msg)
.await?;
}
Ok(HandleControl::Continue)
}
}
fn heartbeat(node_epoch: u64, addr: &str) -> HeartbeatRequest {
HeartbeatRequest {
peer: Some(Peer::new(NODE_ID, addr)),
node_epoch,
..heartbeat_request(Role::Datanode, NODE_ID)
}
}
fn mailbox(pushers: Pushers) -> MailboxRef {
let seq = SequenceBuilder::new("test_seq", Arc::new(MemoryKvBackend::new())).build();
HeartbeatMailbox::create(pushers, seq)
}
async fn push(group: &HeartbeatHandlerGroup) {
mailbox(group.pushers())
.send_oneway(&Channel::Datanode(NODE_ID), MailboxMessage::default())
.await
.unwrap();
}
async fn lease_addr(ctx: &Context) -> Option<String> {
let in_memory = ctx.in_memory.clone() as KvBackendRef;
find_datanode_lease_value(&in_memory, NODE_ID)
.await
.unwrap()
.map(|lease| lease.node_addr)
}
/// Records `node_epoch` as the highest epoch of the node, through the leader cache.
async fn record(ctx: &Context, node_epoch: u64) {
let key = DatanodeEpochKey::new(NODE_ID).to_bytes();
let value = DatanodeEpochValue { node_epoch }
.try_as_raw_value()
.unwrap();
let req = PutRequest::new().with_key(key).with_value(value);
ctx.leader_cached_kv_backend.put(req).await.unwrap();
}
fn epochs() -> (u64, u64) {
let now = current_time_millis() as u64;
(now - 1000, now)
}
#[tokio::test]
async fn test_newer_process_fences_the_older_one() {
let group = new_group();
let mut ctx = TestEnv::new().ctx();
// A mailbox waiting for replies from the datanode.
let pushers = Pushers::default();
let (tx, _rx) = mpsc::channel(8);
pushers
.insert(PusherId::new(Role::Datanode, NODE_ID), Pusher::new(tx))
.await;
ctx.mailbox = mailbox(pushers);
let wait_for_reply = || {
ctx.mailbox.send(
&Channel::Datanode(NODE_ID),
MailboxMessage::default(),
Duration::from_millis(500),
)
};
let (e1, e2) = epochs();
let mut a1 = Stream::open(&group, &ctx).await;
let mut a2 = Stream::open(&group, &ctx).await;
let res = a1.send(heartbeat(e1, "a1")).await;
assert!(res.region_lease.is_some());
// A2 replaces A1 and records its epoch.
let res = a2.send(heartbeat(e2, "a2")).await;
assert!(res.region_lease.is_some());
// A1 is fenced: no lease, no keep-alive, no error that would make it reconnect.
let res = a1.send(heartbeat(e1, "a1")).await;
assert!(res.region_lease.is_none());
assert!(res.header.unwrap().error.is_none());
assert_eq!(lease_addr(&ctx).await.as_deref(), Some("a2"));
// A1's mailbox reply is dropped, A2's is delivered.
let a1_reply = wait_for_reply().await.unwrap();
let a2_reply = wait_for_reply().await.unwrap();
for (stream, reply, epoch) in [(&mut a1, &a1_reply, e1), (&mut a2, &a2_reply, e2)] {
let mut req = heartbeat(epoch, "a");
req.mailbox_message = Some(MailboxMessage {
id: reply.message_id(),
..Default::default()
});
stream.send(req).await;
}
assert!(a1_reply.await.is_err());
assert!(a2_reply.await.is_ok());
}
#[tokio::test]
async fn test_pusher_follows_the_newer_process() {
let group = new_group();
let mut ctx = TestEnv::new().ctx();
ctx.mailbox = mailbox(group.pushers());
let (e1, e2) = epochs();
let mut a1 = Stream::open(&group, &ctx).await;
a1.send(heartbeat(e1, "a1")).await;
// A2 takes over before the handlers of its first heartbeat run, so it receives
// what they broadcast ahead of the response.
let mut a2 = Stream::open(&group, &ctx).await;
assert!(
a2.session
.handle_request(heartbeat(e2, "broadcast"), false)
.await
);
let broadcast = a2.rx.recv().await.unwrap().unwrap();
assert!(broadcast.mailbox_message.is_some());
let res = a2.rx.recv().await.unwrap().unwrap();
assert!(res.region_lease.is_some());
assert!(!a1.received_push());
push(&group).await;
assert!(a2.received_push());
// A1's new stream does not take the pusher back.
let mut a1_again = Stream::open(&group, &ctx).await;
a1_again.send(heartbeat(e1, "a1")).await;
push(&group).await;
assert!(a2.received_push());
// A stream ending removes only its own pusher.
a1.session.run().await;
a1_again.session.run().await;
push(&group).await;
assert!(a2.received_push());
// A2's new stream takes over with the same epoch, and A2's old stream ending
// keeps it.
let mut a2_again = Stream::open(&group, &ctx).await;
a2_again.send(heartbeat(e2, "a2")).await;
a2.session.run().await;
push(&group).await;
assert!(a2_again.received_push());
// A fenced stream removes its own pusher, and a handshake does not take the slot.
record(&ctx, e2 + 1).await;
a2_again.send(heartbeat(e2, "a2")).await;
let pusher_id = PusherId::new(Role::Datanode, NODE_ID);
assert!(!group.contains_pusher(&pusher_id).await);
let _a1_reconnected = Stream::open(&group, &ctx).await;
assert!(!group.contains_pusher(&pusher_id).await);
}
#[tokio::test]
async fn test_newer_process_waits_for_an_older_request() {
let (entered, release) = (Arc::new(Notify::new()), Arc::new(Notify::new()));
let mut builder = group_builder();
let pause = PauseHandler {
entered: entered.clone(),
release: release.clone(),
};
builder
.add_handler_before("DatanodeKeepLeaseHandler", pause)
.unwrap();
let group = Arc::new(builder.build().unwrap());
let ctx = TestEnv::new().ctx();
let (e1, e2) = epochs();
let mut a1 = Stream::open(&group, &ctx).await;
let mut a2 = Stream::open(&group, &ctx).await;
// A1's request passes the check and is held before the keep-alive handler, and
// A2's request waits for it.
let a1_req = a1.session.handle_request(heartbeat(e1, "pause"), false);
let a2_req = async {
entered.notified().await;
let mut a2_req = std::pin::pin!(a2.send(heartbeat(e2, "a2")));
let waiting = tokio::time::timeout(Duration::from_millis(100), &mut a2_req).await;
assert!(waiting.is_err());
release.notify_one();
a2_req.await
};
let (open, res) = tokio::join!(a1_req, a2_req);
assert!(open);
assert!(res.region_lease.is_some());
assert_eq!(lease_addr(&ctx).await.as_deref(), Some("a2"));
// A1 can no longer overwrite the address or get a lease.
a1.rx.recv().await.unwrap().unwrap();
let res = a1.send(heartbeat(e1, "a1")).await;
assert!(res.region_lease.is_none());
assert_eq!(lease_addr(&ctx).await.as_deref(), Some("a2"));
}
#[tokio::test]
async fn test_kv_error_fails_the_request() {
let group = new_group();
let mut ctx = TestEnv::new().ctx();
let kv_backend = MockKvBackendBuilder::default()
.range_fn(Arc::new(|_| {
Ok(RangeResponse {
kvs: vec![],
more: false,
})
}))
.txn(Arc::new(|_| {
UnexpectedSnafu {
err_msg: "injected",
}
.fail()
}))
.build()
.unwrap();
let cached = LeaderCachedKvBackend::with_always_leader(Arc::new(kv_backend));
ctx.leader_cached_kv_backend = Arc::new(cached);
let (e1, _) = epochs();
let mut a1 = Stream::open(&group, &ctx).await;
let req = HeartbeatRequest {
node_epoch: e1,
..heartbeat_request(Role::Datanode, NODE_ID)
};
assert!(a1.session.handle_request(req, false).await);
assert!(a1.rx.recv().await.unwrap().is_err());
}
#[tokio::test]
async fn test_first_epoch_keeps_waiting_messages() {
let group = new_group();
let mut ctx = TestEnv::new().ctx();
ctx.mailbox = mailbox(group.pushers());
let (e1, _) = epochs();
let mut a1 = Stream::open(&group, &ctx).await;
// A request without an epoch registers the pusher first.
a1.send(heartbeat(0, "a1")).await;
let reply = ctx
.mailbox
.send(
&Channel::Datanode(NODE_ID),
MailboxMessage::default(),
Duration::from_secs(10),
)
.await
.unwrap();
let id = a1
.rx
.recv()
.await
.unwrap()
.unwrap()
.mailbox_message
.unwrap()
.id;
tokio::pin!(reply);
// The stream's pusher takes the epoch of its first heartbeat without cancelling
// the message, which a later heartbeat answers.
a1.send(heartbeat(e1, "a1")).await;
assert!(futures::poll!(reply.as_mut()).is_pending());
let mut req = heartbeat(e1, "a1");
req.mailbox_message = Some(MailboxMessage {
id,
..Default::default()
});
a1.send(req).await;
assert!(reply.await.is_ok());
}
#[tokio::test]
async fn test_record_is_reloaded_on_leader_change() {
let store: KvBackendRef = Arc::new(MemoryKvBackend::new());
let cached = LeaderCachedKvBackend::with_always_leader(store.clone());
let now = current_time_millis() as u64;
let (e1, e2) = (now - 1000, now);
assert!(accept_node_epoch(&cached, 1, e1).await.unwrap());
assert!(accept_node_epoch(&cached, 1, e2).await.unwrap());
// The cache is reset and loaded again when this metasrv becomes the leader. The record
// is then served from the loaded cache, even without the store.
cached.reset();
cached.load().await.unwrap();
store
.delete(&DatanodeEpochKey::new(1).to_bytes(), false)
.await
.unwrap();
assert!(!accept_node_epoch(&cached, 1, e1).await.unwrap());
assert!(accept_node_epoch(&cached, 1, e2).await.unwrap());
}
#[tokio::test]
async fn test_epoch_far_ahead_is_not_recorded() {
let kv_backend = MemoryKvBackend::new();
let now = current_time_millis() as u64;
let ahead = now + Duration::from_secs(10 * 60).as_millis() as u64;
assert!(!accept_node_epoch(&kv_backend, 1, ahead).await.unwrap());
let key = DatanodeEpochKey::new(1).to_bytes();
assert!(kv_backend.get(&key).await.unwrap().is_none());
let slightly_ahead = now + Duration::from_secs(60).as_millis() as u64;
assert!(
accept_node_epoch(&kv_backend, 1, slightly_ahead)
.await
.unwrap()
);
}
#[tokio::test]
async fn test_datanode_without_epoch_is_not_checked() {
let group = new_group();
let ctx = TestEnv::new().ctx();
// Even after another process of the node has recorded its epoch.
let (_, e2) = epochs();
record(&ctx, e2).await;
let mut a1 = Stream::open(&group, &ctx).await;
let res = a1.send(heartbeat(0, "a1")).await;
assert!(res.region_lease.is_some());
push(&group).await;
assert!(a1.received_push());
}
}
}
+68 -29
View File
@@ -69,9 +69,10 @@ pub struct LeaderCachedKvBackend {
check_leader: CheckLeaderRef,
store: KvBackendRef,
cache: ResettableKvBackendRef,
/// Incremented before and after each store mutation. The first increment
/// detects overlapping writers; the second invalidates cache fills that
/// read the old value while the store mutation was in flight.
/// Incremented before and after each store mutation, also while the cache is
/// disabled. The first increment detects overlapping writers; the second
/// invalidates cache fills and loads that read the old value while the
/// store mutation was in flight.
version: AtomicUsize,
name: String,
}
@@ -94,11 +95,13 @@ impl LeaderCachedKvBackend {
Self::new(Arc::new(AlwaysLeader), store)
}
/// The caller MUST ensure during the loading, there are no mutation requests reaching the `LeaderCachedKvStore`.
/// Loads the cached prefixes. Writes may run meanwhile: a prefix whose load overlapped a
/// write is dropped from the cache again, and read through later.
pub async fn load(&self) -> Result<()> {
for prefix in &CACHE_KEY_PREFIXES[..] {
let _timer =
metrics::METRIC_META_LEADER_CACHED_KV_LOAD_ELAPSED.with_label_values(&[prefix]);
let ver = self.get_version();
// TODO(weny): Refactors PaginationStream's output to unary output.
let stream = PaginationStream::new(
@@ -110,6 +113,7 @@ impl LeaderCachedKvBackend {
.into_stream();
let kvs = stream.try_collect::<Vec<_>>().await?;
let keys = kvs.iter().map(|kv| kv.key.clone()).collect::<Vec<_>>();
self.cache
.batch_put(BatchPutRequest {
@@ -117,6 +121,10 @@ impl LeaderCachedKvBackend {
prev_kv: false,
})
.await?;
// A write while loading may have changed what was read.
if !self.validate_version(ver) {
self.invalid_keys(keys).await?;
}
}
Ok(())
@@ -210,10 +218,6 @@ impl KvBackend for LeaderCachedKvBackend {
}
async fn put(&self, req: PutRequest) -> Result<PutResponse> {
if !self.is_leader() {
return self.store.put(req).await;
}
let ver = self.create_new_version();
let res = self.store.put(req.clone()).await?;
@@ -229,10 +233,6 @@ impl KvBackend for LeaderCachedKvBackend {
}
async fn batch_put(&self, req: BatchPutRequest) -> Result<BatchPutResponse> {
if !self.is_leader() {
return self.store.batch_put(req).await;
}
let ver = self.create_new_version();
let res = self.store.batch_put(req.clone()).await?;
@@ -305,10 +305,6 @@ impl KvBackend for LeaderCachedKvBackend {
}
async fn compare_and_put(&self, req: CompareAndPutRequest) -> Result<CompareAndPutResponse> {
if !self.is_leader() {
return self.store.compare_and_put(req).await;
}
let _ = self.create_new_version();
let key = req.key.clone();
@@ -323,10 +319,6 @@ impl KvBackend for LeaderCachedKvBackend {
}
async fn delete_range(&self, req: DeleteRangeRequest) -> Result<DeleteRangeResponse> {
if !self.is_leader() {
return self.store.delete_range(req).await;
}
let _ = self.create_new_version();
let res = self.store.delete_range(req.clone()).await?;
@@ -336,10 +328,6 @@ impl KvBackend for LeaderCachedKvBackend {
}
async fn batch_delete(&self, req: BatchDeleteRequest) -> Result<BatchDeleteResponse> {
if !self.is_leader() {
return self.store.batch_delete(req).await;
}
let _ = self.create_new_version();
let res = self.store.batch_delete(req.clone()).await?;
@@ -354,10 +342,6 @@ impl TxnService for LeaderCachedKvBackend {
type Error = Error;
async fn txn(&self, txn: Txn) -> Result<TxnResponse> {
if !self.is_leader() {
return self.store.txn(txn).await;
}
let _ = self.create_new_version();
let res = self.store.txn(txn.clone()).await?;
@@ -531,7 +515,8 @@ mod tests {
Write::DeleteRange,
Write::BatchDelete,
];
const KEY: &[u8] = b"route";
/// Under the first prefix the leader cache loads.
const KEY: &[u8] = b"__table_name/1";
impl Write {
async fn execute(self, store: &LeaderCachedKvBackend, value: &[u8]) {
@@ -733,6 +718,60 @@ mod tests {
}
}
struct EnableCache(std::sync::atomic::AtomicBool);
impl CheckLeader for EnableCache {
fn check(&self) -> bool {
self.0.load(Ordering::Relaxed)
}
}
/// A write made while the cache is disabled overlaps the load: the load reads the old value
/// either before the write commits and installs it afterwards, or while the write is in
/// flight and completes, with the cache enabled, before the write commits.
#[tokio::test]
async fn test_load_overlaps_write() {
for write in WRITES {
for write_first in [false, true] {
let backing = Arc::new(ControlledBackend::default());
backing
.inner
.put(PutRequest::new().with_key(KEY).with_value(b"old"))
.await
.unwrap();
let enabled = Arc::new(EnableCache(Default::default()));
let store = LeaderCachedKvBackend::new(enabled.clone(), backing.clone());
if write_first {
let commit = backing.before_write.arm();
let writer = write.execute(&store, b"new");
tokio::pin!(writer);
assert!(poll!(writer.as_mut()).is_pending());
store.load().await.unwrap();
enabled.0.store(true, Ordering::Relaxed);
commit.send(()).unwrap();
writer.await;
} else {
let return_read = backing.after_read.arm();
let loader = store.load();
tokio::pin!(loader);
assert!(poll!(loader.as_mut()).is_pending());
write.execute(&store, b"new").await;
return_read.send(()).unwrap();
loader.await.unwrap();
enabled.0.store(true, Ordering::Relaxed);
}
let durable = backing.inner.get(KEY).await.unwrap();
assert_eq!(
read(&store, false).await,
durable,
"{write:?}, write_first={write_first}"
);
}
}
}
#[tokio::test]
async fn test_cache_miss_overlaps_reset() {
for batch in [true, false] {
@@ -0,0 +1 @@
{"node_epoch":1761091200000}