meta_srv/handler/
utils.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::future::Future;
16use std::sync::Arc;
17
18use api::v1::meta::{Peer, Role};
19use common_meta::instruction::CacheIdent;
20use common_meta::key::MetadataValue;
21use common_meta::key::node_address::NodeAddressValue;
22use common_meta::rpc::store::PutRequest;
23use common_telemetry::{error, info};
24use dashmap::DashMap;
25use snafu::ResultExt;
26use tokio::sync::Mutex;
27
28use crate::Result;
29use crate::error::{InvalidNodeInfoFormatSnafu, KvBackendSnafu};
30use crate::metasrv::Context;
31
32/// Serializes address changes per node and acknowledges only successful updates.
33/// Node information refreshed on every heartbeat has a separate lifecycle.
34#[derive(Debug, Default)]
35pub(crate) struct NodeAddressUpdater {
36    states: DashMap<(Role, u64), Arc<Mutex<NodeAddressEpochs>>>,
37}
38
39#[derive(Debug, Default)]
40struct NodeAddressEpochs {
41    /// Rejects stale heartbeats even after a newer epoch failed to save.
42    seen_epoch: Option<u64>,
43    /// Suppresses retries only after persistence succeeds.
44    saved_epoch: Option<u64>,
45}
46
47impl NodeAddressUpdater {
48    /// Saves a new or previously failed epoch. Returns true only when this call
49    /// saved the address, or false when the heartbeat is stale or already saved.
50    /// Cache invalidation is separate and does not affect acknowledgement.
51    pub(crate) async fn update_if_needed<F, Fut>(
52        &self,
53        node: (Role, u64),
54        epoch: u64,
55        save: F,
56    ) -> Result<bool>
57    where
58        F: FnOnce() -> Fut,
59        Fut: Future<Output = Result<()>>,
60    {
61        // Never retain a DashMap guard across await, even while waiting for a
62        // different node: its entry may live in the same shard.
63        let state = {
64            let entry = self.states.entry(node).or_default();
65            Arc::clone(entry.value())
66        };
67        let mut state = state.lock().await;
68        if state.seen_epoch.is_some_and(|seen| epoch < seen) || state.saved_epoch == Some(epoch) {
69            return Ok(false);
70        }
71
72        // Seeing an epoch rejects stale heartbeats, but only a successful save
73        // suppresses retries of that same epoch. Cancellation also leaves it retryable.
74        state.seen_epoch = Some(epoch);
75        save().await?;
76        state.saved_epoch = Some(epoch);
77        Ok(true)
78    }
79}
80
81/// Persists an address through the leader cache.
82pub(crate) async fn save_node_address(ctx: &Context, key: Vec<u8>, peer: Peer) -> Result<()> {
83    let address = NodeAddressValue::new(peer);
84    let value = address
85        .try_as_raw_value()
86        .context(InvalidNodeInfoFormatSnafu)?;
87    ctx.leader_cached_kv_backend
88        .put(PutRequest {
89            key,
90            value,
91            prev_kv: false,
92        })
93        .await
94        .context(KvBackendSnafu)?;
95    info!("Successfully updated node address: {:?}", address.peer);
96    Ok(())
97}
98
99/// Preserves the existing best-effort notification behavior. Notification
100/// failures do not undo a persisted address or cause same-epoch rewrites.
101pub(crate) async fn invalidate_address_caches(
102    ctx: &Context,
103    node_id: u64,
104    cache_idents: &[CacheIdent],
105) {
106    if let Err(e) = ctx
107        .cache_invalidator
108        .invalidate(&Default::default(), cache_idents)
109        .await
110    {
111        error!(e; "Failed to invalidate address caches for node {}: {:?}", node_id, cache_idents);
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use std::time::Duration;
118
119    use tokio::sync::Notify;
120    use tokio::time::timeout;
121
122    use super::*;
123
124    #[tokio::test]
125    async fn test_failed_save_retries_same_epoch_and_rejects_stale_heartbeats() {
126        let updater = NodeAddressUpdater::default();
127        let key = (Role::Datanode, 1);
128        assert!(
129            updater
130                .update_if_needed(key, 2, || async {
131                    crate::error::UnexpectedSnafu {
132                        violated: "injected save failure",
133                    }
134                    .fail()
135                })
136                .await
137                .is_err()
138        );
139        assert!(
140            !updater
141                .update_if_needed(key, 1, || async {
142                    panic!("stale heartbeat must not save");
143                })
144                .await
145                .unwrap()
146        );
147        assert!(
148            updater
149                .update_if_needed(key, 2, || async { Ok(()) })
150                .await
151                .unwrap()
152        );
153        assert!(
154            !updater
155                .update_if_needed(key, 2, || async {
156                    panic!("successful epoch must not save again");
157                })
158                .await
159                .unwrap()
160        );
161        assert!(
162            updater
163                .update_if_needed(key, 3, || async { Ok(()) })
164                .await
165                .unwrap()
166        );
167    }
168
169    #[tokio::test]
170    async fn test_cancelled_save_can_be_retried() {
171        let updater = NodeAddressUpdater::default();
172        let key = (Role::Datanode, 1);
173        let entered = Notify::new();
174        {
175            let first = updater.update_if_needed(key, 0, || async {
176                entered.notify_one();
177                std::future::pending::<Result<()>>().await
178            });
179            tokio::pin!(first);
180            tokio::select! {
181                _ = entered.notified() => {},
182                _ = &mut first => panic!("save must remain pending"),
183            }
184        }
185        assert!(
186            timeout(
187                Duration::from_secs(2),
188                updater.update_if_needed(key, 0, || async { Ok(()) })
189            )
190            .await
191            .unwrap()
192            .unwrap()
193        );
194    }
195
196    #[test]
197    fn test_same_node_serialized_other_nodes_progress() {
198        // A watchdog outside the runtime detects synchronous DashMap deadlocks;
199        // a Tokio timeout alone cannot fire if the only worker is blocked.
200        let (progress_tx, progress_rx) = std::sync::mpsc::sync_channel(1);
201        let release = Arc::new(Notify::new());
202        let finish = Arc::new(Notify::new());
203        let worker = std::thread::spawn({
204            let release = release.clone();
205            let finish = finish.clone();
206            move || {
207                tokio::runtime::Builder::new_current_thread()
208                    .enable_time()
209                    .build()
210                    .unwrap()
211                    .block_on(async move {
212                        let updater = Arc::new(NodeAddressUpdater::default());
213                        let entered = Arc::new(Notify::new());
214
215                        let first = tokio::spawn({
216                            let updater = updater.clone();
217                            let entered = entered.clone();
218                            let release = release.clone();
219                            async move {
220                                updater
221                                    .update_if_needed((Role::Datanode, 1), 1, || async {
222                                        entered.notify_one();
223                                        release.notified().await;
224                                        Ok(())
225                                    })
226                                    .await
227                            }
228                        });
229                        entered.notified().await;
230                        let second = tokio::spawn({
231                            let updater = updater.clone();
232                            async move {
233                                updater
234                                    .update_if_needed((Role::Datanode, 1), 2, || async { Ok(()) })
235                                    .await
236                            }
237                        });
238                        tokio::task::yield_now().await;
239                        assert!(!second.is_finished());
240                        // Exercise shared DashMap shards on a single-thread runtime while the
241                        // first node is suspended in external I/O.
242                        timeout(Duration::from_secs(2), async {
243                            for id in 0_u64..128 {
244                                assert!(
245                                    updater
246                                        .update_if_needed((Role::Datanode, id + 2), 1, || async {
247                                            Ok(())
248                                        })
249                                        .await
250                                        .unwrap()
251                                );
252                            }
253                        })
254                        .await
255                        .unwrap();
256                        progress_tx.send(()).unwrap();
257                        finish.notified().await;
258                        timeout(Duration::from_secs(2), async {
259                            assert!(first.await.unwrap().unwrap());
260                            assert!(second.await.unwrap().unwrap());
261                        })
262                        .await
263                        .unwrap();
264                        assert!(
265                            !updater
266                                .update_if_needed((Role::Datanode, 1), 1, || async {
267                                    panic!("old epoch must not overwrite a newer one");
268                                })
269                                .await
270                                .unwrap()
271                        );
272                    });
273            }
274        });
275        let progress = progress_rx.recv_timeout(Duration::from_secs(5));
276        release.notify_one();
277        finish.notify_one();
278        progress.expect("other node updates must progress while one node save is blocked");
279        worker.join().unwrap();
280    }
281}
meta_srv/handler/utils.rs

meta_srv/handler/
utils.rs
