flow/
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 common_meta::key::flow::flow_state::FlowStat;
16use datatypes::value::Value;
17use tokio::sync::{mpsc, oneshot};
18
19use crate::error::InternalSnafu;
20
21pub fn get_value_heap_size(value: &Value) -> usize {
22    match value {
23        Value::Binary(v) => v.len(),
24        Value::String(v) => v.len(),
25        Value::List(v) => v.items().iter().map(get_value_heap_size).sum(),
26        _ => 0,
27    }
28}
29
30#[derive(Clone)]
31pub struct SizeReportSender {
32    inner: mpsc::Sender<oneshot::Sender<FlowStat>>,
33}
34impl SizeReportSender {
35    pub fn new() -> (Self, StateReportHandler) {
36        let (tx, rx) = mpsc::channel(1);
37        (Self { inner: tx }, rx)
38    }
39    pub async fn query(&self, timeout: std::time::Duration) -> crate::Result<FlowStat> {
40        let (tx, rx) = oneshot::channel();
41        self.inner.send(tx).await.map_err(|_| {
42            InternalSnafu {
43                reason: "failed to send size report request",
44            }
45            .build()
46        })?;
47        tokio::time::timeout(timeout, rx)
48            .await
49            .map_err(|_| {
50                InternalSnafu {
51                    reason: "failed to receive size report after timeout",
52                }
53                .build()
54            })?
55            .map_err(|_| {
56                InternalSnafu {
57                    reason: "size report sender dropped",
58                }
59                .build()
60            })
61    }
62}
63pub type StateReportHandler = mpsc::Receiver<oneshot::Sender<FlowStat>>;
flow/utils.rs

flow/
utils.rs
