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

optimizing memory usage #190

Merged
merged 4 commits into from
Sep 14, 2023
Merged
Show file tree
Hide file tree
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
19 changes: 16 additions & 3 deletions teos-common/src/test_utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,9 @@ use rand::distributions::Standard;
use rand::prelude::Distribution;
use rand::Rng;

use bitcoin::consensus;
use bitcoin::hashes::Hash;
use bitcoin::secp256k1::SecretKey;
use bitcoin::Txid;
use bitcoin::{consensus, Script, Transaction, TxOut, Txid};

use crate::appointment::{Appointment, Locator};
use crate::cryptography;
Expand All @@ -32,6 +31,12 @@ pub fn get_random_user_id() -> UserId {
UserId(pk)
}

pub fn get_random_locator() -> Locator {
let mut rng = rand::thread_rng();

Locator::from_slice(&rng.gen::<[u8; 16]>()).unwrap()
}

pub fn generate_random_appointment(dispute_txid: Option<&Txid>) -> Appointment {
let dispute_txid = match dispute_txid {
Some(l) => *l,
Expand All @@ -42,7 +47,15 @@ pub fn generate_random_appointment(dispute_txid: Option<&Txid>) -> Appointment {
};

let tx_bytes = Vec::from_hex(TX_HEX).unwrap();
let penalty_tx = consensus::deserialize(&tx_bytes).unwrap();
let mut penalty_tx: Transaction = consensus::deserialize(&tx_bytes).unwrap();

// Append a random-sized OP_RETURN to make each transcation random in size.
sr-gi marked this conversation as resolved.
Show resolved Hide resolved
penalty_tx.output.push(TxOut {
value: 0,
script_pubkey: Script::new_op_return(&cryptography::get_random_bytes(
get_random_int::<usize>() % 81,
)),
});

let mut raw_locator: [u8; 16] = cryptography::get_random_bytes(16).try_into().unwrap();
raw_locator.copy_from_slice(&dispute_txid[..16]);
Expand Down
21 changes: 15 additions & 6 deletions teos/src/api/http.rs
Original file line number Diff line number Diff line change
Expand Up @@ -660,8 +660,11 @@ mod tests_methods {
};
use super::*;

use crate::extended_appointment::UUID;
use crate::test_utils::{generate_dummy_appointment, ApiConfig, DURATION, SLOTS};
use crate::responder::{ConfirmationStatus, TransactionTracker};
use crate::test_utils::{
generate_dummy_appointment, get_random_tx, ApiConfig, DURATION, SLOTS,
};
use crate::watcher::Breach;

use teos_common::test_utils::get_random_user_id;
use teos_common::{cryptography, UserId};
Expand Down Expand Up @@ -830,14 +833,20 @@ mod tests_methods {
.await
.unwrap();

// Add the appointment to the Responder so it counts as triggered
let appointment = generate_dummy_appointment(None).inner;
let signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();
// Add the appointment to the Responder as a tracker so it counts as triggered
let dispute_tx = get_random_tx();
let tracker = TransactionTracker::new(
Breach::new(dispute_tx.clone(), get_random_tx()),
UserId(user_pk),
ConfirmationStatus::ConfirmedIn(100),
);
internal_api
.get_watcher()
.add_random_tracker_to_responder(UUID::new(appointment.locator, UserId(user_pk)));
.add_dummy_tracker_to_responder(&tracker);

// Try to add it via the http API
let appointment = generate_dummy_appointment(Some(&dispute_tx.txid())).inner;
let signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();
assert_eq!(
check_api_error(
Endpoint::AddAppointment,
Expand Down
77 changes: 43 additions & 34 deletions teos/src/api/internal.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use std::sync::{Arc, Condvar, Mutex};
use tonic::{Code, Request, Response, Status};
use triggered::Trigger;

use crate::extended_appointment::UUID;
use crate::protos as msgs;
use crate::protos::private_tower_services_server::PrivateTowerServices;
use crate::protos::public_tower_services_server::PublicTowerServices;
Expand Down Expand Up @@ -386,10 +387,14 @@ impl PrivateTowerServices for Arc<InternalAPI> {
})?;

match self.watcher.get_user_info(user_id) {
Some(info) => Ok(Response::new(msgs::GetUserResponse {
Some((info, locators)) => Ok(Response::new(msgs::GetUserResponse {
available_slots: info.available_slots,
subscription_expiry: info.subscription_expiry,
appointments: info.appointments.keys().map(|uuid| uuid.to_vec()).collect(),
// TODO: Should make it return locators and make `get_appointments` queryable using the (user_id, locator) pair for consistency.
sr-gi marked this conversation as resolved.
Show resolved Hide resolved
appointments: locators
.into_iter()
.map(|locator| UUID::new(locator, user_id).to_vec())
.collect(),
})),
None => Err(Status::new(Code::NotFound, "User not found")),
}
Expand Down Expand Up @@ -429,11 +434,10 @@ mod tests_private_api {
use bitcoin::hashes::Hash;
use bitcoin::Txid;

use crate::extended_appointment::UUID;
use crate::responder::{ConfirmationStatus, TransactionTracker};
use crate::test_utils::{
create_api, generate_dummy_appointment, generate_uuid, get_random_tx, DURATION, SLOTS,
START_HEIGHT,
create_api, generate_dummy_appointment, generate_dummy_appointment_with_user,
get_random_tx, DURATION, SLOTS, START_HEIGHT,
};
use crate::watcher::Breach;

Expand Down Expand Up @@ -486,9 +490,7 @@ mod tests_private_api {
let (internal_api, _s) = create_api().await;

// Add data to the Responser so we can retrieve it later on
internal_api
.watcher
.add_random_tracker_to_responder(generate_uuid());
internal_api.watcher.add_random_tracker_to_responder();

let response = internal_api
.get_all_appointments(Request::new(()))
Expand Down Expand Up @@ -588,7 +590,7 @@ mod tests_private_api {
);
internal_api
.watcher
.add_dummy_tracker_to_responder(generate_uuid(), &tracker);
.add_dummy_tracker_to_responder(&tracker);
}

let locator = Locator::new(dispute_tx.txid());
Expand Down Expand Up @@ -655,9 +657,7 @@ mod tests_private_api {

// And the Responder
for _ in 0..3 {
internal_api
.watcher
.add_random_tracker_to_responder(generate_uuid());
internal_api.watcher.add_random_tracker_to_responder();
}

let response = internal_api
Expand Down Expand Up @@ -730,12 +730,11 @@ mod tests_private_api {
assert!(response.appointments.is_empty());

// Add an appointment and check back
let appointment = generate_dummy_appointment(None).inner;
let uuid = UUID::new(appointment.locator, user_id);
let user_signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();
let (uuid, appointment) = generate_dummy_appointment_with_user(user_id, None);
let user_signature = cryptography::sign(&appointment.inner.to_vec(), &user_sk).unwrap();
internal_api
.watcher
.add_appointment(appointment.clone(), user_signature)
.add_appointment(appointment.inner, user_signature)
.unwrap();

let response = internal_api
Expand Down Expand Up @@ -786,10 +785,12 @@ mod tests_private_api {
mod tests_public_api {
use super::*;

use crate::extended_appointment::UUID;
use crate::responder::{ConfirmationStatus, TransactionTracker};
use crate::test_utils::{
create_api, create_api_with_config, generate_dummy_appointment, ApiConfig, DURATION, SLOTS,
create_api, create_api_with_config, generate_dummy_appointment, get_random_tx, ApiConfig,
DURATION, SLOTS,
};
use crate::watcher::Breach;
use teos_common::cryptography::{self, get_random_keypair};

#[tokio::test]
Expand Down Expand Up @@ -900,12 +901,12 @@ mod tests_public_api {
internal_api.watcher.register(UserId(user_pk)).unwrap();

let appointment = generate_dummy_appointment(None).inner;
let user_signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();
let signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();

let response = internal_api
.add_appointment(Request::new(common_msgs::AddAppointmentRequest {
appointment: Some(appointment.clone().into()),
signature: user_signature.clone(),
signature,
}))
.await
.unwrap()
Expand All @@ -925,12 +926,12 @@ mod tests_public_api {
let (user_sk, _) = get_random_keypair();

let appointment = generate_dummy_appointment(None).inner;
let user_signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();
let signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();

match internal_api
.add_appointment(Request::new(common_msgs::AddAppointmentRequest {
appointment: Some(appointment.clone().into()),
signature: user_signature.clone(),
signature,
}))
.await
{
Expand All @@ -954,12 +955,12 @@ mod tests_public_api {
internal_api.watcher.register(UserId(user_pk)).unwrap();

let appointment = generate_dummy_appointment(None).inner;
let user_signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();
let signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();

match internal_api
.add_appointment(Request::new(common_msgs::AddAppointmentRequest {
appointment: Some(appointment.clone().into()),
signature: user_signature.clone(),
signature,
}))
.await
{
Expand All @@ -983,12 +984,12 @@ mod tests_public_api {
internal_api.watcher.register(UserId(user_pk)).unwrap();

let appointment = generate_dummy_appointment(None).inner;
let user_signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();
let signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();

match internal_api
.add_appointment(Request::new(common_msgs::AddAppointmentRequest {
appointment: Some(appointment.clone().into()),
signature: user_signature.clone(),
signature,
}))
.await
{
Expand All @@ -1008,16 +1009,24 @@ mod tests_public_api {
let user_id = UserId(user_pk);
internal_api.watcher.register(user_id).unwrap();

let appointment = generate_dummy_appointment(None).inner;
let user_signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();
// Add a tracker to the responder to simulate it being triggered.
let dispute_tx = get_random_tx();
let tracker = TransactionTracker::new(
Breach::new(dispute_tx.clone(), get_random_tx()),
user_id,
ConfirmationStatus::ConfirmedIn(100),
);
internal_api
.watcher
.add_random_tracker_to_responder(UUID::new(appointment.locator, user_id));
.get_watcher()
.add_dummy_tracker_to_responder(&tracker);

// Try to add it again using the API.
let appointment = generate_dummy_appointment(Some(&dispute_tx.txid())).inner;
let signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();
match internal_api
.add_appointment(Request::new(common_msgs::AddAppointmentRequest {
appointment: Some(appointment.clone().into()),
signature: user_signature.clone(),
appointment: Some(appointment.into()),
signature,
}))
.await
{
Expand All @@ -1038,12 +1047,12 @@ mod tests_public_api {

let (user_sk, _) = get_random_keypair();
let appointment = generate_dummy_appointment(None).inner;
let user_signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();
let signature = cryptography::sign(&appointment.to_vec(), &user_sk).unwrap();

match internal_api
.add_appointment(Request::new(common_msgs::AddAppointmentRequest {
appointment: Some(appointment.clone().into()),
signature: user_signature.clone(),
signature,
}))
.await
{
Expand Down
Loading
Loading