Skip to content

Commit

Permalink
feat: add support for --include-partial-seats (#12)
Browse files Browse the repository at this point in the history
* feat: add support for --include-partial-seats

* Represent absence of partial seats by emtpy vec

Instead of representing it by `None`.
This simplifies code as it allows to remove an `Option` in several
places.
  • Loading branch information
mooori authored Nov 13, 2023
1 parent 5409c5d commit 918bc59
Show file tree
Hide file tree
Showing 11 changed files with 601 additions and 24 deletions.
6 changes: 4 additions & 2 deletions Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,11 @@ run:
cargo run -p sim-validator-assignment -- \
run \
--num-blocks 1000 \
--num-shards 1 \
--num-shards 4 \
--seats-per-shard 250 \
--stake-per-seat 1 \
--max-malicious-stake-per-shard 1/2
--max-malicious-stake-per-shard 1/2 \
--include-partial-seats

.PHONY: download
download:
Expand All @@ -28,5 +29,6 @@ run-with:
--seats-per-shard 250 \
--stake-per-seat 50000000000000000000000000000 \
--max-malicious-stake-per-shard 2/3 \
--include-partial-seats \
--validator-data ./validator_data.json

107 changes: 96 additions & 11 deletions sim-validator-assignment/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ use num_rational::Ratio;
use serde::Serialize;
use std::path::PathBuf;

use crate::seat::Seat;
use crate::{partial_seat::PartialSeat, seat::Seat};

#[derive(Args, Serialize, Debug)]
pub struct Config {
Expand All @@ -30,18 +30,24 @@ pub struct Config {
/// data will be used in the simulation.
#[arg(long)]
pub validator_data: Option<PathBuf>,
/// A validator's stake might not entirely cover seats given a particular `stake_per_seat`. This
/// option controls whether remaining stake (not covering a full seat) should be assigned to a
/// partial seat or ignored.
#[arg(long, default_value_t = false)]
pub include_partial_seats: bool,
}

impl Config {
#[cfg(test)]
pub fn new_mock() -> Self {
pub fn new_mock(include_partial_seats: bool) -> Self {
Self {
num_blocks: 1_000,
num_shards: 4,
seats_per_shard: 2,
stake_per_seat: 100,
max_malicious_stake_per_shard: Ratio::new(1, 3),
validator_data: None,
include_partial_seats,
}
}

Expand Down Expand Up @@ -93,31 +99,70 @@ impl Config {

Ok(shard_seats)
}

/// Collect partials seats for `shard_idx` by picking seats from positions with `position %
/// num_shards == shard_idx`.
///
/// # Motivation for assignment via modulo
///
/// Every validator may hold at most 1 partial seat. If a validator's stake covers full seats
/// without remainder, it holds 0 partial seats. Hence `0 <= num_partial_seats <=
/// num_validators`. Assigning seats to shards with `%` allows distributing the existing number
/// of partial seats to shards as evenly as possible.
///
/// # No minimum number of required _partial_ seats per shard
///
/// It is not needed since partial seats are included only to distribute leftover stake (not
/// covering full seats) to shards.
pub fn collect_partial_seats_for_shard<'seats>(
&self,
shard_idx: usize,
partial_seats: &'seats [PartialSeat],
) -> anyhow::Result<Vec<&PartialSeat<'seats>>> {
if shard_idx >= usize::from(self.num_shards) {
anyhow::bail!(
"shard_idx {} is an invalid index for {} shards",
shard_idx,
self.num_shards
)
}

let mut shard_partial_seats = vec![];
for idx in (shard_idx..partial_seats.len()).step_by(self.num_shards.into()) {
shard_partial_seats.push(&partial_seats[idx]);
}

Ok(shard_partial_seats)
}
}

#[cfg(test)]
mod tests {
use std::collections::BTreeMap;

use super::Config;
use crate::validator::tests::new_test_raw_validator_data;
use crate::validator::{new_ordered_seats, parse_raw_validator_data};
use crate::validator::{
new_ordered_partial_seats, new_ordered_seats, parse_raw_validator_data,
};

#[test]
pub fn test_seats_per_stake() {
let config = Config::new_mock();
fn test_seats_per_stake() {
let config = Config::new_mock(false);
assert_eq!(config.seats_per_stake(20), 0);
assert_eq!(config.seats_per_stake(100), 1);
assert_eq!(config.seats_per_stake(530), 5);
}

#[test]
pub fn test_total_seats() {
let config = Config::new_mock();
fn test_total_seats() {
let config = Config::new_mock(false);
assert_eq!(config.total_seats(), 8);
}

#[test]
pub fn test_collect_seats_for_shard() {
let config = Config::new_mock();
fn test_collect_seats_for_shard() {
let config = Config::new_mock(false);
let (_, validators) = parse_raw_validator_data(&config, &new_test_raw_validator_data());
// Using ordered seats as input to have a deterministic result of `collect_seats_for_shard`.
let seats = new_ordered_seats(&validators);
Expand All @@ -137,12 +182,52 @@ mod tests {
}

#[test]
pub fn test_collect_seats_for_shard_errors() {
let config = Config::new_mock();
fn test_collect_seats_for_shard_errors() {
let config = Config::new_mock(false);
let (_, validators) = parse_raw_validator_data(&config, &new_test_raw_validator_data());
let seats = new_ordered_seats(&validators);

insta::assert_debug_snapshot!(config.collect_seats_for_shard(4, &seats));
insta::assert_debug_snapshot!(config.collect_seats_for_shard(0, &[]));
}

#[test]
fn test_collect_partial_seats_for_shard() {
let mut config = Config::new_mock(true);
config.stake_per_seat = 90;
let (_, validators) = parse_raw_validator_data(&config, &new_test_raw_validator_data());
// Using ordered partial seats as input to have a deterministic result of
// `collect_partial_seats_for_shard`.
let partial_seats = new_ordered_partial_seats(&validators, config.stake_per_seat);

// Using `BTreeMap` for deterministic ordering of keys.
let mut assignments = BTreeMap::new();
for shard_idx in 0..config.num_shards {
let assignment = config
.collect_partial_seats_for_shard(shard_idx.into(), &partial_seats)
.unwrap();
assignments.insert(format!("shard_{shard_idx}"), assignment);
}

insta::with_settings!({
info => &(
&config,
"partial_seats:",
&partial_seats
)
}, {
insta::assert_yaml_snapshot!(assignments);
})
}

#[test]
fn test_collect_partial_seats_for_shard_errors() {
let config = Config::new_mock(true);
let (_, validators) = parse_raw_validator_data(&config, &new_test_raw_validator_data());
let partial_seats = new_ordered_partial_seats(&validators, config.stake_per_seat);

insta::assert_debug_snapshot!(
config.collect_partial_seats_for_shard(config.num_shards.into(), &partial_seats)
);
}
}
1 change: 1 addition & 0 deletions sim-validator-assignment/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ use config::Config;
mod download;
use download::{download, DownloadConfig};
mod mocks;
mod partial_seat;
mod run;
mod seat;
mod shard;
Expand Down
49 changes: 49 additions & 0 deletions sim-validator-assignment/src/partial_seat.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
use serde::Serialize;

use crate::validator::Validator;

/// Represents a partial seat filled by a particular validator.
///
/// A partial seat may not outlive the validator it is referrencing.
#[derive(Serialize, PartialEq, Debug, Clone)]
pub struct PartialSeat<'validator> {
/// Reference to the validator holding this particular partial seat.
validator: &'validator Validator,
/// The stake attributed to the partial seat. Note that `0 < weight < stake_per_seat`.
///
/// For example, let validator `V` have a stake of 12 and let `stake_per_seat = 5`. Then `V`
/// holds 2 full seats and a partial seat with `weight = 2`.
weight: u128,
}

impl<'validator> PartialSeat<'validator> {
/// Constructs the partial seat filled by `validator`.
pub fn new(validator: &'validator Validator, weight: u128) -> Self {
Self { validator, weight }
}

pub fn get_is_malicious(&self) -> bool {
self.validator.get_is_malicious()
}

pub fn get_weight(&self) -> u128 {
self.weight
}
}

pub struct ShuffledPartialSeats<'seats> {
partial_seats: &'seats [PartialSeat<'seats>],
}

impl<'seats> ShuffledPartialSeats<'seats> {
/// Shuffles the input `partial_seats`.
// TODO(rand) do all shuffling operations in one generic fn
pub fn new(partial_seats: &'seats mut [PartialSeat<'seats>]) -> Self {
fastrand::shuffle(partial_seats);
Self { partial_seats }
}

pub fn get_partial_seats(&self) -> &[PartialSeat] {
self.partial_seats
}
}
26 changes: 20 additions & 6 deletions sim-validator-assignment/src/run.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,10 @@
use crate::config::Config;
use crate::partial_seat::ShuffledPartialSeats;
use crate::seat::ShuffledSeats;
use crate::shard::Shard;
use crate::validator::{new_ordered_seats, parse_raw_validator_data, RawValidatorData};
use crate::validator::{
new_ordered_partial_seats, new_ordered_seats, parse_raw_validator_data, RawValidatorData,
};
use num_rational::Ratio;
use num_traits::ToPrimitive;
use std::fs::read_to_string;
Expand All @@ -13,7 +16,7 @@ pub fn run(config: &Config) -> anyhow::Result<()> {
None => mock_validator_data(),
};

let (population_stats, validators) = parse_raw_validator_data(&config, &raw_validator_data);
let (population_stats, validators) = parse_raw_validator_data(config, &raw_validator_data);

println!("population_stats: {:?}", population_stats);
println!(
Expand All @@ -40,15 +43,26 @@ pub fn run(config: &Config) -> anyhow::Result<()> {
let mut num_corrupted_shards = 0;

for block_height in 0..config.num_blocks {
let mut ordered_seats = new_ordered_seats(&validators);
let shuffled_seats = ShuffledSeats::new(&mut ordered_seats);
let mut seats = new_ordered_seats(&validators);
let shuffled_seats = ShuffledSeats::new(&mut seats);

let mut partial_seats = if config.include_partial_seats {
new_ordered_partial_seats(&validators, config.stake_per_seat)
} else {
Vec::new()
};
let shuffled_partial_seats = ShuffledPartialSeats::new(&mut partial_seats);

for shard_idx in 0..config.num_shards {
let shard_idx = usize::from(shard_idx);
let shard_seats =
config.collect_seats_for_shard(shard_idx, shuffled_seats.get_seats())?;
let shard = Shard::new(&config, shard_seats)?;
if shard.is_corrupted(&config) {
let shard_partial_seats = config.collect_partial_seats_for_shard(
shard_idx,
shuffled_partial_seats.get_partial_seats(),
)?;
let shard = Shard::new(config, shard_seats, shard_partial_seats)?;
if shard.is_corrupted(config) {
num_corrupted_shards += 1;
}
}
Expand Down
3 changes: 2 additions & 1 deletion sim-validator-assignment/src/seat.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,12 +28,13 @@ pub struct ShuffledSeats<'seats> {

impl<'seats> ShuffledSeats<'seats> {
/// Shuffles the input `seats`.
// TODO(rand) do all shuffling operations in one generic fn
pub fn new(seats: &'seats mut [Seat<'seats>]) -> Self {
fastrand::shuffle(seats);
Self { seats }
}

pub fn get_seats(&self) -> &[Seat] {
&self.seats
self.seats
}
}
22 changes: 21 additions & 1 deletion sim-validator-assignment/src/shard.rs
Original file line number Diff line number Diff line change
@@ -1,17 +1,27 @@
use crate::config::Config;
use crate::partial_seat::PartialSeat;
use crate::seat::Seat;
use num_rational::Ratio;

#[derive(Debug, Default)]
pub struct Shard<'seats> {
seats: Vec<&'seats Seat<'seats>>,
partial_seats: Vec<&'seats PartialSeat<'seats>>,
stake: u128,
malicious_stake: u128,
}

impl<'seats> Shard<'seats> {
pub fn new(config: &Config, seats: Vec<&'seats Seat>) -> anyhow::Result<Self> {
pub fn new(
config: &Config,
seats: Vec<&'seats Seat>,
partial_seats: Vec<&'seats PartialSeat>,
) -> anyhow::Result<Self> {
if seats.len() != usize::try_from(config.seats_per_shard).unwrap() {
// Count only _full_ seats for the minimum number of required seats, since it is not
// clear how a _partial_ seat should be weighted for that concern.
// Validator assignment frameworks might try to minimize the number of partial seats or
// try to ignore them entirely.
anyhow::bail!(
"Shard requires {} seats, received {} seats",
config.seats_per_shard,
Expand All @@ -26,7 +36,17 @@ impl<'seats> Shard<'seats> {
shard.malicious_stake += config.stake_per_seat;
}
}

for ps in partial_seats.iter() {
let weight = ps.get_weight();
shard.stake += weight;
if ps.get_is_malicious() {
shard.malicious_stake += weight;
}
}

shard.seats = seats;
shard.partial_seats = partial_seats;
Ok(shard)
}

Expand Down
Loading

0 comments on commit 918bc59

Please sign in to comment.