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    pub fn enable_leader_cache(&self) -> bool {
79        match &self {
80            State::Leader(leader) => leader.enable_leader_cache,
81            State::Follower(_) => false,
82        }
83    }
84
85    pub fn next_state<F>(&mut self, f: F)
86    where
87        F: FnOnce(&State) -> State,
88    {
89        *self = f(self);
90    }
91}
92
93pub fn become_leader(enable_leader_cache: bool) -> impl FnOnce(&State) -> State {
94    move |prev| match prev {
95        State::Leader(leader) => State::Leader(LeaderState { ..leader.clone() }),
96        State::Follower(follower) => State::Leader(LeaderState {
97            server_addr: follower.server_addr.to_string(),
98            enable_leader_cache,
99        }),
100    }
101}
102
103pub fn become_follower() -> impl FnOnce(&State) -> State {
104    move |prev| match prev {
105        State::Leader(leader) => State::Follower(FollowerState {
106            server_addr: leader.server_addr.to_string(),
107        }),
108        State::Follower(follower) => State::Follower(FollowerState { ..follower.clone() }),
109    }
110}
111
112#[cfg(test)]
113mod tests {
114    use std::assert_matches::assert_matches;
115
116    use crate::state::{become_follower, become_leader, FollowerState, LeaderState, State};
117
118    #[tokio::test]
119    async fn test_next_state() {
120        let mut state = State::follower("test".to_string());
121
122        state.next_state(become_leader(false));
123
124        assert_matches!(
125            state,
126            State::Leader(LeaderState {
127                enable_leader_cache: false,
128                ..
129            })
130        );
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_follower());
143
144        assert_matches!(state, State::Follower(FollowerState { .. }));
145
146        state.next_state(become_follower());
147
148        assert_matches!(state, State::Follower(FollowerState { .. }));
149    }
150}
meta_srv/state.rs

meta_srv/
state.rs
