Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

use a simpler random 2-sampler #716

Merged
merged 2 commits into from
Jan 10, 2023
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 2 additions & 5 deletions tower/src/balance/p2c/service.rs
Original file line number Diff line number Diff line change
@@ -2,7 +2,7 @@ use super::super::error;
use crate::discover::{Change, Discover};
use crate::load::Load;
use crate::ready_cache::{error::Failed, ReadyCache};
use crate::util::rng::{sample_inplace, HasherRng, Rng};
use crate::util::rng::{sample_floyd2, HasherRng, Rng};
use futures_core::ready;
use futures_util::future::{self, TryFutureExt};
use pin_project_lite::pin_project;
@@ -185,10 +185,7 @@ where
len => {
// Get two distinct random indexes (in a random order) and
// compare the loads of the service at each index.
let idxs = sample_inplace(&mut self.rng, len as u32, 2);

let aidx = idxs[0];
let bidx = idxs[1];
let [aidx, bidx] = sample_floyd2(&mut self.rng, len as u64);
debug_assert_ne!(aidx, bidx, "random indices must be distinct");

let aload = self.ready_index_load(aidx as usize);
58 changes: 20 additions & 38 deletions tower/src/util/rng.rs
Original file line number Diff line number Diff line change
@@ -109,35 +109,19 @@ where
}
}

/// An inplace sampler borrowed from the Rand implementation for use internally
/// for the balance middleware.
/// ref: https://github.com/rust-random/rand/blob/b73640705d6714509f8ceccc49e8df996fa19f51/src/seq/index.rs#L425
/// A sampler modified from the Rand implementation for use internally for the balance middleware.
///
/// Docs from rand:
/// It's an implemenetation of Floyd's combination algorithm. with amount fixed at 2. This uses no allocated
/// memory and finishes in constant time (only 2 random calls)
///
/// Randomly sample exactly `amount` indices from `0..length`, using an inplace
/// partial Fisher-Yates method.
/// Sample an amount of indices using an inplace partial fisher yates method.
///
/// This allocates the entire `length` of indices and randomizes only the first `amount`.
/// It then truncates to `amount` and returns.
///
/// This method is not appropriate for large `length` and potentially uses a lot
/// of memory; because of this we only implement for `u32` index (which improves
/// performance in all cases).
///
/// Set-up is `O(length)` time and memory and shuffling is `O(amount)` time.
pub(crate) fn sample_inplace<R: Rng>(rng: &mut R, length: u32, amount: u32) -> Vec<u32> {
debug_assert!(amount <= length);
let mut indices: Vec<u32> = Vec::with_capacity(length as usize);
indices.extend(0..length);
for i in 0..amount {
let j: u64 = rng.next_range(i as u64..length as u64);
indices.swap(i as usize, j as usize);
}
indices.truncate(amount as usize);
debug_assert_eq!(indices.len(), amount as usize);
indices
/// ref: This was borrowed and modified from the following Rand implementation
/// https://github.com/rust-random/rand/blob/b73640705d6714509f8ceccc49e8df996fa19f51/src/seq/index.rs#L375-L411
pub(crate) fn sample_floyd2<R: Rng>(rng: &mut R, length: u64) -> [u64; 2] {
debug_assert!(2 <= length);
let aidx = rng.next_range(0..length - 1);
let bidx = rng.next_range(0..length);
let aidx = if aidx == bidx { length - 1 } else { aidx };
[aidx, bidx]
}

#[cfg(test)]
@@ -167,20 +151,18 @@ mod tests {
TestResult::from_bool(n >= range.start && (n < range.end || range.start == range.end))
}

fn sample_inplace(counter: u64, length: u32, amount: u32) -> TestResult {
if amount > length || length > 256 || amount > 32 {
fn sample_floyd2(counter: u64, length: u64) -> TestResult {
if length < 2 || length > 256 {
return TestResult::discard();
}

let mut rng = HasherRng::default();
rng.counter = counter;

let indxs = super::sample_inplace(&mut rng, length, amount);
let [a, b] = super::sample_floyd2(&mut rng, length);

for indx in indxs {
if indx > length {
return TestResult::failed();
}
if a >= length || b >= length || a == b {
return TestResult::failed();
}

TestResult::passed()
@@ -190,9 +172,9 @@ mod tests {
#[test]
fn sample_inplace_boundaries() {
LucioFranco marked this conversation as resolved.
Show resolved Hide resolved
let mut r = HasherRng::default();

assert_eq!(super::sample_inplace(&mut r, 0, 0).len(), 0);
assert_eq!(super::sample_inplace(&mut r, 1, 0).len(), 0);
assert_eq!(super::sample_inplace(&mut r, 1, 1), vec![0]);
match super::sample_floyd2(&mut r, 2) {
[0, 1] | [1, 0] => (),
err => panic!("{err:?}"),
}
}
}