Skip to main content

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}