flow/compute/
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::cell::RefCell;
16use std::collections::{BTreeMap, VecDeque};
17use std::rc::Rc;
18
19use dfir_rs::scheduled::SubgraphId;
20use dfir_rs::scheduled::graph::Dfir;
21use get_size2::GetSize;
22
23use crate::compute::types::ErrCollector;
24use crate::repr::{self, Timestamp};
25use crate::utils::{ArrangeHandler, Arrangement};
26
27/// input/output of a dataflow
28/// One `ComputeState` manage the input/output/schedule of one `Dfir`
29#[derive(Debug, Default)]
30pub struct DataflowState {
31    /// it is important to use a deque to maintain the order of subgraph here
32    /// TODO(discord9): consider dedup? Also not necessary for hydroflow itself also do dedup when schedule
33    schedule_subgraph: Rc<RefCell<BTreeMap<Timestamp, VecDeque<SubgraphId>>>>,
34    /// Frontier (in sys time) before which updates should not be emitted.
35    ///
36    /// We *must* apply it to sinks, to ensure correct outputs.
37    /// We *should* apply it to sources and imported shared state, because it improves performance.
38    /// Which means it's also the current time in temporal filter to get current correct result
39    as_of: Rc<RefCell<Timestamp>>,
40    /// error collector local to this `ComputeState`,
41    /// useful for distinguishing errors from different `Dfir`
42    err_collector: ErrCollector,
43    /// save all used arrange in this dataflow, since usually there is no delete operation
44    /// we can just keep track of all used arrange and schedule subgraph when they need to be updated
45    arrange_used: Vec<ArrangeHandler>,
46    /// the time arrangement need to be expired after a certain time in milliseconds
47    expire_after: Option<Timestamp>,
48    /// the last time each subgraph executed
49    last_exec_time: Option<Timestamp>,
50    /// the time the flow first executed, in unix timestamp milliseconds
51    start_time: Option<Timestamp>,
52}
53
54impl DataflowState {
55    pub fn new_arrange(&mut self, name: Option<Vec<String>>) -> ArrangeHandler {
56        let arrange = name.map(Arrangement::new_with_name).unwrap_or_default();
57
58        let arr = ArrangeHandler::from(arrange);
59        // mark this arrange as used in this dataflow
60        self.arrange_used.push(
61            arr.clone_future_only()
62                .expect("No write happening at this point"),
63        );
64        arr
65    }
66
67    /// schedule all subgraph that need to run with time <= `as_of` and run_available()
68    ///
69    /// return true if any subgraph actually executed
70    #[allow(clippy::swap_with_temporary)]
71    pub fn run_available_with_schedule(&mut self, df: &mut Dfir) -> bool {
72        // first split keys <= as_of into another map
73        let mut before = self
74            .schedule_subgraph
75            .borrow_mut()
76            .split_off(&(*self.as_of.borrow() + 1));
77        std::mem::swap(&mut before, &mut self.schedule_subgraph.borrow_mut());
78        for (_, v) in before {
79            for subgraph in v {
80                df.schedule_subgraph(subgraph);
81            }
82        }
83        df.run_available()
84    }
85    pub fn get_scheduler(&self) -> Scheduler {
86        Scheduler {
87            schedule_subgraph: self.schedule_subgraph.clone(),
88            cur_subgraph: Rc::new(RefCell::new(None)),
89        }
90    }
91
92    /// return a handle to the current time, will update when `as_of` is updated
93    ///
94    /// so it can keep track of the current time even in a closure that is called later
95    pub fn current_time_ref(&self) -> Rc<RefCell<Timestamp>> {
96        self.as_of.clone()
97    }
98
99    pub fn current_ts(&self) -> Timestamp {
100        *self.as_of.borrow()
101    }
102
103    pub fn set_current_ts(&mut self, ts: Timestamp) {
104        self.as_of.replace(ts);
105    }
106
107    pub fn get_err_collector(&self) -> ErrCollector {
108        self.err_collector.clone()
109    }
110
111    pub fn set_expire_after(&mut self, after: Option<repr::Duration>) {
112        self.expire_after = after;
113    }
114
115    pub fn expire_after(&self) -> Option<Timestamp> {
116        self.expire_after
117    }
118
119    pub fn get_state_size(&self) -> usize {
120        self.arrange_used.iter().map(|x| x.read().get_size()).sum()
121    }
122
123    pub fn set_last_exec_time(&mut self, time: Timestamp) {
124        self.last_exec_time = Some(time);
125        if self.start_time.is_none() {
126            // start_time is recorded at the completion of the first execution
127            // (post-execution), consistent with how last_exec_time is recorded.
128            self.start_time = Some(time);
129        }
130    }
131
132    pub fn last_exec_time(&self) -> Option<Timestamp> {
133        self.last_exec_time
134    }
135
136    /// Returns the time the flow first executed, in unix timestamp milliseconds.
137    pub fn start_time(&self) -> Option<Timestamp> {
138        self.start_time
139    }
140}
141
142#[derive(Debug, Clone)]
143pub struct Scheduler {
144    // this scheduler is shared with `DataflowState`, so it can schedule subgraph
145    schedule_subgraph: Rc<RefCell<BTreeMap<Timestamp, VecDeque<SubgraphId>>>>,
146    cur_subgraph: Rc<RefCell<Option<SubgraphId>>>,
147}
148
149impl Scheduler {
150    pub fn schedule_at(&self, next_run_time: Timestamp) {
151        let mut schedule_subgraph = self.schedule_subgraph.borrow_mut();
152        let subgraph = self.cur_subgraph.borrow();
153        let subgraph = subgraph.as_ref().expect("Set SubgraphId before schedule");
154        let subgraph_queue = schedule_subgraph.entry(next_run_time).or_default();
155        subgraph_queue.push_back(*subgraph);
156    }
157
158    pub fn schedule_for_arrange(&self, arrange: &Arrangement, now: Timestamp) {
159        if let Some(i) = arrange.get_next_update_time(&now) {
160            self.schedule_at(i)
161        }
162    }
163
164    pub fn set_cur_subgraph(&self, subgraph: SubgraphId) {
165        self.cur_subgraph.replace(Some(subgraph));
166    }
167}
flow/compute/state.rs

flow/compute/
state.rs
