common_meta/key/flow/
flow_state.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::collections::BTreeMap;
16use std::sync::Arc;
17
18use serde::{Deserialize, Serialize};
19use tokio::sync::Mutex;
20
21use crate::error::{self, Result};
22use crate::key::flow::FlowScoped;
23use crate::key::{FlowId, MetadataKey, MetadataValue};
24use crate::kv_backend::KvBackendRef;
25use crate::rpc::store::{PutRequest, RangeRequest};
26
27/// The entire FlowId to Flow Size's Map is stored directly in the value part of the key.
28pub const FLOW_STATE_KEY: &str = "state";
29
30/// The inner prefix (under `state/`) of the per-flownode flow state keys.
31pub const FLOW_STATE_NODE_KEY_PREFIX: &str = "node";
32
33/// The key of flow state.
34#[derive(Debug, Clone, Copy, PartialEq)]
35struct FlowStateKeyInner;
36
37impl FlowStateKeyInner {
38    pub fn new() -> Self {
39        Self
40    }
41}
42
43impl<'a> MetadataKey<'a, FlowStateKeyInner> for FlowStateKeyInner {
44    fn to_bytes(&self) -> Vec<u8> {
45        FLOW_STATE_KEY.as_bytes().to_vec()
46    }
47
48    fn from_bytes(bytes: &'a [u8]) -> Result<FlowStateKeyInner> {
49        let key = std::str::from_utf8(bytes).map_err(|e| {
50            error::InvalidMetadataSnafu {
51                err_msg: format!(
52                    "FlowInfoKeyInner '{}' is not a valid UTF8 string: {e}",
53                    String::from_utf8_lossy(bytes)
54                ),
55            }
56            .build()
57        })?;
58        if key != FLOW_STATE_KEY {
59            return Err(error::InvalidMetadataSnafu {
60                err_msg: format!("Invalid FlowStateKeyInner '{key}'"),
61            }
62            .build());
63        }
64        Ok(FlowStateKeyInner::new())
65    }
66}
67
68/// The key stores the state size of the flow.
69///
70/// The layout: `__flow/state`.
71pub struct FlowStateKey(FlowScoped<FlowStateKeyInner>);
72
73impl FlowStateKey {
74    /// Returns the [FlowStateKey].
75    pub fn new() -> FlowStateKey {
76        let inner = FlowStateKeyInner::new();
77        FlowStateKey(FlowScoped::new(inner))
78    }
79}
80
81impl Default for FlowStateKey {
82    fn default() -> Self {
83        Self::new()
84    }
85}
86
87impl<'a> MetadataKey<'a, FlowStateKey> for FlowStateKey {
88    fn to_bytes(&self) -> Vec<u8> {
89        self.0.to_bytes()
90    }
91
92    fn from_bytes(bytes: &'a [u8]) -> Result<FlowStateKey> {
93        Ok(FlowStateKey(FlowScoped::<FlowStateKeyInner>::from_bytes(
94            bytes,
95        )?))
96    }
97}
98
99/// The inner key of a per-flownode flow state entry: `state/node/{node_id}`.
100///
101/// `node_id` is the operator-configured flownode id (unique within the
102/// cluster; a flownode requires `node_id` in its config, see
103/// `src/cmd/src/flownode.rs`). It is the same value reported as
104/// `HeartbeatRequest.header.member_id` — the canonical identity metasrv uses
105/// for flownodes, see `get_node_id` in `src/meta-srv/src/service/heartbeat.rs`
106/// — and as `HeartbeatRequest.peer.id`.
107#[derive(Debug, Clone, PartialEq)]
108struct FlowStateNodeKeyInner {
109    node_id: u64,
110}
111
112impl FlowStateNodeKeyInner {
113    pub fn new(node_id: u64) -> Self {
114        Self { node_id }
115    }
116}
117
118impl<'a> MetadataKey<'a, FlowStateNodeKeyInner> for FlowStateNodeKeyInner {
119    fn to_bytes(&self) -> Vec<u8> {
120        format!(
121            "{FLOW_STATE_KEY}/{FLOW_STATE_NODE_KEY_PREFIX}/{}",
122            self.node_id
123        )
124        .into_bytes()
125    }
126
127    fn from_bytes(bytes: &'a [u8]) -> Result<FlowStateNodeKeyInner> {
128        let key = std::str::from_utf8(bytes).map_err(|e| {
129            error::InvalidMetadataSnafu {
130                err_msg: format!(
131                    "FlowStateNodeKeyInner '{}' is not a valid UTF8 string: {e}",
132                    String::from_utf8_lossy(bytes)
133                ),
134            }
135            .build()
136        })?;
137        let prefix = format!("{FLOW_STATE_KEY}/{FLOW_STATE_NODE_KEY_PREFIX}/");
138        let Some(node_id) = key.strip_prefix(&prefix) else {
139            return Err(error::InvalidMetadataSnafu {
140                err_msg: format!("Invalid FlowStateNodeKeyInner '{key}'"),
141            }
142            .build());
143        };
144        let node_id = node_id.parse::<u64>().map_err(|_| {
145            error::InvalidMetadataSnafu {
146                err_msg: format!("Invalid node id '{node_id}' in FlowStateNodeKeyInner '{key}'"),
147            }
148            .build()
149        })?;
150        Ok(FlowStateNodeKeyInner::new(node_id))
151    }
152}
153
154/// The key stores the per-flownode flow state report.
155///
156/// The layout: `__flow/state/node/{node_id}`.
157///
158/// Per-node keys live in the in-memory KV (same as the global `__flow/state`
159/// key), so they are automatically cleared when metasrv resets the in-memory
160/// KV on leader change; no separate cleanup is needed.
161pub struct FlowStateNodeKey(FlowScoped<FlowStateNodeKeyInner>);
162
163impl FlowStateNodeKey {
164    /// Returns the [FlowStateNodeKey] of the given node.
165    pub fn new(node_id: u64) -> FlowStateNodeKey {
166        FlowStateNodeKey(FlowScoped::new(FlowStateNodeKeyInner::new(node_id)))
167    }
168
169    /// Returns the full key prefix of all per-node flow state keys:
170    /// `__flow/state/node/`.
171    pub fn prefix() -> Vec<u8> {
172        format!(
173            "{}{FLOW_STATE_KEY}/{FLOW_STATE_NODE_KEY_PREFIX}/",
174            FlowScoped::<FlowStateNodeKeyInner>::PREFIX
175        )
176        .into_bytes()
177    }
178}
179
180impl<'a> MetadataKey<'a, FlowStateNodeKey> for FlowStateNodeKey {
181    fn to_bytes(&self) -> Vec<u8> {
182        self.0.to_bytes()
183    }
184
185    fn from_bytes(bytes: &'a [u8]) -> Result<FlowStateNodeKey> {
186        Ok(FlowStateNodeKey(
187            FlowScoped::<FlowStateNodeKeyInner>::from_bytes(bytes)?,
188        ))
189    }
190}
191
192/// The value of flow state size
193#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
194pub struct FlowStateValue {
195    /// For each key, the bytes of the state in memory
196    pub state_size: BTreeMap<FlowId, usize>,
197    /// For each key, the last execution time of flow in unix timestamp milliseconds.
198    pub last_exec_time_map: BTreeMap<FlowId, i64>,
199    /// For each flow, the time the flow first executed, in unix timestamp milliseconds.
200    /// TODO(#7987-followup): not yet propagated via the heartbeat wire format in distributed mode.
201    #[serde(default)]
202    pub start_time_map: BTreeMap<FlowId, i64>,
203}
204
205impl FlowStateValue {
206    pub fn new(
207        state_size: BTreeMap<FlowId, usize>,
208        last_exec_time_map: BTreeMap<FlowId, i64>,
209        start_time_map: BTreeMap<FlowId, i64>,
210    ) -> Self {
211        Self {
212            state_size,
213            last_exec_time_map,
214            start_time_map,
215        }
216    }
217}
218
219pub type FlowStateManagerRef = Arc<FlowStateManager>;
220
221/// The manager of [FlowStateKey]. Since state size changes frequently, we store it in memory.
222///
223/// This is only used in distributed mode. When meta-srv use heartbeat to update the flow stat report
224/// and frontned use get to get the latest flow stat report.
225///
226/// Per-flownode reports are stored under `__flow/state/node/{node_id}` keys in
227/// the in-memory KV (not in a separate in-process map), so a metasrv leader
228/// change — which resets the in-memory KV — automatically clears all per-node
229/// state without leaving stale entries behind.
230pub struct FlowStateManager {
231    in_memory: KvBackendRef,
232    /// Serializes the critical section of [`FlowStateManager::merge`]
233    /// (write per-node key -> scan -> aggregate -> write global key). It holds
234    /// no long-lived state; per-node reports live in the in-memory KV.
235    merge_lock: Mutex<()>,
236}
237
238impl FlowStateManager {
239    pub fn new(in_memory: KvBackendRef) -> Self {
240        Self {
241            in_memory,
242            merge_lock: Mutex::new(()),
243        }
244    }
245
246    pub async fn get(&self) -> Result<Option<FlowStateValue>> {
247        let key = FlowStateKey::new().to_bytes();
248        self.in_memory
249            .get(&key)
250            .await?
251            .map(|x| FlowStateValue::try_from_raw_value(&x.value))
252            .transpose()
253    }
254
255    pub async fn put(&self, value: FlowStateValue) -> Result<()> {
256        let key = FlowStateKey::new().to_bytes();
257        let value = value.try_as_raw_value()?;
258        let req = PutRequest::new().with_key(key).with_value(value);
259        self.in_memory.put(req).await?;
260        Ok(())
261    }
262
263    /// Merges a flow state report from a single flownode into the global view.
264    ///
265    /// `node_id` is the operator-configured flownode id (unique within the
266    /// cluster; reported as `HeartbeatRequest.header.member_id`, the canonical
267    /// metasrv identity of a flownode, and equal to `peer.id`).
268    ///
269    /// Reports are tracked per node under `__flow/state/node/{node_id}` in the
270    /// in-memory KV. A new report from the same node unconditionally replaces
271    /// that node's previous entry: reports are processed strictly in arrival
272    /// order, so no epoch comparison is made and a wall-clock rollback after a
273    /// node restart cannot permanently drop the node's later reports. The
274    /// global `FlowStateValue` is then aggregated over all per-node entries:
275    /// for the same flow, `last_exec_time_map` takes the max reported
276    /// timestamp and `state_size` takes the max reported size across nodes (a
277    /// flow normally runs on a single active flownode, so max is a safe
278    /// approximation). The aggregated value is written into the in-memory KV
279    /// under the global key `__flow/state` via the same path as `put`, keeping
280    /// `get()` behavior unchanged.
281    ///
282    /// Per-node keys are cleared automatically when the in-memory KV is reset
283    /// on a leader change. Known limitation: entries of dropped flows are not
284    /// proactively removed here. Once a flow is dropped its metadata is gone,
285    /// so the flows table join simply can't see the stale entry; it becomes
286    /// user-invisible until the owning flownode reports again (or stops
287    /// heartbeating forever, in which case the stale flow id lingers in the
288    /// aggregate but is never joined against any flow metadata).
289    pub async fn merge(&self, node_id: u64, incoming: FlowStateValue) -> Result<()> {
290        let _guard = self.merge_lock.lock().await;
291
292        // 1. Store this node's latest report under its per-node key.
293        let node_key = FlowStateNodeKey::new(node_id).to_bytes();
294        let value = incoming.try_as_raw_value()?;
295        let req = PutRequest::new().with_key(node_key).with_value(value);
296        self.in_memory.put(req).await?;
297
298        // 2. Read back every per-node report and aggregate them.
299        let req = RangeRequest::new().with_prefix(FlowStateNodeKey::prefix());
300        let resp = self.in_memory.range(req).await?;
301        let mut state_size = BTreeMap::new();
302        let mut last_exec_time_map = BTreeMap::new();
303        let mut start_time_map = BTreeMap::new();
304        for kv in resp.kvs {
305            let state = FlowStateValue::try_from_raw_value(&kv.value)?;
306            for (flow_id, size) in state.state_size {
307                state_size
308                    .entry(flow_id)
309                    .and_modify(|v: &mut usize| *v = (*v).max(size))
310                    .or_insert(size);
311            }
312            for (flow_id, ts) in state.last_exec_time_map {
313                last_exec_time_map
314                    .entry(flow_id)
315                    .and_modify(|v: &mut i64| *v = (*v).max(ts))
316                    .or_insert(ts);
317            }
318            for (flow_id, ts) in state.start_time_map {
319                start_time_map
320                    .entry(flow_id)
321                    .and_modify(|v: &mut i64| *v = (*v).max(ts))
322                    .or_insert(ts);
323            }
324        }
325
326        // 3. Write the aggregated value to the global key.
327        let aggregated = FlowStateValue::new(state_size, last_exec_time_map, start_time_map);
328        let key = FlowStateKey::new().to_bytes();
329        let value = aggregated.try_as_raw_value()?;
330        let req = PutRequest::new().with_key(key).with_value(value);
331        self.in_memory.put(req).await?;
332        Ok(())
333    }
334}
335
336/// Flow's state report, send regularly through heartbeat message
337#[derive(Debug, Clone, Default)]
338pub struct FlowStat {
339    /// For each key, the bytes of the state in memory
340    pub state_size: BTreeMap<u32, usize>,
341    /// For each key, the last execution time of flow in unix timestamp milliseconds.
342    pub last_exec_time_map: BTreeMap<FlowId, i64>,
343    /// For each flow, the time the flow first executed, in unix timestamp milliseconds.
344    /// TODO(#7987-followup): not yet propagated via the heartbeat wire format in distributed mode.
345    pub start_time_map: BTreeMap<FlowId, i64>,
346}
347
348impl From<FlowStateValue> for FlowStat {
349    fn from(value: FlowStateValue) -> Self {
350        Self {
351            state_size: value.state_size,
352            last_exec_time_map: value.last_exec_time_map,
353            start_time_map: value.start_time_map,
354        }
355    }
356}
357
358impl From<FlowStat> for FlowStateValue {
359    fn from(value: FlowStat) -> Self {
360        Self {
361            state_size: value.state_size,
362            last_exec_time_map: value.last_exec_time_map,
363            start_time_map: value.start_time_map,
364        }
365    }
366}
367
368#[cfg(test)]
369mod tests {
370    use std::collections::BTreeMap;
371
372    use crate::key::FlowId;
373    use crate::key::flow::flow_state::FlowStateValue;
374
375    #[test]
376    fn test_deserialize_legacy_flow_state_value() {
377        // Legacy format: only state_size and last_exec_time_map are present,
378        // without the start_time_map field added in PR #8392.
379        let legacy_json =
380            r#"{"state_size":{"1":1024,"2":2048},"last_exec_time_map":{"1":1700000000000}}"#;
381        let value: FlowStateValue = serde_json::from_str(legacy_json).unwrap();
382
383        let mut expected_state_size = BTreeMap::new();
384        expected_state_size.insert(FlowId::from(1u32), 1024usize);
385        expected_state_size.insert(FlowId::from(2u32), 2048usize);
386        assert_eq!(value.state_size, expected_state_size);
387
388        let mut expected_last_exec_time_map = BTreeMap::new();
389        expected_last_exec_time_map.insert(FlowId::from(1u32), 1700000000000i64);
390        assert_eq!(value.last_exec_time_map, expected_last_exec_time_map);
391
392        // serde(default) kicks in: old persisted data must not break,
393        // and the new field defaults to empty.
394        assert!(value.start_time_map.is_empty());
395    }
396
397    #[test]
398    fn test_flow_state_value_roundtrip_includes_start_time_map() {
399        let mut state_size = BTreeMap::new();
400        state_size.insert(FlowId::from(1u32), 1024usize);
401        let mut last_exec_time_map = BTreeMap::new();
402        last_exec_time_map.insert(FlowId::from(1u32), 1700000000000i64);
403        let mut start_time_map = BTreeMap::new();
404        start_time_map.insert(FlowId::from(1u32), 1700000000000i64);
405
406        let value = FlowStateValue {
407            state_size,
408            last_exec_time_map,
409            start_time_map,
410        };
411
412        let json = serde_json::to_string(&value).unwrap();
413        assert!(json.contains("start_time_map"));
414
415        let decoded: FlowStateValue = serde_json::from_str(&json).unwrap();
416        assert_eq!(decoded, value);
417    }
418
419    use std::sync::Arc;
420
421    use super::*;
422    use crate::kv_backend::memory::MemoryKvBackend;
423
424    fn state(last_exec_time_map: BTreeMap<FlowId, i64>) -> FlowStateValue {
425        FlowStateValue::new(BTreeMap::new(), last_exec_time_map, BTreeMap::new())
426    }
427
428    #[tokio::test]
429    async fn test_merge_keeps_reports_from_different_nodes() {
430        let manager = FlowStateManager::new(Arc::new(MemoryKvBackend::default()));
431
432        // Node A reports flow 1 executed at t1.
433        manager
434            .merge(1, state(BTreeMap::from([(1, 100)])))
435            .await
436            .unwrap();
437        // Node B reports flow 2 executed at t2. Before the per-node merge this
438        // would have wiped out node A's flow 1 entry.
439        manager
440            .merge(2, state(BTreeMap::from([(2, 200)])))
441            .await
442            .unwrap();
443        // Node A reports flow 1 executed at t3.
444        manager
445            .merge(1, state(BTreeMap::from([(1, 300)])))
446            .await
447            .unwrap();
448
449        let value = manager.get().await.unwrap().unwrap();
450        // flow 1 and flow 2 are both present, and flow 1 takes max(t1, t3).
451        assert_eq!(value.last_exec_time_map.get(&1), Some(&300));
452        assert_eq!(value.last_exec_time_map.get(&2), Some(&200));
453    }
454
455    #[tokio::test]
456    async fn test_merge_replaces_same_node_state() {
457        let manager = FlowStateManager::new(Arc::new(MemoryKvBackend::default()));
458
459        manager
460            .merge(1, state(BTreeMap::from([(1, 100)])))
461            .await
462            .unwrap();
463        // A new report from the same node replaces the previous one. There is
464        // no epoch comparison: arrival order alone decides, so this is also
465        // what a restarted node (new epoch) hits.
466        manager
467            .merge(1, state(BTreeMap::from([(2, 200)])))
468            .await
469            .unwrap();
470
471        let value = manager.get().await.unwrap().unwrap();
472        assert!(!value.last_exec_time_map.contains_key(&1));
473        assert_eq!(value.last_exec_time_map.get(&2), Some(&200));
474    }
475
476    #[tokio::test]
477    async fn test_merge_accepts_clock_rollback_from_same_node() {
478        let manager = FlowStateManager::new(Arc::new(MemoryKvBackend::default()));
479
480        // Simulates a flownode restart whose wall clock rolled back: the node
481        // first reports flow 1 at t1, then (after restart) reports a *smaller*
482        // timestamp t0. Because reports are processed strictly in arrival
483        // order (no epoch comparison), the later report must win instead of
484        // being permanently rejected.
485        manager
486            .merge(1, state(BTreeMap::from([(1, 100)])))
487            .await
488            .unwrap();
489        manager
490            .merge(1, state(BTreeMap::from([(1, 50)])))
491            .await
492            .unwrap();
493
494        let value = manager.get().await.unwrap().unwrap();
495        assert_eq!(value.last_exec_time_map.get(&1), Some(&50));
496    }
497
498    #[tokio::test]
499    async fn test_merge_aggregates_state_size_and_last_exec_time_by_max() {
500        let manager = FlowStateManager::new(Arc::new(MemoryKvBackend::default()));
501
502        // Both nodes report the same flow; the aggregate must take the max of
503        // state_size and last_exec_time_map across nodes.
504        manager
505            .merge(
506                1,
507                FlowStateValue::new(
508                    BTreeMap::from([(1, 1024)]),
509                    BTreeMap::from([(1, 100)]),
510                    BTreeMap::new(),
511                ),
512            )
513            .await
514            .unwrap();
515        manager
516            .merge(
517                2,
518                FlowStateValue::new(
519                    BTreeMap::from([(1, 2048)]),
520                    BTreeMap::from([(1, 50)]),
521                    BTreeMap::new(),
522                ),
523            )
524            .await
525            .unwrap();
526
527        let value = manager.get().await.unwrap().unwrap();
528        assert_eq!(value.state_size.get(&1), Some(&2048));
529        assert_eq!(value.last_exec_time_map.get(&1), Some(&100));
530    }
531
532    #[tokio::test]
533    async fn test_merge_concurrent_reports_no_lost_update() {
534        let manager = Arc::new(FlowStateManager::new(Arc::new(MemoryKvBackend::default())));
535
536        // Two nodes report concurrently; the merge lock must serialize the
537        // read-modify-write so neither node's report is lost.
538        let m1 = manager.clone();
539        let h1 = tokio::spawn(async move {
540            m1.merge(1, state(BTreeMap::from([(1, 100)])))
541                .await
542                .unwrap();
543        });
544        let m2 = manager.clone();
545        let h2 = tokio::spawn(async move {
546            m2.merge(2, state(BTreeMap::from([(2, 200)])))
547                .await
548                .unwrap();
549        });
550        h1.await.unwrap();
551        h2.await.unwrap();
552
553        let value = manager.get().await.unwrap().unwrap();
554        assert_eq!(value.last_exec_time_map.get(&1), Some(&100));
555        assert_eq!(value.last_exec_time_map.get(&2), Some(&200));
556    }
557
558    #[tokio::test]
559    async fn test_merge_empty_report_removes_own_flows_keeps_others() {
560        let manager = FlowStateManager::new(Arc::new(MemoryKvBackend::default()));
561
562        // Node A reports flow 1, node B reports flow 2.
563        manager
564            .merge(1, state(BTreeMap::from([(1, 100)])))
565            .await
566            .unwrap();
567        manager
568            .merge(2, state(BTreeMap::from([(2, 200)])))
569            .await
570            .unwrap();
571
572        // Node A reports an empty map: its own flow 1 disappears from the
573        // aggregate while node B's flow 2 is retained.
574        manager.merge(1, state(BTreeMap::new())).await.unwrap();
575
576        let value = manager.get().await.unwrap().unwrap();
577        assert!(!value.last_exec_time_map.contains_key(&1));
578        assert_eq!(value.last_exec_time_map.get(&2), Some(&200));
579    }
580
581    #[tokio::test]
582    async fn test_merge_state_cleared_on_in_memory_kv_reset() {
583        let backend = Arc::new(MemoryKvBackend::default());
584        let manager = FlowStateManager::new(backend.clone());
585
586        manager
587            .merge(1, state(BTreeMap::from([(1, 100)])))
588            .await
589            .unwrap();
590        assert!(manager.get().await.unwrap().is_some());
591
592        // Simulate a metasrv leader change, which clears the in-memory KV
593        // (including the per-node keys, since they live in the same KV).
594        backend.clear();
595        assert!(manager.get().await.unwrap().is_none());
596
597        // A fresh report works again after the reset.
598        manager
599            .merge(2, state(BTreeMap::from([(2, 200)])))
600            .await
601            .unwrap();
602        let value = manager.get().await.unwrap().unwrap();
603        assert_eq!(value.last_exec_time_map.get(&2), Some(&200));
604    }
605
606    #[tokio::test]
607    async fn test_merge_writes_global_key() {
608        let backend = Arc::new(MemoryKvBackend::default());
609        let manager = FlowStateManager::new(backend.clone());
610
611        manager
612            .merge(1, state(BTreeMap::from([(1, 100)])))
613            .await
614            .unwrap();
615
616        // The global key is still `__flow/state` and holds the serialized
617        // FlowStateValue, so existing get()/client readers are unchanged.
618        let dump = backend.dump();
619        let global_key = "__flow/state".as_bytes().to_vec();
620        assert!(dump.contains_key(&global_key));
621        let value = FlowStateValue::try_from_raw_value(dump.get(&global_key).unwrap()).unwrap();
622        assert_eq!(value.last_exec_time_map.get(&1), Some(&100));
623    }
624}
common_meta/key/flow/flow_state.rs

common_meta/key/flow/
flow_state.rs
