mito2/lease.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::LazyLock;
16use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
17use std::time::Instant;
18
19/// The instant lease deadlines count from, so that a deadline fits in an [AtomicU64].
20static LEASE_EPOCH: LazyLock<Instant> = LazyLock::new(Instant::now);
21
22fn lease_nanos(instant: Instant) -> u64 {
23 instant.saturating_duration_since(*LEASE_EPOCH).as_nanos() as u64
24}
25
26/// A lease deadline in nanoseconds since [LEASE_EPOCH]. It only moves forward; 0 means none.
27#[derive(Debug, Default)]
28pub(crate) struct Deadline(AtomicU64);
29
30impl Deadline {
31 pub(crate) fn extend(&self, deadline: Instant) {
32 self.0.fetch_max(lease_nanos(deadline), Ordering::Relaxed);
33 }
34
35 pub(crate) fn get(&self) -> u64 {
36 self.0.load(Ordering::Relaxed)
37 }
38}
39
40/// Lease of this node: the latest deadline of the heartbeat responses that carry region leases.
41#[derive(Debug, Default)]
42pub(crate) struct NodeLease {
43 /// Whether writes need a lease. Only a datanode the metasrv controls sets it.
44 enabled: AtomicBool,
45 deadline: Deadline,
46 /// The time of the lease checks while a test freezes it, in the unit of [Deadline]; 0 when
47 /// it is not frozen.
48 #[cfg(any(test, feature = "test"))]
49 frozen_now: AtomicU64,
50}
51
52impl NodeLease {
53 pub(crate) fn enable(&self) {
54 self.enabled.store(true, Ordering::Relaxed);
55 }
56
57 pub(crate) fn is_enabled(&self) -> bool {
58 self.enabled.load(Ordering::Relaxed)
59 }
60
61 pub(crate) fn renew(&self, deadline: Instant) {
62 self.deadline.extend(deadline);
63 }
64
65 /// Returns the current time in the unit of [Deadline].
66 pub(crate) fn now(&self) -> u64 {
67 #[cfg(any(test, feature = "test"))]
68 {
69 let frozen_now = self.frozen_now.load(Ordering::Relaxed);
70 if frozen_now != 0 {
71 return frozen_now;
72 }
73 }
74 lease_nanos(Instant::now())
75 }
76
77 #[cfg(any(test, feature = "test"))]
78 pub(crate) fn freeze_clock(&self, at: Instant) {
79 self.frozen_now
80 .store(lease_nanos(at).max(1), Ordering::Relaxed);
81 }
82
83 /// Lets a region created or upgraded on this node write until the node's deadline. The
84 /// metasrv granted that deadline before the creation or upgrade, so a later migration of the
85 /// region waits past it. Later responses do not extend it; the region's Leader grants do.
86 pub(crate) fn seed(&self, leader_deadline: &Deadline) {
87 leader_deadline
88 .0
89 .fetch_max(self.deadline.get(), Ordering::Relaxed);
90 }
91}