meta_srv/
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::sync::{Arc, RwLock};
16
17pub type StateRef = Arc<RwLock<State>>;
18
19/// State transition.
20/// ```text
21///                     +------------------------------+
22///                     |                              |
23///                     |                              |
24///                     |                              |
25/// +-------------------v--------------------+         |
26/// | LeaderState{enable_leader_cache:false} |         |
27/// +-------------------+--------------------+         |
28///                     |                              |
29///                     |                              |
30///           +---------v---------+                    |
31///           | Init Leader Cache |                    |
32///           +---------+---------+                    |
33///                     |                              |
34///                     |                              |
35/// +-------------------v-------------------+          |
36/// | LeaderState{enable_leader_cache:true} |          |
37/// +-------------------+-------------------+          |
38///                     |                              |
39///                     |                              |
40///             +-------v-------+                      |
41///             | FollowerState |                      |
42///             +-------+-------+                      |
43///                     |                              |
44///                     |                              |
45///                     +------------------------------+
46///```
47#[derive(Debug, Clone)]
48pub enum State {
49    Leader(LeaderState),
50    Follower(FollowerState),
51}
52
53#[derive(Debug, Clone)]
54pub struct LeaderState {
55    // Disables the leader cache during initiation
56    pub enable_leader_cache: bool,
57
58    pub server_addr: String,
59}
60
61#[derive(Debug, Clone)]
62pub struct FollowerState {
63    pub server_addr: String,
64}
65
66impl State {
67    pub fn follower(server_addr: String) -> State {
68        Self::Follower(FollowerState { server_addr })
69    }
70
71    pub fn leader(server_addr: String, enable_leader_cache: bool) -> State {
72        Self::Leader(LeaderState {
73            enable_leader_cache,
74            server_addr,
75        })
76    }
77
78    /// Returns true if the current state is a leader.
79    pub fn is_leader(&self) -> bool {
80        matches!(self, State::Leader(_))
81    }
82
83    /// Returns true if the leader cache is enabled.
84    pub fn enable_leader_cache(&self) -> bool {
85        match &self {
86            State::Leader(leader) => leader.enable_leader_cache,
87            State::Follower(_) => false,
88        }
89    }
90
91    pub fn next_state<F>(&mut self, f: F)
92    where
93        F: FnOnce(&State) -> State,
94    {
95        *self = f(self);
96    }
97}
98
99pub fn become_leader(enable_leader_cache: bool) -> impl FnOnce(&State) -> State {
100    move |prev| match prev {
101        State::Leader(leader) => State::Leader(LeaderState { ..leader.clone() }),
102        State::Follower(follower) => State::Leader(LeaderState {
103            server_addr: follower.server_addr.clone(),
104            enable_leader_cache,
105        }),
106    }
107}
108
109pub fn become_follower() -> impl FnOnce(&State) -> State {
110    move |prev| match prev {
111        State::Leader(leader) => State::Follower(FollowerState {
112            server_addr: leader.server_addr.clone(),
113        }),
114        State::Follower(follower) => State::Follower(FollowerState { ..follower.clone() }),
115    }
116}
117
118#[cfg(test)]
119mod tests {
120    use std::assert_matches::assert_matches;
121
122    use crate::state::{FollowerState, LeaderState, State, become_follower, become_leader};
123
124    #[tokio::test]
125    async fn test_next_state() {
126        let mut state = State::follower("test".to_string());
127
128        state.next_state(become_leader(false));
129
130        assert_matches!(
131            state,
132            State::Leader(LeaderState {
133                enable_leader_cache: false,
134                ..
135            })
136        );
137
138        state.next_state(become_leader(false));
139
140        assert_matches!(
141            state,
142            State::Leader(LeaderState {
143                enable_leader_cache: false,
144                ..
145            })
146        );
147
148        state.next_state(become_follower());
149
150        assert_matches!(state, State::Follower(FollowerState { .. }));
151
152        state.next_state(become_follower());
153
154        assert_matches!(state, State::Follower(FollowerState { .. }));
155    }
156}
meta_srv/state.rs

meta_srv/
state.rs
