use chrono::{DateTime, Duration, NaiveDateTime, Utc};
use near_primitives_core::types::ProtocolVersion;
use std::{env, ops::RangeInclusive};
const NEAR_TESTS_PROTOCOL_UPGRADE_OVERRIDE: &str = "NEAR_TESTS_PROTOCOL_UPGRADE_OVERRIDE";
const PROTOCOL_UPGRADE_OVERRIDE_NOW: &str = "now";
const PROTOCOL_UPGRADE_OVERRIDE_SEQUENTIAL: &str = "sequential";
#[derive(thiserror::Error, Clone, Debug)]
pub enum ProtocolUpgradeVotingScheduleError {
#[error(
"The final upgrade must be the client protocol version! final version: {0}, client version: {1}"
)]
InvalidFinalUpgrade(ProtocolVersion, ProtocolVersion),
#[error("The upgrades must be sorted by datetime!")]
InvalidDateTimeOrder,
#[error("The upgrades must be sorted and increasing by one!")]
InvalidProtocolVersionOrder,
#[error("The environment override has an invalid format! Input: {0} Error: {1}")]
InvalidOverrideFormat(String, String),
}
type ProtocolUpgradeVotingScheduleRaw = Vec<(chrono::DateTime<Utc>, ProtocolVersion)>;
#[derive(Clone, Debug, PartialEq)]
pub struct ProtocolUpgradeVotingSchedule {
client_protocol_version: ProtocolVersion,
schedule: ProtocolUpgradeVotingScheduleRaw,
}
impl ProtocolUpgradeVotingSchedule {
pub fn new_immediate(client_protocol_version: ProtocolVersion) -> Self {
Self { client_protocol_version, schedule: vec![] }
}
pub fn new_sequential(
start: ProtocolVersion,
end: ProtocolVersion,
) -> ProtocolUpgradeVotingScheduleRaw {
RangeInclusive::new(start, end)
.enumerate()
.map(|(i, version)| {
let time = Utc::now() - Duration::minutes((end - start + 1) as i64 - i as i64);
(time, version)
})
.collect()
}
pub fn new_from_env_or_schedule(
min_supported_protocol_version: ProtocolVersion,
client_protocol_version: ProtocolVersion,
mut schedule: ProtocolUpgradeVotingScheduleRaw,
) -> Result<Self, ProtocolUpgradeVotingScheduleError> {
let env_override = env::var(NEAR_TESTS_PROTOCOL_UPGRADE_OVERRIDE).ok();
if let Some(env_override) = env_override {
schedule = Self::parse_override(
&env_override,
min_supported_protocol_version,
client_protocol_version,
)?;
tracing::warn!(
target: "protocol_upgrade",
?schedule,
"setting protocol upgrade override, this is fine in tests but should be avoided otherwise"
);
}
if let Some((_, version)) = schedule.last() {
if *version != client_protocol_version {
return Err(ProtocolUpgradeVotingScheduleError::InvalidFinalUpgrade(
*version,
client_protocol_version,
));
}
}
for i in 1..schedule.len() {
let prev_time = schedule[i - 1].0;
let this_time = schedule[i].0;
if !(prev_time < this_time) {
return Err(ProtocolUpgradeVotingScheduleError::InvalidDateTimeOrder);
}
}
for i in 1..schedule.len() {
let prev_protocol_version = schedule[i - 1].1;
let this_protocol_version = schedule[i].1;
if prev_protocol_version + 1 != this_protocol_version {
return Err(ProtocolUpgradeVotingScheduleError::InvalidProtocolVersionOrder);
}
}
tracing::debug!(target: "protocol_upgrade", ?schedule, "created protocol upgrade schedule");
Ok(Self { client_protocol_version, schedule })
}
#[cfg(feature = "clock")]
pub(crate) fn protocol_version_to_vote_for_at_date(
&self,
now: DateTime<Utc>,
next_epoch_protocol_version: ProtocolVersion,
) -> ProtocolVersion {
if next_epoch_protocol_version >= self.client_protocol_version {
return self.client_protocol_version;
}
if self.schedule.is_empty() {
return self.client_protocol_version;
}
let mut result = next_epoch_protocol_version;
for (time, version) in &self.schedule {
if now < *time {
break;
}
result = *version;
if *version > next_epoch_protocol_version {
break;
}
}
result
}
#[cfg(feature = "clock")]
pub fn protocol_version_to_vote_for(
&self,
current_time: near_time::Utc,
next_epoch_protocol_version: ProtocolVersion,
) -> ProtocolVersion {
let date_time = chrono::DateTime::from_timestamp(
current_time.unix_timestamp(),
current_time.nanosecond(),
);
self.protocol_version_to_vote_for_at_date(
date_time.unwrap_or_default(),
next_epoch_protocol_version,
)
}
pub fn schedule(&self) -> &Vec<(chrono::DateTime<Utc>, ProtocolVersion)> {
&self.schedule
}
pub fn parse_datetime(s: &str) -> Result<DateTime<Utc>, chrono::ParseError> {
let datetime = NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S")?;
let datetime = DateTime::<Utc>::from_naive_utc_and_offset(datetime, Utc);
Ok(datetime)
}
fn parse_override(
override_str: &str,
min_supported_protocol_version: ProtocolVersion,
client_protocol_version: ProtocolVersion,
) -> Result<ProtocolUpgradeVotingScheduleRaw, ProtocolUpgradeVotingScheduleError> {
match override_str.to_lowercase().as_str() {
PROTOCOL_UPGRADE_OVERRIDE_NOW => return Ok(vec![]),
PROTOCOL_UPGRADE_OVERRIDE_SEQUENTIAL => {
return Ok(Self::new_sequential(
min_supported_protocol_version + 1,
client_protocol_version,
));
}
_ => {}
}
let mut result = vec![];
let datetime_and_version_vec = override_str.split(',').collect::<Vec<_>>();
for datetime_and_version in datetime_and_version_vec {
let datetime_and_version = datetime_and_version.split('=').collect::<Vec<_>>();
let [datetime, version] = datetime_and_version[..] else {
let input = format!("{:?}", datetime_and_version);
let error = "The override must be in the format datetime=version!".to_string();
return Err(ProtocolUpgradeVotingScheduleError::InvalidOverrideFormat(
input, error,
));
};
let datetime = Self::parse_datetime(datetime).map_err(|err| {
ProtocolUpgradeVotingScheduleError::InvalidOverrideFormat(
datetime.to_string(),
err.to_string(),
)
})?;
let version = version.parse::<u32>().map_err(|err| {
ProtocolUpgradeVotingScheduleError::InvalidOverrideFormat(
version.to_string(),
err.to_string(),
)
})?;
result.push((datetime, version));
}
Ok(result)
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::vec;
fn make_simple_voting_schedule(
min_supported_protocol_version: ProtocolVersion,
client_protocol_version: u32,
datetime_str: &str,
) -> ProtocolUpgradeVotingSchedule {
make_voting_schedule(
min_supported_protocol_version,
client_protocol_version,
vec![(datetime_str, client_protocol_version)],
)
}
fn make_voting_schedule(
min_supported_protocol_version: ProtocolVersion,
client_protocol_version: ProtocolVersion,
schedule: Vec<(&str, u32)>,
) -> ProtocolUpgradeVotingSchedule {
make_voting_schedule_impl(schedule, min_supported_protocol_version, client_protocol_version)
.unwrap()
}
fn make_voting_schedule_impl(
schedule: Vec<(&str, u32)>,
min_supported_protocol_version: ProtocolVersion,
client_protocol_version: ProtocolVersion,
) -> Result<ProtocolUpgradeVotingSchedule, ProtocolUpgradeVotingScheduleError> {
let parse = |(datetime_str, version)| {
let datetime = ProtocolUpgradeVotingSchedule::parse_datetime(datetime_str);
let datetime = datetime.unwrap();
(datetime, version)
};
let schedule = schedule.into_iter().map(parse).collect::<Vec<_>>();
let schedule = ProtocolUpgradeVotingSchedule::new_from_env_or_schedule(
min_supported_protocol_version,
client_protocol_version,
schedule,
);
schedule
}
#[test]
fn test_default_upgrade_schedule() {
let now = chrono::Utc::now();
let client_protocol_version = 100;
let schedule = ProtocolUpgradeVotingSchedule::new_immediate(client_protocol_version);
assert_eq!(
client_protocol_version,
schedule.protocol_version_to_vote_for_at_date(now, client_protocol_version - 2)
);
assert_eq!(
client_protocol_version,
schedule.protocol_version_to_vote_for_at_date(now, client_protocol_version)
);
assert_eq!(
client_protocol_version,
schedule.protocol_version_to_vote_for_at_date(now, client_protocol_version + 2)
);
}
#[test]
fn test_before_scheduled_time() {
let now = chrono::Utc::now();
let client_protocol_version = 100;
let min_supported_protocol_version = 89;
let schedule = make_simple_voting_schedule(
min_supported_protocol_version,
client_protocol_version,
"3000-01-01 00:00:00",
);
let next_epoch_protocol_version = client_protocol_version - 2;
assert_eq!(
next_epoch_protocol_version,
schedule.protocol_version_to_vote_for_at_date(now, next_epoch_protocol_version)
);
let next_epoch_protocol_version = client_protocol_version;
assert_eq!(
next_epoch_protocol_version,
schedule.protocol_version_to_vote_for_at_date(now, next_epoch_protocol_version)
);
let next_epoch_protocol_version = client_protocol_version + 2;
assert_eq!(
client_protocol_version,
schedule.protocol_version_to_vote_for_at_date(now, next_epoch_protocol_version)
);
}
#[test]
fn test_after_scheduled_time() {
let now = chrono::Utc::now();
let client_protocol_version = 100;
let min_supported_protocol_version = 89;
let schedule = make_simple_voting_schedule(
min_supported_protocol_version,
client_protocol_version,
"1900-01-01 00:00:00",
);
assert_eq!(
client_protocol_version,
schedule.protocol_version_to_vote_for_at_date(now, client_protocol_version - 2)
);
assert_eq!(
client_protocol_version,
schedule.protocol_version_to_vote_for_at_date(now, client_protocol_version)
);
assert_eq!(
client_protocol_version,
schedule.protocol_version_to_vote_for_at_date(now, client_protocol_version + 2)
);
}
#[test]
fn test_double_upgrade_schedule() {
let current_protocol_version = 100;
let client_protocol_version = 102;
let min_supported_protocol_version = 89;
let schedule = vec![
("2000-01-10 00:00:00", current_protocol_version + 1),
("2000-01-15 00:00:00", current_protocol_version + 2),
];
let schedule =
make_voting_schedule(min_supported_protocol_version, client_protocol_version, schedule);
let now = ProtocolUpgradeVotingSchedule::parse_datetime("2000-01-05 00:00:00").unwrap();
assert_eq!(
current_protocol_version,
schedule.protocol_version_to_vote_for_at_date(now, current_protocol_version)
);
let now = ProtocolUpgradeVotingSchedule::parse_datetime("2000-01-10 00:00:00").unwrap();
assert_eq!(
current_protocol_version + 1,
schedule.protocol_version_to_vote_for_at_date(now, current_protocol_version)
);
let now = ProtocolUpgradeVotingSchedule::parse_datetime("2000-01-12 00:00:00").unwrap();
assert_eq!(
current_protocol_version + 1,
schedule.protocol_version_to_vote_for_at_date(now, current_protocol_version)
);
let now = ProtocolUpgradeVotingSchedule::parse_datetime("2000-01-15 00:00:00").unwrap();
assert_eq!(
current_protocol_version + 2,
schedule.protocol_version_to_vote_for_at_date(now, current_protocol_version + 1)
);
let now = ProtocolUpgradeVotingSchedule::parse_datetime("2000-01-20 00:00:00").unwrap();
assert_eq!(
current_protocol_version + 2,
schedule.protocol_version_to_vote_for_at_date(now, current_protocol_version + 1)
);
}
#[test]
fn test_upgrades_are_voted_one_at_a_time() {
let current_protocol_version = 100;
let client_protocol_version = 103;
let min_supported_protocol_version = 89;
let schedule = vec![
("2000-01-10 00:00:00", current_protocol_version + 1),
("2000-01-10 01:00:00", current_protocol_version + 2),
("2000-01-10 02:00:00", current_protocol_version + 3),
];
let schedule =
make_voting_schedule(min_supported_protocol_version, client_protocol_version, schedule);
let now = ProtocolUpgradeVotingSchedule::parse_datetime("2000-01-05 00:00:00").unwrap();
assert_eq!(
current_protocol_version,
schedule.protocol_version_to_vote_for_at_date(now, current_protocol_version)
);
let now = ProtocolUpgradeVotingSchedule::parse_datetime("2000-01-10 10:00:00").unwrap();
assert_eq!(
current_protocol_version + 1,
schedule.protocol_version_to_vote_for_at_date(now, current_protocol_version)
);
let now = ProtocolUpgradeVotingSchedule::parse_datetime("2000-01-10 10:00:00").unwrap();
assert_eq!(
current_protocol_version + 2,
schedule.protocol_version_to_vote_for_at_date(now, current_protocol_version + 1)
);
let now = ProtocolUpgradeVotingSchedule::parse_datetime("2000-01-10 10:00:00").unwrap();
assert_eq!(
current_protocol_version + 3,
schedule.protocol_version_to_vote_for_at_date(now, current_protocol_version + 2)
);
}
#[test]
fn test_errors() {
let client_protocol_version = 110;
let min_supported_protocol_version = 89;
let schedule = vec![("2000-01-10 00:00:00", 108), ("2000-01-15 00:00:00", 109)];
let schedule = make_voting_schedule_impl(
schedule,
min_supported_protocol_version,
client_protocol_version,
);
assert!(schedule.is_err());
let schedule = vec![("2000-01-10 00:00:00", 111), ("2000-01-15 00:00:00", 110)];
let schedule = make_voting_schedule_impl(
schedule,
min_supported_protocol_version,
client_protocol_version,
);
assert!(schedule.is_err());
let schedule = vec![("2000-01-10 00:00:00", 108), ("2000-01-15 00:00:00", 110)];
let schedule = make_voting_schedule_impl(
schedule,
min_supported_protocol_version,
client_protocol_version,
);
assert!(schedule.is_err());
let schedule = vec![("2000-01-15 00:00:00", 109), ("2000-01-10 00:00:00", 110)];
let schedule = make_voting_schedule_impl(
schedule,
min_supported_protocol_version,
client_protocol_version,
);
assert!(schedule.is_err());
}
#[test]
fn test_parse() {
assert!(ProtocolUpgradeVotingSchedule::parse_datetime("2001-02-03 23:59:59").is_ok());
assert!(ProtocolUpgradeVotingSchedule::parse_datetime("123").is_err());
}
#[test]
fn test_parse_override() {
let client_protocol_version = 103;
let min_supported_protocol_version = 100;
let datetime_str_1 = "2001-01-01 23:59:59";
let datetime_str_2 = "2001-01-02 23:59:59";
let datetime_str_3 = "2001-01-03 23:59:59";
let datetime_1 = ProtocolUpgradeVotingSchedule::parse_datetime(datetime_str_1).unwrap();
let datetime_2 = ProtocolUpgradeVotingSchedule::parse_datetime(datetime_str_2).unwrap();
let datetime_3 = ProtocolUpgradeVotingSchedule::parse_datetime(datetime_str_3).unwrap();
let datetime_version_str_1 = format!("{}={}", datetime_str_1, 101);
let datetime_version_str_2 = format!("{}={}", datetime_str_2, 102);
let datetime_version_str_3 = format!("{}={}", datetime_str_3, 103);
let override_str = PROTOCOL_UPGRADE_OVERRIDE_NOW;
let raw_schedule = ProtocolUpgradeVotingSchedule::parse_override(
override_str,
min_supported_protocol_version,
client_protocol_version,
)
.unwrap();
assert_eq!(raw_schedule.len(), 0);
let override_str = PROTOCOL_UPGRADE_OVERRIDE_SEQUENTIAL;
let raw_schedule = ProtocolUpgradeVotingSchedule::parse_override(
override_str,
min_supported_protocol_version,
client_protocol_version,
)
.unwrap();
assert_eq!(raw_schedule.len(), 3);
assert_eq!(raw_schedule[0].1, min_supported_protocol_version + 1);
assert_eq!(raw_schedule[1].1, min_supported_protocol_version + 2);
assert_eq!(raw_schedule[2].1, min_supported_protocol_version + 3);
let override_str = datetime_version_str_1.clone();
let raw_schedule = ProtocolUpgradeVotingSchedule::parse_override(
&override_str,
min_supported_protocol_version,
client_protocol_version,
)
.unwrap();
assert_eq!(raw_schedule.len(), 1);
assert_eq!(raw_schedule[0].0, datetime_1);
assert_eq!(raw_schedule[0].1, 101);
let override_str =
[datetime_version_str_1.clone(), datetime_version_str_2.clone()].join(",");
let raw_schedule = ProtocolUpgradeVotingSchedule::parse_override(
&override_str,
min_supported_protocol_version,
client_protocol_version,
)
.unwrap();
assert_eq!(raw_schedule.len(), 2);
assert_eq!(raw_schedule[0].0, datetime_1);
assert_eq!(raw_schedule[0].1, 101);
assert_eq!(raw_schedule[1].0, datetime_2);
assert_eq!(raw_schedule[1].1, 102);
let override_str =
[datetime_version_str_1, datetime_version_str_2, datetime_version_str_3].join(",");
let raw_schedule = ProtocolUpgradeVotingSchedule::parse_override(
&override_str,
min_supported_protocol_version,
client_protocol_version,
)
.unwrap();
assert_eq!(raw_schedule.len(), 3);
assert_eq!(raw_schedule[0].0, datetime_1);
assert_eq!(raw_schedule[0].1, 101);
assert_eq!(raw_schedule[1].0, datetime_2);
assert_eq!(raw_schedule[1].1, 102);
assert_eq!(raw_schedule[2].0, datetime_3);
assert_eq!(raw_schedule[2].1, 103);
}
#[test]
fn test_env_override() {
let client_protocol_version = 100;
let min_supported_protocol_version = 97;
unsafe {
std::env::set_var(NEAR_TESTS_PROTOCOL_UPGRADE_OVERRIDE, PROTOCOL_UPGRADE_OVERRIDE_NOW);
}
let schedule = make_simple_voting_schedule(
min_supported_protocol_version,
client_protocol_version,
"2999-02-03 23:59:59",
);
assert_eq!(schedule, ProtocolUpgradeVotingSchedule::new_immediate(client_protocol_version));
unsafe {
std::env::set_var(
NEAR_TESTS_PROTOCOL_UPGRADE_OVERRIDE,
PROTOCOL_UPGRADE_OVERRIDE_SEQUENTIAL,
);
}
let schedule = make_simple_voting_schedule(
min_supported_protocol_version,
client_protocol_version,
"2999-02-03 23:59:59",
);
assert_eq!(schedule.schedule().len(), 3);
assert_eq!(schedule.schedule()[0].1, min_supported_protocol_version + 1);
let mut date = schedule.schedule()[0].0;
assert_eq!(schedule.schedule()[1].1, min_supported_protocol_version + 2);
assert!(schedule.schedule()[1].0 > date);
date = schedule.schedule()[1].0;
assert_eq!(schedule.schedule()[2].1, min_supported_protocol_version + 3);
assert!(schedule.schedule()[2].0 > date);
assert!(schedule.schedule()[2].0 < Utc::now());
let datetime_override = "2000-01-01 23:59:59";
unsafe {
std::env::set_var(
NEAR_TESTS_PROTOCOL_UPGRADE_OVERRIDE,
format!("{}={}", datetime_override, client_protocol_version),
);
}
let schedule = make_simple_voting_schedule(
min_supported_protocol_version,
client_protocol_version,
"2999-02-03 23:59:59",
);
assert_eq!(
schedule.schedule()[0].0,
ProtocolUpgradeVotingSchedule::parse_datetime(datetime_override).unwrap()
);
assert_eq!(schedule.schedule()[0].1, client_protocol_version);
unsafe {
std::env::remove_var(NEAR_TESTS_PROTOCOL_UPGRADE_OVERRIDE);
}
}
}