use {
crate::{
bridge::{KeyedTransactionMeta, ScheduleBatch, SchedulerBindingsBridge, TransactionKey},
handshake::{AgaveSession, ClientLogon, client, server::Server},
responses_region::{execution_responses_from_iter, resolve_responses_from_iter},
transaction_ptr::TransactionPtrBatch,
},
agave_scheduler_bindings::{
ProgressMessage, SharablePubkeys, SharableTransactionBatchRegion,
SharableTransactionRegion, TpuToPackMessage, TransactionResponseRegion,
WorkerToPackMessage, pack_message_flags, processed_codes,
worker_message_types::{
CheckResponse, ExecutionResponse, fee_payer_balance_flags, not_included_reasons,
resolve_flags, status_check_flags,
},
},
solana_pubkey::Pubkey,
solana_transaction::versioned::VersionedTransaction,
std::ops::{Deref, DerefMut},
};
pub struct TestBridge<M>
where
M: Copy,
{
bridge: SchedulerBindingsBridge<M>,
agave: AgaveSession,
}
impl<M> Deref for TestBridge<M>
where
M: Copy,
{
type Target = SchedulerBindingsBridge<M>;
fn deref(&self) -> &Self::Target {
&self.bridge
}
}
impl<M> DerefMut for TestBridge<M>
where
M: Copy,
{
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.bridge
}
}
impl<M> TestBridge<M>
where
M: Copy,
{
#[must_use]
pub fn new(worker_count: usize, worker_req_cap: usize) -> Self {
assert!(
worker_req_cap.is_power_of_two(),
"shaq requires power of 2 queue sizes"
);
let logon = ClientLogon {
worker_count,
allocator_size: 64 * 1024 * 1024,
allocator_handles: 1,
tpu_to_pack_capacity: 1024,
progress_tracker_capacity: 256,
pack_to_worker_capacity: worker_req_cap,
worker_to_pack_capacity: 1024,
flags: 0,
};
let (agave, files) = Server::setup_session(logon).unwrap();
let client_session = client::setup_session(&logon, files).unwrap();
Self {
bridge: SchedulerBindingsBridge::new(client_session),
agave,
}
}
#[must_use]
pub fn tx_count(&self) -> usize {
self.bridge.state().len()
}
#[must_use]
pub fn contains_tx(&self, key: TransactionKey) -> bool {
self.bridge.state().contains_key(key)
}
pub fn queue_progress(&mut self, progress: ProgressMessage) {
self.agave.progress_tracker.try_write(progress).unwrap();
self.agave.progress_tracker.commit();
}
pub fn queue_tpu(&mut self, tx: &VersionedTransaction) {
let serialized = wincode::serialize(tx).unwrap();
let allocator = &self.agave.tpu_to_pack.allocator;
let ptr = allocator
.allocate(serialized.len().try_into().unwrap())
.unwrap();
unsafe {
std::ptr::copy_nonoverlapping(serialized.as_ptr(), ptr.as_ptr(), serialized.len());
}
let offset = unsafe { allocator.offset(ptr) };
let msg = TpuToPackMessage {
transaction: SharableTransactionRegion {
offset,
length: serialized.len() as u32,
},
flags: 0,
src_addr: [0; 16],
};
self.agave.tpu_to_pack.producer.try_write(msg).unwrap();
self.agave.tpu_to_pack.producer.commit();
}
pub fn queue_check_response_ok(
&mut self,
batch: &ScheduleBatch<Vec<KeyedTransactionMeta<M>>>,
index: usize,
keys: Option<Vec<Pubkey>>,
) {
self.queue_check_response(batch, index, keys, self.check_ok());
}
pub fn queue_check_response(
&mut self,
batch: &ScheduleBatch<Vec<KeyedTransactionMeta<M>>>,
index: usize,
keys: Option<Vec<Pubkey>>,
mut response: CheckResponse,
) {
let worker_idx = batch.worker;
if let Some(keys) = keys {
let worker = &mut self.agave.workers[worker_idx];
let pubkeys_ptr = worker
.allocator
.allocate(
(keys
.len()
.checked_mul(std::mem::size_of::<Pubkey>())
.unwrap())
.try_into()
.unwrap(),
)
.unwrap();
unsafe {
std::ptr::copy_nonoverlapping(
keys.as_ptr().cast::<u8>(),
pubkeys_ptr.as_ptr(),
keys.len()
.checked_mul(std::mem::size_of::<Pubkey>())
.unwrap(),
);
}
let offset = unsafe { worker.allocator.offset(pubkeys_ptr) };
response.resolved_pubkeys = SharablePubkeys {
offset,
num_pubkeys: keys.len() as u32,
};
}
let batch_region = self.build_single_tx_batch_region(batch, index, worker_idx);
let worker = &mut self.agave.workers[worker_idx];
let responses_region =
resolve_responses_from_iter(&worker.allocator, [response].into_iter()).unwrap();
let msg = WorkerToPackMessage {
batch: batch_region,
processed_code: processed_codes::PROCESSED,
responses: responses_region,
};
worker.worker_to_pack.try_write(msg).unwrap();
worker.worker_to_pack.commit();
}
pub fn queue_all_checks_ok(&mut self) {
while let Some(batch) = self.pop_schedule() {
assert_eq!(batch.flags & 1, pack_message_flags::CHECK);
for i in 0..batch.transactions.len() {
self.queue_check_response_ok(&batch, i, None);
}
}
}
pub fn queue_execute_response(
&mut self,
batch: &ScheduleBatch<Vec<KeyedTransactionMeta<M>>>,
index: usize,
response: ExecutionResponse,
) {
let worker_idx = batch.worker;
let batch_region = self.build_single_tx_batch_region(batch, index, worker_idx);
let worker = &mut self.agave.workers[worker_idx];
let responses_region =
execution_responses_from_iter(&worker.allocator, [response].into_iter()).unwrap();
let msg = WorkerToPackMessage {
batch: batch_region,
processed_code: processed_codes::PROCESSED,
responses: responses_region,
};
worker.worker_to_pack.try_write(msg).unwrap();
worker.worker_to_pack.commit();
}
pub fn queue_unprocessed_response(
&mut self,
batch: &ScheduleBatch<Vec<KeyedTransactionMeta<M>>>,
index: usize,
) {
let worker_idx = batch.worker;
let batch_region = self.build_single_tx_batch_region(batch, index, worker_idx);
let worker = &mut self.agave.workers[worker_idx];
let msg = WorkerToPackMessage {
batch: batch_region,
processed_code: processed_codes::MAX_WORKING_SLOT_EXCEEDED,
responses: TransactionResponseRegion {
tag: 0,
num_transaction_responses: 0,
transaction_responses_offset: 0,
},
};
worker.worker_to_pack.try_write(msg).unwrap();
worker.worker_to_pack.commit();
}
pub fn pop_schedule(&mut self) -> Option<ScheduleBatch<Vec<KeyedTransactionMeta<M>>>> {
for (worker_idx, worker) in self.agave.workers.iter_mut().enumerate() {
worker.pack_to_worker.sync();
if let Some(msg) = worker.pack_to_worker.try_read() {
let msg = *msg;
worker.pack_to_worker.finalize();
let batch = unsafe {
TransactionPtrBatch::<KeyedTransactionMeta<M>>::from_sharable_transaction_batch_region(
&msg.batch,
self.bridge.allocator(),
)
};
let transactions: Vec<_> = batch.iter().map(|(_tx_ptr, meta)| meta).collect();
unsafe { batch.free() };
return Some(ScheduleBatch {
worker: worker_idx,
transactions,
max_working_slot: msg.max_working_slot,
flags: msg.flags,
});
}
}
None
}
pub fn check_ok(&self) -> CheckResponse {
let progress = self.bridge.progress();
CheckResponse {
parsing_and_sanitization_flags: 0,
status_check_flags: status_check_flags::REQUESTED | status_check_flags::PERFORMED,
fee_payer_balance_flags: fee_payer_balance_flags::REQUESTED
| fee_payer_balance_flags::PERFORMED,
resolve_flags: resolve_flags::REQUESTED | resolve_flags::PERFORMED,
included_slot: progress.current_slot,
balance_slot: progress.current_slot,
fee_payer_balance: u64::from(u32::MAX),
resolution_slot: progress.current_slot,
min_alt_deactivation_slot: u64::MAX,
resolved_pubkeys: SharablePubkeys {
offset: 0,
num_pubkeys: 0,
},
}
}
#[must_use]
pub fn execute_ok(&self) -> ExecutionResponse {
ExecutionResponse {
execution_slot: self.bridge.progress().current_slot,
not_included_reason: not_included_reasons::NONE,
cost_units: 0,
fee_payer_balance: u64::from(u32::MAX),
}
}
#[must_use]
pub fn execute_err(&self, reason: u8) -> ExecutionResponse {
ExecutionResponse {
execution_slot: self.bridge.progress().current_slot,
not_included_reason: reason,
cost_units: 0,
fee_payer_balance: u64::from(u32::MAX),
}
}
fn build_single_tx_batch_region(
&self,
batch: &ScheduleBatch<Vec<KeyedTransactionMeta<M>>>,
index: usize,
worker_idx: usize,
) -> SharableTransactionBatchRegion {
type Batch<'a, M> = TransactionPtrBatch<'a, KeyedTransactionMeta<M>>;
let meta = batch.transactions[index];
let worker_allocator = &self.agave.workers[worker_idx].allocator;
let batch_ptr = worker_allocator
.allocate(Batch::<M>::TRANSACTION_META_END as u32)
.unwrap();
let batch_offset = unsafe { worker_allocator.offset(batch_ptr) };
let tx_state = self.bridge.transaction(meta.key);
let tx_region = unsafe {
tx_state
.data
.inner_data()
.to_sharable_transaction_region(self.bridge.allocator())
};
let tx_ptr = batch_ptr.cast::<SharableTransactionRegion>();
unsafe { tx_ptr.as_ptr().write(tx_region) };
let meta_ptr = unsafe {
batch_ptr
.as_ptr()
.byte_add(Batch::<M>::TRANSACTION_META_START)
.cast::<KeyedTransactionMeta<M>>()
};
unsafe { meta_ptr.write(meta) };
SharableTransactionBatchRegion {
num_transactions: 1,
transactions_offset: batch_offset,
}
}
}