Skip to content
This repository has been archived by the owner on Dec 4, 2024. It is now read-only.

Commit

Permalink
fix: remove allocation management from pool
Browse files Browse the repository at this point in the history
  • Loading branch information
@Theodus
Theodus committed May 8, 2024
1 parent b12e197 commit e94e0f1
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 157 deletions.
154 changes: 16 additions & 138 deletions src/pool.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,24 +13,13 @@ const SIGNATURE_RANGE: Range = next_range::<Signature>(RECEIPT_ID_RANGE);
const UNLOCKED_FEE_RANGE: Range = next_range::<U256>(SIGNATURE_RANGE);
pub const BORROWED_RECEIPT_LEN: usize = UNLOCKED_FEE_RANGE.end;

/// A collection of installed allocation that can borrow or generate receipts.
#[derive(Default, Debug, PartialEq, Eq)]
pub struct ReceiptPool {
allocations: Vec<Allocation>,
}

/// A in-flight state for an allocation on-chain.
// This must never implement Clone
/// A per-allocation collection that can borrow or generate receipts.
#[derive(Debug, PartialEq, Eq)]
struct Allocation {
pub struct ReceiptPool {
pub allocation: Address,
/// Receipts that can be folded. These contain an unbroken chain
/// of agreed upon history between the Indexer and Gateway.
receipt_cache: Vec<PooledReceipt>,
/// Signer: Signs the receipts. Each allocation must have a globally unique ppk pair.
/// If keys are shared across multiple allocations it will allow the Indexer to
/// double-collect
signer: SecretKey,
allocation_id: Address,
}

#[derive(Copy, Clone, PartialEq, Eq, Debug)]
Expand Down Expand Up @@ -72,9 +61,10 @@ impl From<SignError> for BorrowFail {
}

impl ReceiptPool {
pub fn new() -> Self {
pub fn new(allocation: Address) -> Self {
Self {
allocations: Vec::new(),
allocation,
receipt_cache: Default::default(),
}
}

Expand All @@ -84,80 +74,22 @@ impl ReceiptPool {
#[cfg(test)]
pub fn known_unlocked_fees(&self) -> U256 {
let mut result = U256::zero();
for fee in self
.allocations
.iter()
.flat_map(|a| &a.receipt_cache)
.map(|r| r.unlocked_fee)
{
for fee in self.receipt_cache.iter().map(|r| r.unlocked_fee) {
result += fee;
}
result
}

pub fn add_allocation(&mut self, signer: SecretKey, allocation_id: Address) {
// Defensively ensure we don't already have this allocation.
for allocation in self.allocations.iter() {
if allocation.allocation_id == allocation_id {
return;
}
}

let allocation = Allocation {
signer,
receipt_cache: Vec::new(),
allocation_id,
};
self.allocations.push(allocation)
}

pub fn remove_allocation(&mut self, allocation_id: &Address) {
if let Some(index) = self
.allocations
.iter()
.position(|a| &a.allocation_id == allocation_id)
{
self.allocations.swap_remove(index);
}
}

pub fn has_collateral_for(&self) -> bool {
self.allocations.len() != 0
}

fn select_allocation(&mut self) -> Result<&mut Allocation, BorrowFail> {
// Prefer the one most recently added
self.allocations.last_mut().ok_or(BorrowFail::NoAllocation)
}

pub fn contains_allocation(&self, allocation_id: &Address) -> bool {
self.allocations
.iter()
.any(|a| &a.allocation_id == allocation_id)
}

pub fn addresses(&self) -> Vec<Address> {
self.allocations.iter().map(|a| a.allocation_id).collect()
}

fn allocation_by_id_mut(&mut self, allocation_id: &Address) -> Option<&mut Allocation> {
self.allocations
.iter_mut()
.find(|a| &a.allocation_id == allocation_id)
}

pub fn commit(&mut self, locked_fee: U256) -> Result<Vec<u8>, BorrowFail> {
let allocation = self.select_allocation()?;

let receipt = if allocation.receipt_cache.len() == 0 {
pub fn commit(&mut self, signer: &SecretKey, locked_fee: U256) -> Result<Vec<u8>, BorrowFail> {
let receipt = if self.receipt_cache.len() == 0 {
let mut receipt_id = ReceiptId::default();
rng().fill_bytes(&mut receipt_id);
PooledReceipt {
receipt_id,
unlocked_fee: U256::zero(),
}
} else {
let receipts = &mut allocation.receipt_cache;
let receipts = &mut self.receipt_cache;
let index = rng().gen_range(0..receipts.len());
receipts.swap_remove(index)
};
Expand All @@ -170,7 +102,7 @@ impl ReceiptPool {
// This is: [allocation_id, fee, receipt_id, signature]
let mut commitment = Vec::with_capacity(BORROWED_RECEIPT_LEN);
let fee = receipt.unlocked_fee + locked_fee;
commitment.extend_from_slice(&allocation.allocation_id);
commitment.extend_from_slice(&self.allocation);
commitment.extend_from_slice(&to_be_bytes(fee));
commitment.extend_from_slice(&receipt.receipt_id);

Expand All @@ -183,7 +115,7 @@ impl ReceiptPool {
// The part of the message that needs to be signed in the fee and receipt id only.
let signature = sign(
&commitment[ALLOCATION_ID_RANGE.start..RECEIPT_ID_RANGE.end],
&allocation.signer,
signer,
)?;
commitment.extend_from_slice(&signature);

Expand All @@ -198,15 +130,6 @@ impl ReceiptPool {

pub fn release(&mut self, bytes: &[u8], status: QueryStatus) {
assert_eq!(bytes.len(), BORROWED_RECEIPT_LEN);
let allocation_id: Address = bytes[ALLOCATION_ID_RANGE].try_into().unwrap();

// Try to find the allocation. If there is no allocation, it means it's been uninstalled.
// In that case, drop the receipt.
let allocation = if let Some(allocation) = self.allocation_by_id_mut(&allocation_id) {
allocation
} else {
return;
};

let unlocked_fee = if status == QueryStatus::Success {
U256::from_big_endian(&bytes[FEE_RANGE])
Expand All @@ -218,7 +141,7 @@ impl ReceiptPool {
unlocked_fee,
receipt_id: bytes[RECEIPT_ID_RANGE].try_into().unwrap(),
};
allocation.receipt_cache.push(receipt);
self.receipt_cache.push(receipt);
}
}

Expand All @@ -227,23 +150,16 @@ mod tests {
use super::*;
use crate::tests::*;

#[track_caller]
fn assert_failed_borrow(pool: &mut ReceiptPool, fee: impl Into<U256>) {
let receipt = pool.commit(fee.into());
assert_eq!(receipt, Err(BorrowFail::NoAllocation));
}

#[track_caller]
fn assert_successful_borrow(pool: &mut ReceiptPool, fee: impl Into<U256>) -> Vec<u8> {
pool.commit(U256::from(fee.into()))
pool.commit(&test_signer(), U256::from(fee.into()))
.expect("Should be able to borrow")
}

// Simple happy-path case of paying for requests in a loop.
#[test]
pub fn can_pay_for_requests() {
let mut pool = ReceiptPool::new();
pool.add_allocation(test_signer(), bytes(1));
let mut pool = ReceiptPool::new(bytes(1));

for i in 1..=10 {
let borrow = assert_successful_borrow(&mut pool, i);
Expand All @@ -254,48 +170,10 @@ mod tests {
}
}

#[test]
pub fn selects_allocation() {
let mut pool = ReceiptPool::new();

pool.add_allocation(test_signer(), bytes(1));
pool.add_allocation(test_signer(), bytes(2));
pool.add_allocation(test_signer(), bytes(3));
pool.add_allocation(test_signer(), bytes(4));
pool.add_allocation(test_signer(), bytes(5));
pool.add_allocation(test_signer(), bytes(6));
pool.add_allocation(test_signer(), bytes(7));
pool.add_allocation(test_signer(), bytes(8));

assert_successful_borrow(&mut pool, 2);
assert_successful_borrow(&mut pool, 4);
assert_successful_borrow(&mut pool, 3);
assert_successful_borrow(&mut pool, 1);
assert_successful_borrow(&mut pool, 2);
assert_successful_borrow(&mut pool, 3);
assert_successful_borrow(&mut pool, 1);
assert_successful_borrow(&mut pool, 4);
}

// Any uninstalled allocation cannot be used to pay for queries.
#[test]
fn removed_allocation_cannot_pay() {
let mut pool = ReceiptPool::new();
pool.add_allocation(test_signer(), bytes(2));
pool.add_allocation(test_signer(), bytes(1));

pool.remove_allocation(&bytes(2));
assert_successful_borrow(&mut pool, 5);
pool.remove_allocation(&bytes(1));
assert_failed_borrow(&mut pool, 5);
}

// Tests modes for returning collateral
#[test]
fn collateral_return() {
let mut pool = ReceiptPool::new();

pool.add_allocation(test_signer(), bytes(2));
let mut pool = ReceiptPool::new(bytes(2));

let borrow3 = assert_successful_borrow(&mut pool, 3);
assert_eq!(pool.known_unlocked_fees(), 0.into());
Expand Down
30 changes: 11 additions & 19 deletions src/tests.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,17 +17,14 @@ fn debug_hex(bytes: &[u8]) {
// This is just useful for constructing a value to test with.
#[test]
pub fn make_receipt() {
let mut pool = ReceiptPool::new();
let signer = test_signer();

pool.add_allocation(signer, bytes(100));
let mut pool = ReceiptPool::new(bytes(100));

println!("Receipt 0: value 5");
let commit0 = pool.commit(U256::from(5)).unwrap();
let commit0 = pool.commit(&test_signer(), U256::from(5)).unwrap();
debug_hex(&commit0);

println!("Receipt 1: value 8");
let commit1 = pool.commit(U256::from(8)).unwrap();
let commit1 = pool.commit(&test_signer(), U256::from(8)).unwrap();
debug_hex(&commit1);
}

Expand All @@ -46,17 +43,15 @@ pub fn test_signer() -> SecretKey {
#[test]
#[ignore = "Benchmark"]
fn speed() {
let mut pool = ReceiptPool::new();
pool.add_allocation(test_signer(), bytes(0));
pool.add_allocation(test_signer(), bytes(1));
let mut pool = ReceiptPool::new(bytes(0));

let mut borrows = Vec::<Vec<u8>>::new();

let start = Instant::now();

for _ in 0..2700 {
for _ in 0..10 {
let commitment = pool.commit(U256::from(100)).unwrap();
let commitment = pool.commit(&test_signer(), U256::from(100)).unwrap();
borrows.push(commitment)
}
while let Some(borrow) = borrows.pop() {
Expand All @@ -74,12 +69,11 @@ fn attempt_to_double_collect_with_partial_voucher_rejects() {
let allocation_id = bytes(1);

// Create a bunch of receipts
let mut pool = ReceiptPool::new();
pool.add_allocation(test_signer(), allocation_id);
let mut pool = ReceiptPool::new(allocation_id);
let mut borrows = Vec::<Vec<u8>>::new();
for _ in 0..10 {
let fee = U256::from(1);
let commitment = pool.commit(fee).unwrap();
let commitment = pool.commit(&test_signer(), fee).unwrap();
borrows.push(commitment);
}

Expand Down Expand Up @@ -112,8 +106,7 @@ fn vouchers() {
let allocation_id = bytes(1);

// Create a bunch of receipts
let mut pool = ReceiptPool::new();
pool.add_allocation(test_signer(), allocation_id);
let mut pool = ReceiptPool::new(allocation_id);
let mut borrows = Vec::<Vec<u8>>::new();
let mut fees = U256::zero();
for i in 2..10 {
Expand All @@ -123,7 +116,7 @@ fn vouchers() {
for _ in 0..i {
let fee = U256::from(1);
fees += fee;
let commitment = pool.commit(fee).unwrap();
let commitment = pool.commit(&test_signer(), fee).unwrap();
borrows.push(commitment);
}
}
Expand Down Expand Up @@ -232,11 +225,10 @@ fn partial_vouchers_combine() {
}

fn create_receipts(allocation_id: Address, count: usize) -> Vec<u8> {
let mut pool = ReceiptPool::new();
pool.add_allocation(test_signer(), allocation_id);
let mut pool = ReceiptPool::new(allocation_id);
let mut borrows = Vec::<Vec<u8>>::new();
for _ in 1..=count {
let commitment = pool.commit(U256::from(1)).unwrap();
let commitment = pool.commit(&test_signer(), U256::from(1)).unwrap();
borrows.push(commitment);
}
receipts_from_borrows(borrows)
Expand Down

0 comments on commit e94e0f1

Please sign in to comment.