meta_srv/handler/
collect_stats_handler.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::cmp::Ordering;
16use std::sync::Arc;
17
18use api::v1::meta::{HeartbeatRequest, Peer, Role};
19use common_meta::datanode::{DatanodeStatKey, DatanodeStatValue, Stat};
20use common_meta::instruction::CacheIdent;
21use common_meta::key::MetadataKey;
22use common_meta::key::node_address::NodeAddressKey;
23use common_meta::rpc::store::PutRequest;
24use common_telemetry::{error, warn};
25use dashmap::DashMap;
26use snafu::ResultExt;
27use tokio::sync::Mutex;
28
29use crate::error::{self, Result};
30use crate::handler::utils::{self, NodeAddressUpdater};
31use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler};
32use crate::metasrv::Context;
33
34#[derive(Debug, Default)]
35struct EpochStats {
36    stats: Vec<Stat>,
37    epoch: Option<u64>,
38}
39
40impl EpochStats {
41    #[inline]
42    fn drain_all(&mut self) -> Vec<Stat> {
43        std::mem::take(&mut self.stats)
44    }
45
46    #[inline]
47    fn clear_stats(&mut self) {
48        self.stats.clear();
49    }
50
51    #[inline]
52    fn push_stat(&mut self, stat: Stat) {
53        self.stats.push(stat);
54    }
55
56    #[inline]
57    fn len(&self) -> usize {
58        self.stats.len()
59    }
60
61    #[inline]
62    fn epoch(&self) -> Option<u64> {
63        self.epoch
64    }
65
66    #[inline]
67    fn set_epoch(&mut self, epoch: u64) {
68        self.epoch = Some(epoch);
69    }
70}
71
72const DEFAULT_FLUSH_STATS_FACTOR: usize = 3;
73
74pub struct CollectStatsHandler {
75    stats_cache: DashMap<DatanodeStatKey, Arc<Mutex<EpochStats>>>,
76    flush_stats_factor: usize,
77    address_updater: NodeAddressUpdater,
78}
79
80impl Default for CollectStatsHandler {
81    fn default() -> Self {
82        Self::new(None)
83    }
84}
85
86impl CollectStatsHandler {
87    pub fn new(flush_stats_factor: Option<usize>) -> Self {
88        Self {
89            flush_stats_factor: flush_stats_factor.unwrap_or(DEFAULT_FLUSH_STATS_FACTOR),
90            stats_cache: DashMap::default(),
91            address_updater: NodeAddressUpdater::default(),
92        }
93    }
94}
95
96#[async_trait::async_trait]
97impl HeartbeatHandler for CollectStatsHandler {
98    fn is_acceptable(&self, role: Role) -> bool {
99        role == Role::Datanode
100    }
101
102    async fn handle(
103        &self,
104        _req: &HeartbeatRequest,
105        ctx: &mut Context,
106        acc: &mut HeartbeatAccumulator,
107    ) -> Result<HandleControl> {
108        let Some(current_stat) = acc.stat.take() else {
109            return Ok(HandleControl::Continue);
110        };
111
112        let key = current_stat.stat_key();
113        let state = {
114            let entry = self
115                .stats_cache
116                .entry(key)
117                .or_insert_with(|| Arc::new(Mutex::new(EpochStats::default())));
118            Arc::clone(entry.value())
119        };
120        let mut epoch_stats = state.lock().await;
121
122        let key: Vec<u8> = key.into();
123
124        let current_epoch = current_stat.node_epoch;
125        let refresh = if let Some(epoch) = epoch_stats.epoch() {
126            match current_stat.node_epoch.cmp(&epoch) {
127                Ordering::Greater => {
128                    // This node may have been redeployed.
129                    epoch_stats.clear_stats();
130                    epoch_stats.set_epoch(current_stat.node_epoch);
131                    epoch_stats.push_stat(current_stat);
132                    true
133                }
134                Ordering::Equal => {
135                    epoch_stats.push_stat(current_stat);
136                    false
137                }
138                Ordering::Less => {
139                    warn!("Ignore stale heartbeat: {:?}", current_stat);
140                    false
141                }
142            }
143        } else {
144            epoch_stats.set_epoch(current_stat.node_epoch);
145            epoch_stats.push_stat(current_stat);
146            // If the epoch is empty, it indicates that the current node sending the heartbeat
147            // for the first time to the current meta leader, so it is necessary to save
148            // the data to the KV store as soon as possible.
149            true
150        };
151
152        // Keep address updates and statistics in the original per-node order.
153        // Address acknowledgement is independent of the statistics epoch so a
154        // failed save can be retried by the next heartbeat of the same epoch.
155        if epoch_stats.epoch() == Some(current_epoch)
156            && let Some(stat) = epoch_stats.stats.last()
157        {
158            update_datanode_address(&self.address_updater, ctx, stat).await;
159        }
160
161        if !refresh && epoch_stats.len() < self.flush_stats_factor {
162            return Ok(HandleControl::Continue);
163        }
164
165        let value: Vec<u8> = DatanodeStatValue {
166            stats: epoch_stats.drain_all(),
167        }
168        .try_into()
169        .context(error::InvalidDatanodeStatFormatSnafu {})?;
170        let put = PutRequest {
171            key,
172            value,
173            prev_kv: false,
174        };
175
176        let _ = ctx
177            .in_memory
178            .put(put)
179            .await
180            .context(error::KvBackendSnafu)?;
181
182        Ok(HandleControl::Continue)
183    }
184}
185
186/// Applies datanode-specific address and table-cache updates. The caller keeps
187/// its statistics lock so notifications and statistics retain their ordering.
188async fn update_datanode_address(updater: &NodeAddressUpdater, ctx: &Context, stat: &Stat) {
189    let updated = updater
190        .update_if_needed((Role::Datanode, stat.id), stat.node_epoch, || async {
191            let peer = Peer { id: stat.id, addr: stat.addr.clone() };
192            let key = NodeAddressKey::with_datanode(stat.id).to_bytes();
193            utils::save_node_address(ctx, key, peer).await
194        })
195        .await
196        .inspect_err(|e| {
197            error!(e; "Failed to update datanode address, node: {}, address: {}", stat.id, stat.addr);
198        });
199    if let Ok(true) = updated {
200        let cache_idents = stat
201            .table_ids()
202            .into_iter()
203            .map(CacheIdent::TableId)
204            .collect::<Vec<_>>();
205        utils::invalidate_address_caches(ctx, stat.id, &cache_idents).await;
206    }
207}
208
209#[cfg(test)]
210mod tests {
211    use std::any::Any;
212    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering as AtomicOrdering};
213    use std::sync::{Arc, Mutex as StdMutex, mpsc};
214    use std::thread;
215    use std::time::Duration;
216
217    use common_meta::datanode::DatanodeStatKey;
218    use common_meta::error::{Error as MetaError, Result as MetaResult};
219    use common_meta::kv_backend::{KvBackend, KvBackendRef, ResettableKvBackend, TxnService};
220    use common_meta::rpc::store::{
221        BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse,
222        BatchPutRequest, BatchPutResponse, DeleteRangeRequest, DeleteRangeResponse, PutResponse,
223        RangeRequest, RangeResponse,
224    };
225    use tokio::sync::Semaphore;
226    use tokio::time::{sleep, timeout};
227
228    use super::*;
229    use crate::handler::test_utils::TestEnv;
230    use crate::service::store::cached_kv::LeaderCachedKvBackend;
231
232    struct ControlledKvBackend {
233        recorded_puts: StdMutex<Vec<PutRequest>>,
234        put_entered: Semaphore,
235        put_release: Semaphore,
236        block_next_put: AtomicBool,
237        delay_puts: AtomicBool,
238        active_puts: AtomicUsize,
239        max_active_puts: AtomicUsize,
240    }
241
242    impl ControlledKvBackend {
243        fn new() -> Self {
244            Self {
245                recorded_puts: StdMutex::new(Vec::new()),
246                put_entered: Semaphore::new(0),
247                put_release: Semaphore::new(0),
248                block_next_put: AtomicBool::new(false),
249                delay_puts: AtomicBool::new(false),
250                active_puts: AtomicUsize::new(0),
251                max_active_puts: AtomicUsize::new(0),
252            }
253        }
254
255        fn block_next_put(&self) {
256            self.block_next_put.store(true, AtomicOrdering::Relaxed);
257        }
258
259        async fn wait_for_blocked_put(&self) {
260            self.put_entered.acquire().await.unwrap().forget();
261        }
262
263        fn release_one_put(&self) {
264            self.put_release.add_permits(1);
265        }
266
267        fn set_delay_puts(&self, delay: bool) {
268            self.delay_puts.store(delay, AtomicOrdering::Relaxed);
269        }
270
271        fn clear_recorded_puts(&self) {
272            self.recorded_puts.lock().unwrap().clear();
273        }
274
275        fn recorded_puts(&self) -> Vec<PutRequest> {
276            self.recorded_puts.lock().unwrap().clone()
277        }
278
279        fn max_active_puts(&self) -> usize {
280            self.max_active_puts.load(AtomicOrdering::Relaxed)
281        }
282
283        fn start_put(&self, req: &PutRequest) -> PutGuard<'_> {
284            self.recorded_puts.lock().unwrap().push(req.clone());
285            let active = self.active_puts.fetch_add(1, AtomicOrdering::Relaxed) + 1;
286            self.max_active_puts
287                .fetch_max(active, AtomicOrdering::Relaxed);
288            PutGuard { backend: self }
289        }
290    }
291
292    struct PutGuard<'a> {
293        backend: &'a ControlledKvBackend,
294    }
295
296    impl Drop for PutGuard<'_> {
297        fn drop(&mut self) {
298            self.backend
299                .active_puts
300                .fetch_sub(1, AtomicOrdering::Relaxed);
301        }
302    }
303
304    #[async_trait::async_trait]
305    impl TxnService for ControlledKvBackend {
306        type Error = MetaError;
307    }
308
309    #[async_trait::async_trait]
310    impl KvBackend for ControlledKvBackend {
311        fn name(&self) -> &str {
312            "controlled"
313        }
314
315        fn as_any(&self) -> &dyn Any {
316            self
317        }
318
319        async fn range(&self, _req: RangeRequest) -> MetaResult<RangeResponse> {
320            unimplemented!()
321        }
322
323        async fn put(&self, req: PutRequest) -> MetaResult<PutResponse> {
324            let _guard = self.start_put(&req);
325            if self.block_next_put.swap(false, AtomicOrdering::Relaxed) {
326                self.put_entered.add_permits(1);
327                self.put_release.acquire().await.unwrap().forget();
328            }
329            if self.delay_puts.load(AtomicOrdering::Relaxed) {
330                sleep(Duration::from_millis(10)).await;
331            }
332            Ok(PutResponse::default())
333        }
334
335        async fn batch_put(&self, _req: BatchPutRequest) -> MetaResult<BatchPutResponse> {
336            unimplemented!()
337        }
338
339        async fn batch_get(&self, _req: BatchGetRequest) -> MetaResult<BatchGetResponse> {
340            unimplemented!()
341        }
342
343        async fn delete_range(&self, _req: DeleteRangeRequest) -> MetaResult<DeleteRangeResponse> {
344            unimplemented!()
345        }
346
347        async fn batch_delete(&self, _req: BatchDeleteRequest) -> MetaResult<BatchDeleteResponse> {
348            unimplemented!()
349        }
350    }
351
352    impl ResettableKvBackend for ControlledKvBackend {
353        fn reset(&self) {
354            self.clear_recorded_puts();
355        }
356
357        fn as_kv_backend_ref(self: Arc<Self>) -> KvBackendRef {
358            self
359        }
360    }
361
362    fn stat(node_id: u64, epoch: u64, marker: u64, addr: &str) -> Stat {
363        Stat {
364            timestamp_millis: marker as i64,
365            id: node_id,
366            addr: addr.to_string(),
367            region_num: marker,
368            node_epoch: epoch,
369            ..Default::default()
370        }
371    }
372
373    async fn handle_stat(
374        handler: Arc<CollectStatsHandler>,
375        mut ctx: Context,
376        stat: Stat,
377    ) -> Result<HandleControl> {
378        let mut acc = HeartbeatAccumulator {
379            stat: Some(stat),
380            ..Default::default()
381        };
382        handler
383            .handle(&HeartbeatRequest::default(), &mut ctx, &mut acc)
384            .await
385    }
386
387    fn use_controlled_address_backend(ctx: &mut Context) -> Arc<ControlledKvBackend> {
388        let backend = Arc::new(ControlledKvBackend::new());
389        ctx.leader_cached_kv_backend =
390            Arc::new(LeaderCachedKvBackend::with_always_leader(backend.clone()));
391        backend
392    }
393
394    #[tokio::test]
395    async fn test_handle_datanode_stats() {
396        let env = TestEnv::new();
397        let ctx = env.ctx();
398
399        let handler = CollectStatsHandler::default();
400        handle_request_many_times(ctx.clone(), &handler, 1).await;
401
402        let key = DatanodeStatKey { node_id: 101 };
403        let key: Vec<u8> = key.into();
404        let res = ctx.in_memory.get(&key).await.unwrap();
405        let kv = res.unwrap();
406        let key: DatanodeStatKey = kv.key.clone().try_into().unwrap();
407        assert_eq!(101, key.node_id);
408        let val: DatanodeStatValue = kv.value.try_into().unwrap();
409        // first new stat must be set in kv store immediately
410        assert_eq!(1, val.stats.len());
411        assert_eq!(1, val.stats[0].region_num);
412
413        handle_request_many_times(ctx.clone(), &handler, 10).await;
414
415        let key: Vec<u8> = key.into();
416        let res = ctx.in_memory.get(&key).await.unwrap();
417        let kv = res.unwrap();
418        let val: DatanodeStatValue = kv.value.try_into().unwrap();
419        assert_eq!(handler.flush_stats_factor, val.stats.len());
420    }
421
422    #[test]
423    fn test_same_datanode_wait_keeps_current_thread_runtime_responsive() {
424        let (backend_tx, backend_rx) = mpsc::sync_channel(1);
425        let (timer_tx, timer_rx) = mpsc::sync_channel(1);
426        let (done_tx, done_rx) = mpsc::sync_channel(1);
427
428        let worker = thread::spawn(move || {
429            tokio::runtime::Builder::new_current_thread()
430                .enable_time()
431                .build()
432                .unwrap()
433                .block_on(async move {
434                    let env = TestEnv::new();
435                    let mut ctx = env.ctx();
436                    let address_backend = use_controlled_address_backend(&mut ctx);
437                    address_backend.block_next_put();
438                    backend_tx.send(address_backend.clone()).unwrap();
439
440                    let handler = Arc::new(CollectStatsHandler::default());
441                    let first = tokio::spawn(handle_stat(
442                        handler.clone(),
443                        ctx.clone(),
444                        stat(101, 1, 1, "dn-101-v1"),
445                    ));
446                    address_backend.wait_for_blocked_put().await;
447
448                    let second =
449                        tokio::spawn(handle_stat(handler, ctx, stat(101, 1, 2, "dn-101-v1")));
450                    tokio::spawn(async move {
451                        sleep(Duration::from_millis(10)).await;
452                        timer_tx.send(()).unwrap();
453                    });
454                    tokio::task::yield_now().await;
455
456                    first.await.unwrap().unwrap();
457                    second.await.unwrap().unwrap();
458                    done_tx.send(()).unwrap();
459                });
460        });
461
462        let address_backend = backend_rx.recv_timeout(Duration::from_secs(1)).unwrap();
463        let timer_result = timer_rx.recv_timeout(Duration::from_secs(1));
464        address_backend.release_one_put();
465        timer_result.expect("the current-thread runtime must remain responsive");
466        done_rx
467            .recv_timeout(Duration::from_secs(1))
468            .expect("both heartbeat handlers must complete after releasing the write");
469        worker.join().unwrap();
470    }
471
472    #[tokio::test]
473    async fn test_concurrent_flush_persists_every_stat_once() {
474        let env = TestEnv::new();
475        let mut ctx = env.ctx();
476        let stats_backend = Arc::new(ControlledKvBackend::new());
477        ctx.in_memory = stats_backend.clone();
478        stats_backend.set_delay_puts(true);
479
480        let flush_stats_factor = 3;
481        let handler = Arc::new(CollectStatsHandler::new(Some(flush_stats_factor)));
482        handle_stat(handler.clone(), ctx.clone(), stat(101, 1, 0, "dn-101"))
483            .await
484            .unwrap();
485        stats_backend.clear_recorded_puts();
486
487        let mut tasks = Vec::with_capacity(2 * flush_stats_factor);
488        for marker in 1..=(2 * flush_stats_factor) {
489            tasks.push(tokio::spawn(handle_stat(
490                handler.clone(),
491                ctx.clone(),
492                stat(101, 1, marker as u64, "dn-101"),
493            )));
494        }
495        for task in tasks {
496            timeout(Duration::from_secs(1), task)
497                .await
498                .unwrap()
499                .unwrap()
500                .unwrap();
501        }
502
503        let puts = stats_backend.recorded_puts();
504        assert_eq!(2, puts.len());
505        let mut markers = puts
506            .into_iter()
507            .flat_map(|put| {
508                let value: DatanodeStatValue = put.value.try_into().unwrap();
509                value
510                    .stats
511                    .into_iter()
512                    .map(|stat| stat.region_num)
513                    .collect::<Vec<_>>()
514            })
515            .collect::<Vec<_>>();
516        markers.sort_unstable();
517        assert_eq!((1..=6).collect::<Vec<_>>(), markers);
518        assert_eq!(1, stats_backend.max_active_puts());
519    }
520
521    async fn handle_request_many_times(
522        mut ctx: Context,
523        handler: &CollectStatsHandler,
524        loop_times: i32,
525    ) {
526        let req = HeartbeatRequest::default();
527        for i in 1..=loop_times {
528            let mut acc = HeartbeatAccumulator {
529                stat: Some(Stat {
530                    id: 101,
531                    region_num: i as _,
532                    ..Default::default()
533                }),
534                ..Default::default()
535            };
536            handler.handle(&req, &mut ctx, &mut acc).await.unwrap();
537        }
538    }
539
540    #[tokio::test]
541    async fn test_persistence_failure_retries_but_notification_failure_does_not() {
542        use crate::handler::collect_stats_handler::CollectStatsHandler;
543        use crate::handler::test_utils::FailingCacheInvalidator;
544        use crate::handler::{HandleControl, HeartbeatAccumulator, HeartbeatHandler};
545        use crate::service::store::cached_kv::LeaderCachedKvBackend;
546        use api::v1::meta::HeartbeatRequest;
547        use common_meta::key::MetadataValue;
548        use common_meta::key::node_address::NodeAddressValue;
549        use common_meta::kv_backend::test_util::MockKvBackendBuilder;
550        use common_meta::rpc::store::PutResponse;
551        let mut ctx = TestEnv::new().ctx();
552        let puts = Arc::new(std::sync::Mutex::new(Vec::<PutRequest>::new()));
553        let backend = MockKvBackendBuilder::default()
554            .put_fn(Arc::new({
555                let puts = puts.clone();
556                move |req: PutRequest| {
557                    let mut puts = puts.lock().unwrap();
558                    puts.push(req);
559                    if puts.len() == 1 {
560                        return common_meta::error::UnexpectedSnafu {
561                            err_msg: "injected persistence failure",
562                        }
563                        .fail();
564                    }
565                    Ok(PutResponse::default())
566                }
567            }) as _)
568            .build()
569            .unwrap();
570        ctx.leader_cached_kv_backend =
571            Arc::new(LeaderCachedKvBackend::with_always_leader(Arc::new(backend)));
572        let invalidator = Arc::new(FailingCacheInvalidator::default());
573        ctx.cache_invalidator = invalidator.clone();
574        let handler = CollectStatsHandler::default();
575        let mut stat = Stat {
576            id: 101,
577            addr: "dn-original".into(),
578            node_epoch: 7,
579            ..Default::default()
580        };
581        let req = HeartbeatRequest::default();
582
583        // A backend error must not interrupt subsequent heartbeat handlers.
584        assert!(matches!(
585            handler
586                .handle(
587                    &req,
588                    &mut ctx,
589                    &mut HeartbeatAccumulator {
590                        stat: Some(stat.clone()),
591                        ..Default::default()
592                    }
593                )
594                .await
595                .unwrap(),
596            HandleControl::Continue
597        ));
598        assert_eq!(puts.lock().unwrap().len(), 1);
599        assert_eq!(invalidator.attempts(), 0);
600
601        let stats_key: Vec<u8> = common_meta::datanode::DatanodeStatKey { node_id: 101 }.into();
602        assert!(ctx.in_memory.get(&stats_key).await.unwrap().is_some());
603
604        // Same epoch retries persistence using the current heartbeat's address.
605        stat.addr = "dn-retry".into();
606        handler
607            .handle(
608                &req,
609                &mut ctx,
610                &mut HeartbeatAccumulator {
611                    stat: Some(stat.clone()),
612                    ..Default::default()
613                },
614            )
615            .await
616            .unwrap();
617        {
618            let puts = puts.lock().unwrap();
619            assert_eq!(puts.len(), 2);
620            assert_eq!(puts[1].key, NodeAddressKey::with_datanode(101).to_bytes());
621            assert_eq!(
622                NodeAddressValue::try_from_raw_value(&puts[1].value)
623                    .unwrap()
624                    .peer
625                    .addr,
626                "dn-retry"
627            );
628        }
629        assert_eq!(invalidator.attempts(), 1);
630
631        // Notification remains best-effort: its failure neither rewrites nor retries.
632        handler
633            .handle(
634                &req,
635                &mut ctx,
636                &mut HeartbeatAccumulator {
637                    stat: Some(stat.clone()),
638                    ..Default::default()
639                },
640            )
641            .await
642            .unwrap();
643        assert_eq!(puts.lock().unwrap().len(), 2);
644        assert_eq!(invalidator.attempts(), 1);
645
646        stat.node_epoch = 8;
647        handler
648            .handle(
649                &req,
650                &mut ctx,
651                &mut HeartbeatAccumulator {
652                    stat: Some(stat.clone()),
653                    ..Default::default()
654                },
655            )
656            .await
657            .unwrap();
658        assert_eq!(puts.lock().unwrap().len(), 3);
659        assert_eq!(invalidator.attempts(), 2);
660    }
661}
meta_srv/handler/collect_stats_handler.rs

meta_srv/handler/
collect_stats_handler.rs
