third_party/raft-rs/tests/test_util/mod.rs - incubator-teaclave-sgx-sdk - Git at Google

 // Copyright 2018 PingCAP, Inc.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //     http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // See the License for the specific language governing permissions and
 // limitations under the License.

 // Copyright 2015 CoreOS, Inc.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //     http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.

 use env_logger;
 use protobuf::RepeatedField;
 use raft::eraftpb::*;
 use raft::storage::MemStorage;
 use raft::*;
 use rand;
 use std::collections::HashMap;
 use std::ops::*;

 /// Do any common test initialization. Eg set up logging, setup fail-rs.
 pub fn setup_for_test() {
     let _ = env_logger::try_init();
 }

 pub fn ltoa(raft_log: &RaftLog<MemStorage>) -> String {
     let mut s = format!("committed: {}\n", raft_log.committed);
     s = s + &format!("applied: {}\n", raft_log.applied);
     for (i, e) in raft_log.all_entries().iter().enumerate() {
         s = s + &format!("#{}: {:?}\n", i, e);
     }
     s
 }

 pub fn new_storage() -> MemStorage {
     MemStorage::new()
 }

 pub fn new_test_config(
     id: u64,
     peers: Vec<u64>,
     election_tick: usize,
     heartbeat_tick: usize,
 ) -> Config {
     Config {
         id,
         peers,
         election_tick,
         heartbeat_tick,
         max_size_per_msg: NO_LIMIT,
         max_inflight_msgs: 256,
         ..Default::default()
     }
 }

 /// Compare to upstream, we use struct instead of trait here.
 /// Because to be able to cast Interface later, we have to make
 /// Raft derive Any, which will require a lot of dependencies to derive Any.
 /// That's not worthy for just testing purpose.
 pub struct Interface {
     pub raft: Option<Raft<MemStorage>>,
 }

 impl Interface {
     pub fn new(r: Raft<MemStorage>) -> Interface {
         Interface { raft: Some(r) }
     }

     pub fn step(&mut self, m: Message) -> Result<()> {
         match self.raft {
             Some(_) => Raft::step(self, m),
             None => Ok(()),
         }
     }

     pub fn read_messages(&mut self) -> Vec<Message> {
         match self.raft {
             Some(_) => self.msgs.drain(..).collect(),
             None => vec![],
         }
     }

     fn initial(&mut self, id: u64, ids: &[u64]) {
         if self.raft.is_some() {
             self.id = id;
             let prs = self.take_prs();
             self.set_prs(ProgressSet::with_capacity(
                 ids.len(),
                 prs.learner_ids().len(),
             ));
             for id in ids {
                 if prs.learner_ids().contains(id) {
                     let progress = Progress {
                         is_learner: true,
                         ..Default::default()
                     };
                     if let Err(e) = self.mut_prs().insert_learner(*id, progress) {
                         panic!("{}", e);
                     }
                 } else {
                     let progress = Progress {
                         ..Default::default()
                     };
                     if let Err(e) = self.mut_prs().insert_voter(*id, progress) {
                         panic!("{}", e);
                     }
                 }
             }
             let term = self.term;
             self.reset(term);
         }
     }
 }

 impl Deref for Interface {
     type Target = Raft<MemStorage>;
     fn deref(&self) -> &Raft<MemStorage> {
         self.raft.as_ref().unwrap()
     }
 }

 impl DerefMut for Interface {
     fn deref_mut(&mut self) -> &mut Raft<MemStorage> {
         self.raft.as_mut().unwrap()
     }
 }

 pub fn new_test_raft(
     id: u64,
     peers: Vec<u64>,
     election: usize,
     heartbeat: usize,
     storage: MemStorage,
 ) -> Interface {
     Interface::new(Raft::new(
         &new_test_config(id, peers, election, heartbeat),
         storage,
     ))
 }

 pub fn new_test_raft_with_prevote(
     id: u64,
     peers: Vec<u64>,
     election: usize,
     heartbeat: usize,
     storage: MemStorage,
     pre_vote: bool,
 ) -> Interface {
     let mut config = new_test_config(id, peers, election, heartbeat);
     config.pre_vote = pre_vote;
     new_test_raft_with_config(&config, storage)
 }

 pub fn new_test_raft_with_config(config: &Config, storage: MemStorage) -> Interface {
     Interface::new(Raft::new(config, storage))
 }

 pub fn hard_state(t: u64, c: u64, v: u64) -> HardState {
     let mut hs = HardState::new();
     hs.set_term(t);
     hs.set_commit(c);
     hs.set_vote(v);
     hs
 }

 pub const SOME_DATA: Option<&'static str> = Some("somedata");

 pub fn new_message_with_entries(from: u64, to: u64, t: MessageType, ents: Vec<Entry>) -> Message {
     let mut m = Message::new();
     m.set_from(from);
     m.set_to(to);
     m.set_msg_type(t);
     if !ents.is_empty() {
         m.set_entries(RepeatedField::from_vec(ents));
     }
     m
 }

 pub fn new_message(from: u64, to: u64, t: MessageType, n: usize) -> Message {
     let mut m = new_message_with_entries(from, to, t, vec![]);
     if n > 0 {
         let mut ents = Vec::with_capacity(n);
         for _ in 0..n {
             ents.push(new_entry(0, 0, SOME_DATA));
         }
         m.set_entries(RepeatedField::from_vec(ents));
     }
     m
 }

 pub fn new_entry(term: u64, index: u64, data: Option<&str>) -> Entry {
     let mut e = Entry::new();
     e.set_index(index);
     e.set_term(term);
     if let Some(d) = data {
         e.set_data(d.as_bytes().to_vec());
     }
     e
 }

 pub fn empty_entry(term: u64, index: u64) -> Entry {
     new_entry(term, index, None)
 }

 pub fn new_snapshot(index: u64, term: u64, nodes: Vec<u64>) -> Snapshot {
     let mut s = Snapshot::new();
     s.mut_metadata().set_index(index);
     s.mut_metadata().set_term(term);
     s.mut_metadata().mut_conf_state().set_nodes(nodes);
     s
 }

 #[derive(Default, Debug, PartialEq, Eq, Hash)]
 struct Connem {
     from: u64,
     to: u64,
 }

 #[allow(declare_interior_mutable_const)]
 pub const NOP_STEPPER: Option<Interface> = Some(Interface { raft: None });

 #[derive(Default)]
 pub struct Network {
     pub peers: HashMap<u64, Interface>,
     pub storage: HashMap<u64, MemStorage>,
     dropm: HashMap<Connem, f64>,
     ignorem: HashMap<MessageType, bool>,
 }

 impl Network {
     // initializes a network from peers.
     // A nil node will be replaced with a new *stateMachine.
     // A *stateMachine will get its k, id.
     // When using stateMachine, the address list is always [1, n].
     pub fn new(peers: Vec<Option<Interface>>) -> Network {
         Network::new_with_config(peers, false)
     }

     // new_with_config is like new but sets the configuration pre_vote explicitly
     // for any state machines it creates.
     pub fn new_with_config(mut peers: Vec<Option<Interface>>, pre_vote: bool) -> Network {
         let size = peers.len();
         let peer_addrs: Vec<u64> = (1..size as u64 + 1).collect();
         let mut nstorage = HashMap::new();
         let mut npeers = HashMap::new();
         for (p, id) in peers.drain(..).zip(peer_addrs.clone()) {
             match p {
                 None => {
                     nstorage.insert(id, new_storage());
                     let r = new_test_raft_with_prevote(
                         id,
                         peer_addrs.clone(),
                         10,
                         1,
                         nstorage[&id].clone(),
                         pre_vote,
                     );
                     npeers.insert(id, r);
                 }
                 Some(mut p) => {
                     p.initial(id, &peer_addrs);
                     npeers.insert(id, p);
                 }
             }
         }
         Network {
             peers: npeers,
             storage: nstorage,
             ..Default::default()
         }
     }

     pub fn ignore(&mut self, t: MessageType) {
         self.ignorem.insert(t, true);
     }

     pub fn filter(&self, mut msgs: Vec<Message>) -> Vec<Message> {
         msgs.drain(..)
             .filter(|m| {
                 if self
                     .ignorem
                     .get(&m.get_msg_type())
                     .cloned()
                     .unwrap_or(false)
                 {
                     return false;
                 }
                 // hups never go over the network, so don't drop them but panic
                 assert_ne!(m.get_msg_type(), MessageType::MsgHup, "unexpected msgHup");
                 let perc = self
                     .dropm
                     .get(&Connem {
                         from: m.get_from(),
                         to: m.get_to(),
                     }).cloned()
                     .unwrap_or(0f64);
                 rand::random::<f64>() >= perc
             }).collect()
     }

     pub fn send(&mut self, msgs: Vec<Message>) {
         let mut msgs = msgs;
         while !msgs.is_empty() {
             let mut new_msgs = vec![];
             for m in msgs.drain(..) {
                 let resp = {
                     let p = self.peers.get_mut(&m.get_to()).unwrap();
                     let _ = p.step(m);
                     p.read_messages()
                 };
                 new_msgs.append(&mut self.filter(resp));
             }
             msgs.append(&mut new_msgs);
         }
     }

     pub fn drop(&mut self, from: u64, to: u64, perc: f64) {
         self.dropm.insert(Connem { from, to }, perc);
     }

     pub fn cut(&mut self, one: u64, other: u64) {
         self.drop(one, other, 1f64);
         self.drop(other, one, 1f64);
     }

     pub fn isolate(&mut self, id: u64) {
         for i in 0..self.peers.len() as u64 {
             let nid = i + 1;
             if nid != id {
                 self.drop(id, nid, 1.0);
                 self.drop(nid, id, 1.0);
             }
         }
     }

     pub fn recover(&mut self) {
         self.dropm = HashMap::new();
         self.ignorem = HashMap::new();
     }
 }
	// Copyright 2018 PingCAP, Inc.
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// See the License for the specific language governing permissions and
	// limitations under the License.

	// Copyright 2015 CoreOS, Inc.
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	use env_logger;
	use protobuf::RepeatedField;
	use raft::eraftpb::*;
	use raft::storage::MemStorage;
	use raft::*;
	use rand;
	use std::collections::HashMap;
	use std::ops::*;

	/// Do any common test initialization. Eg set up logging, setup fail-rs.
	pub fn setup_for_test() {
	let _ = env_logger::try_init();
	}

	pub fn ltoa(raft_log: &RaftLog<MemStorage>) -> String {
	let mut s = format!("committed: {}\n", raft_log.committed);
	s = s + &format!("applied: {}\n", raft_log.applied);
	for (i, e) in raft_log.all_entries().iter().enumerate() {
	s = s + &format!("#{}: {:?}\n", i, e);
	}
	s
	}

	pub fn new_storage() -> MemStorage {
	MemStorage::new()
	}

	pub fn new_test_config(
	id: u64,
	peers: Vec<u64>,
	election_tick: usize,
	heartbeat_tick: usize,
	) -> Config {
	Config {
	id,
	peers,
	election_tick,
	heartbeat_tick,
	max_size_per_msg: NO_LIMIT,
	max_inflight_msgs: 256,
	..Default::default()
	}
	}

	/// Compare to upstream, we use struct instead of trait here.
	/// Because to be able to cast Interface later, we have to make
	/// Raft derive Any, which will require a lot of dependencies to derive Any.
	/// That's not worthy for just testing purpose.
	pub struct Interface {
	pub raft: Option<Raft<MemStorage>>,
	}

	impl Interface {
	pub fn new(r: Raft<MemStorage>) -> Interface {
	Interface { raft: Some(r) }
	}

	pub fn step(&mut self, m: Message) -> Result<()> {
	match self.raft {
	Some(_) => Raft::step(self, m),
	None => Ok(()),
	}
	}

	pub fn read_messages(&mut self) -> Vec<Message> {
	match self.raft {
	Some(_) => self.msgs.drain(..).collect(),
	None => vec![],
	}
	}

	fn initial(&mut self, id: u64, ids: &[u64]) {
	if self.raft.is_some() {
	self.id = id;
	let prs = self.take_prs();
	self.set_prs(ProgressSet::with_capacity(
	ids.len(),
	prs.learner_ids().len(),
	));
	for id in ids {
	if prs.learner_ids().contains(id) {
	let progress = Progress {
	is_learner: true,
	..Default::default()
	};
	if let Err(e) = self.mut_prs().insert_learner(*id, progress) {
	panic!("{}", e);
	}
	} else {
	let progress = Progress {
	..Default::default()
	};
	if let Err(e) = self.mut_prs().insert_voter(*id, progress) {
	panic!("{}", e);
	}
	}
	}
	let term = self.term;
	self.reset(term);
	}
	}
	}

	impl Deref for Interface {
	type Target = Raft<MemStorage>;
	fn deref(&self) -> &Raft<MemStorage> {
	self.raft.as_ref().unwrap()
	}
	}

	impl DerefMut for Interface {
	fn deref_mut(&mut self) -> &mut Raft<MemStorage> {
	self.raft.as_mut().unwrap()
	}
	}

	pub fn new_test_raft(
	id: u64,
	peers: Vec<u64>,
	election: usize,
	heartbeat: usize,
	storage: MemStorage,
	) -> Interface {
	Interface::new(Raft::new(
	&new_test_config(id, peers, election, heartbeat),
	storage,
	))
	}

	pub fn new_test_raft_with_prevote(
	id: u64,
	peers: Vec<u64>,
	election: usize,
	heartbeat: usize,
	storage: MemStorage,
	pre_vote: bool,
	) -> Interface {
	let mut config = new_test_config(id, peers, election, heartbeat);
	config.pre_vote = pre_vote;
	new_test_raft_with_config(&config, storage)
	}

	pub fn new_test_raft_with_config(config: &Config, storage: MemStorage) -> Interface {
	Interface::new(Raft::new(config, storage))
	}

	pub fn hard_state(t: u64, c: u64, v: u64) -> HardState {
	let mut hs = HardState::new();
	hs.set_term(t);
	hs.set_commit(c);
	hs.set_vote(v);
	hs
	}

	pub const SOME_DATA: Option<&'static str> = Some("somedata");

	pub fn new_message_with_entries(from: u64, to: u64, t: MessageType, ents: Vec<Entry>) -> Message {
	let mut m = Message::new();
	m.set_from(from);
	m.set_to(to);
	m.set_msg_type(t);
	if !ents.is_empty() {
	m.set_entries(RepeatedField::from_vec(ents));
	}
	m
	}

	pub fn new_message(from: u64, to: u64, t: MessageType, n: usize) -> Message {
	let mut m = new_message_with_entries(from, to, t, vec![]);
	if n > 0 {
	let mut ents = Vec::with_capacity(n);
	for _ in 0..n {
	ents.push(new_entry(0, 0, SOME_DATA));
	}
	m.set_entries(RepeatedField::from_vec(ents));
	}
	m
	}

	pub fn new_entry(term: u64, index: u64, data: Option<&str>) -> Entry {
	let mut e = Entry::new();
	e.set_index(index);
	e.set_term(term);
	if let Some(d) = data {
	e.set_data(d.as_bytes().to_vec());
	}
	e
	}

	pub fn empty_entry(term: u64, index: u64) -> Entry {
	new_entry(term, index, None)
	}

	pub fn new_snapshot(index: u64, term: u64, nodes: Vec<u64>) -> Snapshot {
	let mut s = Snapshot::new();
	s.mut_metadata().set_index(index);
	s.mut_metadata().set_term(term);
	s.mut_metadata().mut_conf_state().set_nodes(nodes);
	s
	}

	#[derive(Default, Debug, PartialEq, Eq, Hash)]
	struct Connem {
	from: u64,
	to: u64,
	}

	#[allow(declare_interior_mutable_const)]
	pub const NOP_STEPPER: Option<Interface> = Some(Interface { raft: None });

	#[derive(Default)]
	pub struct Network {
	pub peers: HashMap<u64, Interface>,
	pub storage: HashMap<u64, MemStorage>,
	dropm: HashMap<Connem, f64>,
	ignorem: HashMap<MessageType, bool>,
	}

	impl Network {
	// initializes a network from peers.
	// A nil node will be replaced with a new *stateMachine.
	// A *stateMachine will get its k, id.
	// When using stateMachine, the address list is always [1, n].
	pub fn new(peers: Vec<Option<Interface>>) -> Network {
	Network::new_with_config(peers, false)
	}

	// new_with_config is like new but sets the configuration pre_vote explicitly
	// for any state machines it creates.
	pub fn new_with_config(mut peers: Vec<Option<Interface>>, pre_vote: bool) -> Network {
	let size = peers.len();
	let peer_addrs: Vec<u64> = (1..size as u64 + 1).collect();
	let mut nstorage = HashMap::new();
	let mut npeers = HashMap::new();
	for (p, id) in peers.drain(..).zip(peer_addrs.clone()) {
	match p {
	None => {
	nstorage.insert(id, new_storage());
	let r = new_test_raft_with_prevote(
	id,
	peer_addrs.clone(),
	10,
	1,
	nstorage[&id].clone(),
	pre_vote,
	);
	npeers.insert(id, r);
	}
	Some(mut p) => {
	p.initial(id, &peer_addrs);
	npeers.insert(id, p);
	}
	}
	}
	Network {
	peers: npeers,
	storage: nstorage,
	..Default::default()
	}
	}

	pub fn ignore(&mut self, t: MessageType) {
	self.ignorem.insert(t, true);
	}

	pub fn filter(&self, mut msgs: Vec<Message>) -> Vec<Message> {
	msgs.drain(..)
	.filter(\|m\| {
	if self
	.ignorem
	.get(&m.get_msg_type())
	.cloned()
	.unwrap_or(false)
	{
	return false;
	}
	// hups never go over the network, so don't drop them but panic
	assert_ne!(m.get_msg_type(), MessageType::MsgHup, "unexpected msgHup");
	let perc = self
	.dropm
	.get(&Connem {
	from: m.get_from(),
	to: m.get_to(),
	}).cloned()
	.unwrap_or(0f64);
	rand::random::<f64>() >= perc
	}).collect()
	}

	pub fn send(&mut self, msgs: Vec<Message>) {
	let mut msgs = msgs;
	while !msgs.is_empty() {
	let mut new_msgs = vec![];
	for m in msgs.drain(..) {
	let resp = {
	let p = self.peers.get_mut(&m.get_to()).unwrap();
	let _ = p.step(m);
	p.read_messages()
	};
	new_msgs.append(&mut self.filter(resp));
	}
	msgs.append(&mut new_msgs);
	}
	}

	pub fn drop(&mut self, from: u64, to: u64, perc: f64) {
	self.dropm.insert(Connem { from, to }, perc);
	}

	pub fn cut(&mut self, one: u64, other: u64) {
	self.drop(one, other, 1f64);
	self.drop(other, one, 1f64);
	}

	pub fn isolate(&mut self, id: u64) {
	for i in 0..self.peers.len() as u64 {
	let nid = i + 1;
	if nid != id {
	self.drop(id, nid, 1.0);
	self.drop(nid, id, 1.0);
	}
	}
	}

	pub fn recover(&mut self) {
	self.dropm = HashMap::new();
	self.ignorem = HashMap::new();
	}
	}