use std::cmp::Reverse;
use std::collections::BinaryHeap;
use std::collections::HashMap;
use std::hash::Hash;
use crate::Context;
use crate::merge::MergePolicy;
use crate::relay::{BackpressurePolicy, BoundDim, Overflow, RelayCell};
#[derive(Debug, Clone)]
pub struct RatePolicy {
capacity: u64,
tokens: u64,
refill_per_tick: u64,
}
impl RatePolicy {
pub fn new(capacity: u64, refill_per_tick: u64) -> Self {
Self {
capacity,
tokens: capacity,
refill_per_tick,
}
}
pub fn tokens(&self) -> u64 {
self.tokens
}
pub fn try_egress(&mut self) -> bool {
if self.tokens > 0 {
self.tokens -= 1;
true
} else {
false
}
}
pub fn tick(&mut self) {
self.tokens = (self.tokens + self.refill_per_tick).min(self.capacity);
}
}
#[derive(Debug, Clone)]
pub struct WindowPolicy {
window_ops: u64,
pending: u64,
}
impl WindowPolicy {
pub fn new(window_ops: u64) -> Self {
Self {
window_ops: window_ops.max(1),
pending: 0,
}
}
pub fn on_ingress(&mut self) -> bool {
self.pending += 1;
if self.pending >= self.window_ops {
self.pending = 0;
true
} else {
false
}
}
pub fn tick(&mut self) -> bool {
if self.pending > 0 {
self.pending = 0;
true
} else {
false
}
}
}
#[derive(Debug, Clone)]
pub struct ExpiryPolicy {
ttl: u64,
now: u64,
}
impl ExpiryPolicy {
pub fn new(ttl: u64) -> Self {
Self { ttl, now: 0 }
}
pub fn advance(&mut self, by: u64) {
self.now += by;
}
pub fn now(&self) -> u64 {
self.now
}
pub fn is_live(&self, stamped_at: u64) -> bool {
self.now.saturating_sub(stamped_at) <= self.ttl
}
pub fn retain_live<T>(&self, batch: Vec<(u64, T)>) -> Vec<T> {
batch
.into_iter()
.filter(|(ts, _)| self.is_live(*ts))
.map(|(_, v)| v)
.collect()
}
}
#[derive(Debug, Default)]
pub struct PriorityStorage<T> {
heap: BinaryHeap<(u64, Reverse<u64>, usize)>,
items: Vec<Option<T>>,
seq: u64,
}
impl<T> PriorityStorage<T> {
pub fn new() -> Self {
Self {
heap: BinaryHeap::new(),
items: Vec::new(),
seq: 0,
}
}
pub fn push(&mut self, priority: u64, value: T) {
let idx = self.items.len();
self.items.push(Some(value));
self.heap.push((priority, Reverse(self.seq), idx));
self.seq += 1;
}
pub fn pop(&mut self) -> Option<T> {
while let Some((_, _, idx)) = self.heap.pop() {
if let Some(v) = self.items[idx].take() {
return Some(v);
}
}
None
}
pub fn len(&self) -> usize {
self.heap.len()
}
pub fn is_empty(&self) -> bool {
self.heap.is_empty()
}
}
pub struct KeyedRelay<K, T, M> {
shards: HashMap<K, RelayCell<T, M>>,
high_water: u64,
overflow: Overflow,
}
impl<K, T, M> KeyedRelay<K, T, M>
where
K: Eq + Hash + Clone,
T: Clone + PartialEq + 'static,
M: MergePolicy<T>,
{
pub fn new(high_water: u64, overflow: Overflow) -> Self {
Self {
shards: HashMap::new(),
high_water,
overflow,
}
}
pub fn ingress(&mut self, ctx: &Context, key: K, op: T) {
let hw = self.high_water;
let ov = self.overflow;
let relay = self.shards.entry(key).or_insert_with(|| {
RelayCell::new(
ctx,
BackpressurePolicy::new(ctx, BoundDim::Count, hw, hw / 2, ov),
)
.expect("keyed relay shard config")
});
relay.ingress(ctx, op);
}
pub fn drain(&self, ctx: &Context, key: &K) -> Option<T> {
self.shards.get(key).and_then(|r| r.drain(ctx))
}
pub fn keys(&self) -> impl Iterator<Item = &K> {
self.shards.keys()
}
}