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) => {
102            let mut new_leader = leader.clone();
103            new_leader.enable_leader_cache = enable_leader_cache;
104            State::Leader(new_leader)
105        }
106        State::Follower(follower) => State::Leader(LeaderState {
107            server_addr: follower.server_addr.clone(),
108            enable_leader_cache,
109        }),
110    }
111}
112
113pub fn become_follower() -> impl FnOnce(&State) -> State {
114    move |prev| match prev {
115        State::Leader(leader) => State::Follower(FollowerState {
116            server_addr: leader.server_addr.clone(),
117        }),
118        State::Follower(follower) => State::Follower(FollowerState { ..follower.clone() }),
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use std::assert_matches;
125
126    use crate::state::{FollowerState, LeaderState, State, become_follower, become_leader};
127
128    #[tokio::test]
129    async fn test_next_state() {
130        let mut state = State::follower("test".to_string());
131
132        state.next_state(become_leader(false));
133
134        assert_matches!(
135            state,
136            State::Leader(LeaderState {
137                enable_leader_cache: false,
138                ..
139            })
140        );
141
142        state.next_state(become_leader(false));
143
144        assert_matches!(
145            state,
146            State::Leader(LeaderState {
147                enable_leader_cache: false,
148                ..
149            })
150        );
151
152        state.next_state(become_follower());
153
154        assert_matches!(state, State::Follower(FollowerState { .. }));
155
156        state.next_state(become_follower());
157
158        assert_matches!(state, State::Follower(FollowerState { .. }));
159    }
160}
meta_srv/state.rs

meta_srv/
state.rs
