diff --git a/atra-ob/src/api/service.rs b/atra-ob/src/api/service.rs index ade4b14..0a2437f 100644 --- a/atra-ob/src/api/service.rs +++ b/atra-ob/src/api/service.rs @@ -1,24 +1,88 @@ -use std::sync::{Arc, Mutex}; -use tonic::{transport::Server, Request, Response, Status}; +use crate::core::MatchingEngine; +use crate::core::{Order, OrderType, Side}; +use crate::proto; +use crate::proto::order_book_service_server::{ + OrderBookService as GrpcService, OrderBookServiceServer, +}; +use crate::proto::{ + CancelOrderRequest, GetOrderBookRequest, GetOrderStatusRequest, GetTradeHistoryRequest, OrderBatchRequest, + OrderBatchResponse, OrderRequest, OrderResponse, Trade as ProtoTrade, TradeHistoryResponse, +}; +use prost_types::Timestamp; use rust_decimal::Decimal; +use std::collections::{HashMap, HashSet}; use std::str::FromStr; -use prost_types::Timestamp; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use tokio::sync::{mpsc, oneshot, Mutex, RwLock}; +use tonic::{transport::Server, Request, Response, Status}; -use crate::core::MatchingEngine; // Updated import -use crate::proto; -use crate::proto::order_book_service_server::{OrderBookService as GrpcService, OrderBookServiceServer}; -use crate::proto::{OrderRequest, OrderResponse, OrderBatchRequest, OrderBatchResponse, GetOrderBookRequest, OrderBookResponse, GetOrderStatusRequest, CancelOrderRequest}; -use crate::core::{Order, OrderType, Side}; -use crate::proto::{GetTradeHistoryRequest, TradeHistoryResponse, Trade as ProtoTrade}; +#[derive(Clone, Copy)] +pub struct SequencerConfig { + pub strict_sequence_validation: bool, +} + +impl SequencerConfig { + pub fn from_env() -> Self { + let strict = std::env::var("ATRA_STRICT_SEQUENCE_VALIDATION") + .ok() + .map(|v| matches!(v.as_str(), "1" | "true" | "TRUE" | "yes" | "YES")) + .unwrap_or(false); + Self { + strict_sequence_validation: strict, + } + } +} + +struct InstrumentState { + next_sequence: AtomicU64, +} + +impl InstrumentState { + fn new() -> Self { + Self { + next_sequence: AtomicU64::new(1), + } + } +} + +enum WorkerCommand { + Place { + order: Order, + response: oneshot::Sender>, + }, + Cancel { + order_id: u64, + response: oneshot::Sender>, + }, + Snapshot { + depth: usize, + response: oneshot::Sender>, + }, + Status { + order_id: u64, + response: oneshot::Sender>, + }, + Trades { + limit: usize, + response: oneshot::Sender, Status>>, + }, +} pub struct OrderBookService { - engine: Arc>, + lanes: Arc>>>, + lane_states: Arc>>>, + lane_count: u32, + config: SequencerConfig, } impl OrderBookService { - pub fn new(engine: MatchingEngine) -> Self { + pub fn new(lane_count: u32, config: SequencerConfig) -> Self { Self { - engine: Arc::new(Mutex::new(engine)) + lanes: Arc::new(RwLock::new(HashMap::new())), + lane_states: Arc::new(RwLock::new(HashMap::new())), + lane_count: lane_count.max(1), + config, } } @@ -30,32 +94,201 @@ impl OrderBookService { .await?; Ok(()) } -} -#[tonic::async_trait] -impl GrpcService for OrderBookService { - async fn place_order( - &self, - request: Request, - ) -> Result, Status> { - let req = request.into_inner(); - let price = Decimal::from_str(&req.price) - .map_err(|_| Status::invalid_argument("Invalid price format"))?; + fn lane_for_instrument(&self, instrument_id: u32) -> u32 { + instrument_id % self.lane_count + } + + async fn lane_sender_for_instrument(&self, instrument_id: u32) -> mpsc::Sender { + let lane_id = self.lane_for_instrument(instrument_id); + if let Some(sender) = self.lanes.read().await.get(&lane_id).cloned() { + return sender; + } + + let mut lanes = self.lanes.write().await; + if let Some(sender) = lanes.get(&lane_id).cloned() { + return sender; + } + + let (tx, rx) = mpsc::channel(4096); + tokio::spawn(run_lane_worker(rx)); + lanes.insert(lane_id, tx.clone()); + tx + } + + async fn state_for_instrument(&self, instrument_id: u32) -> Arc { + if let Some(state) = self.lane_states.read().await.get(&instrument_id).cloned() { + return state; + } + let mut states = self.lane_states.write().await; + if let Some(state) = states.get(&instrument_id).cloned() { + return state; + } + let state = Arc::new(InstrumentState::new()); + states.insert(instrument_id, state.clone()); + state + } + + async fn build_engine_order(&self, req: OrderRequest) -> Result { + let price = + Decimal::from_str(&req.price).map_err(|_| Status::invalid_argument("Invalid price format"))?; let quantity = Decimal::from_str(&req.quantity) .map_err(|_| Status::invalid_argument("Invalid quantity format"))?; + let state = self.state_for_instrument(req.instrument_id).await; + let next_expected = state.next_sequence.load(Ordering::SeqCst); + let sequence = if req.sequence_number == 0 { + state.next_sequence.fetch_add(1, Ordering::SeqCst) + } else { + if self.config.strict_sequence_validation && req.sequence_number != next_expected { + return Err(Status::failed_precondition(format!( + "Out-of-order sequence for instrument {}: expected {}, got {}", + req.instrument_id, next_expected, req.sequence_number + ))); + } + if req.sequence_number >= next_expected { + state + .next_sequence + .store(req.sequence_number.saturating_add(1), Ordering::SeqCst); + } + req.sequence_number + }; - let order = Order::new( + let mut order = Order::new( req.id, + req.instrument_id, + sequence, price, quantity, if req.side == 0 { Side::Bid } else { Side::Ask }, - if req.order_type == 0 { OrderType::Limit } else { OrderType::Market }, + if req.order_type == 0 { + OrderType::Limit + } else { + OrderType::Market + }, ); + order.ingress_timestamp_ns = req.ingress_timestamp_ns; + order.idempotency_key = req.idempotency_key; + Ok(order) + } +} - let result = self.engine - .lock() - .map_err(|_| Status::internal("Lock error"))? - .place_order(order); +async fn run_lane_worker(mut rx: mpsc::Receiver) { + let engines: Arc>> = Arc::new(Mutex::new(HashMap::new())); + let mut seen_idempotency: HashSet = HashSet::new(); + while let Some(cmd) = rx.recv().await { + match cmd { + WorkerCommand::Place { order, response } => { + if let Some(key) = &order.idempotency_key { + if seen_idempotency.contains(key) { + let _ = response.send(Err(Status::already_exists("Duplicate idempotency key"))); + continue; + } + seen_idempotency.insert(key.clone()); + } + let mut engines_locked = engines.lock().await; + let engine = engines_locked + .entry(order.instrument_id) + .or_insert_with(MatchingEngine::new); + let placed = engine.place_order(order); + let _ = response.send(Ok(placed)); + } + WorkerCommand::Cancel { order_id, response } => { + let mut engines_locked = engines.lock().await; + let mut cancelled = None; + for engine in engines_locked.values_mut() { + if let Some(order) = engine.cancel_order(order_id) { + cancelled = Some(order); + break; + } + } + let _ = response.send( + cancelled.ok_or_else(|| Status::not_found("Order not found or cannot be cancelled")), + ); + } + WorkerCommand::Snapshot { depth, response } => { + let mut bids = Vec::new(); + let mut asks = Vec::new(); + let engines_locked = engines.lock().await; + for engine in engines_locked.values() { + let (mut engine_bids, mut engine_asks) = engine.get_order_book(depth); + bids.append(&mut engine_bids); + asks.append(&mut engine_asks); + } + let _ = response.send(Ok(proto::OrderBookResponse { + bids: bids + .into_iter() + .map(|(price, qty)| proto::OrderBookLevel { + price: price.to_string(), + quantity: qty.to_string(), + }) + .collect(), + asks: asks + .into_iter() + .map(|(price, qty)| proto::OrderBookLevel { + price: price.to_string(), + quantity: qty.to_string(), + }) + .collect(), + })); + } + WorkerCommand::Status { order_id, response } => { + let engines_locked = engines.lock().await; + let mut found = None; + for engine in engines_locked.values() { + if let Some(order) = engine.get_order_status(order_id) { + found = Some(order.clone()); + break; + } + } + let _ = response.send(found.ok_or_else(|| Status::not_found("Order not found"))); + } + WorkerCommand::Trades { limit, response } => { + let engines_locked = engines.lock().await; + let mut trades = Vec::new(); + for engine in engines_locked.values() { + let mut history = engine + .get_trade_history(Some(limit)) + .into_iter() + .map(|trade| ProtoTrade { + maker_order_id: trade.maker_order_id, + taker_order_id: trade.taker_order_id, + price: trade.price.to_string(), + quantity: trade.quantity.to_string(), + side: trade.side as i32, + timestamp: trade.timestamp.map(|ts| Timestamp { + seconds: ts.timestamp(), + nanos: ts.timestamp_subsec_nanos() as i32, + }), + maker_sequence_number: trade.maker_sequence, + taker_sequence_number: trade.taker_sequence, + ingress_timestamp_ns: trade.ingress_timestamp_ns, + }) + .collect::>(); + trades.append(&mut history); + } + let _ = response.send(Ok(trades)); + } + } + } +} + +#[tonic::async_trait] +impl GrpcService for OrderBookService { + async fn place_order( + &self, + request: Request, + ) -> Result, Status> { + let req = request.into_inner(); + let order = self.build_engine_order(req).await?; + let lane_sender = self.lane_sender_for_instrument(order.instrument_id).await; + let (tx, rx) = oneshot::channel(); + lane_sender + .send(WorkerCommand::Place { order, response: tx }) + .await + .map_err(|_| Status::internal("Lane worker unavailable"))?; + let result = rx + .await + .map_err(|_| Status::internal("Lane worker response dropped"))??; Ok(Response::new(OrderResponse { id: result.id, @@ -65,10 +298,14 @@ impl GrpcService for OrderBookService { side: result.side as i32, order_type: result.order_type as i32, status: result.status as i32, - timestamp: Some(Timestamp { - seconds: result.timestamp.timestamp(), - nanos: result.timestamp.timestamp_subsec_nanos() as i32, + timestamp: result.timestamp.map(|ts| Timestamp { + seconds: ts.timestamp(), + nanos: ts.timestamp_subsec_nanos() as i32, }), + instrument_id: result.instrument_id, + sequence_number: result.sequence, + ingress_timestamp_ns: result.ingress_timestamp_ns, + idempotency_key: result.idempotency_key, })) } @@ -77,7 +314,15 @@ impl GrpcService for OrderBookService { &self, request: Request, ) -> Result, Status> { - Ok(Response::new(OrderBatchResponse{orders:Vec::new()})) + let mut responses = Vec::new(); + for order in request.into_inner().orders { + let placed = self + .place_order(Request::new(order)) + .await? + .into_inner(); + responses.push(placed); + } + Ok(Response::new(OrderBatchResponse { orders: responses })) } @@ -86,11 +331,21 @@ impl GrpcService for OrderBookService { request: Request, ) -> Result, Status> { let order_id = request.into_inner().order_id; - let cancelled_order = self.engine - .lock() - .map_err(|_| Status::internal("Lock error"))? - .cancel_order(order_id) - .ok_or_else(|| Status::not_found("Order not found or cannot be cancelled"))?; + let mut cancelled_order = None; + for lane in self.lanes.read().await.values().cloned() { + let (tx, rx) = oneshot::channel(); + lane.send(WorkerCommand::Cancel { + order_id, + response: tx, + }) + .await + .map_err(|_| Status::internal("Lane worker unavailable"))?; + if let Ok(Ok(order)) = rx.await { + cancelled_order = Some(order); + break; + } + } + let cancelled_order = cancelled_order.ok_or_else(|| Status::not_found("Order not found or cannot be cancelled"))?; Ok(Response::new(OrderResponse { id: cancelled_order.id, price: cancelled_order.price.to_string(), @@ -99,37 +354,36 @@ impl GrpcService for OrderBookService { side: cancelled_order.side as i32, order_type: cancelled_order.order_type as i32, status: cancelled_order.status as i32, - timestamp: Some(Timestamp { - seconds: cancelled_order.timestamp.timestamp(), - nanos: cancelled_order.timestamp.timestamp_subsec_nanos() as i32, + timestamp: cancelled_order.timestamp.map(|ts| Timestamp { + seconds: ts.timestamp(), + nanos: ts.timestamp_subsec_nanos() as i32, }), + instrument_id: cancelled_order.instrument_id, + sequence_number: cancelled_order.sequence, + ingress_timestamp_ns: cancelled_order.ingress_timestamp_ns, + idempotency_key: cancelled_order.idempotency_key, })) } async fn get_order_book( &self, request: Request, - ) -> Result, Status> { + ) -> Result, Status> { let depth = request.into_inner().depth as usize; - let (bids, asks) = self.engine - .lock() - .map_err(|_| Status::internal("Lock error"))? - .get_order_book(depth); - - Ok(Response::new(OrderBookResponse { - bids: bids.into_iter() - .map(|(price, qty)| proto::OrderBookLevel { - price: price.to_string(), - quantity: qty.to_string(), - }) - .collect(), - asks: asks.into_iter() - .map(|(price, qty)| proto::OrderBookLevel { - price: price.to_string(), - quantity: qty.to_string(), - }) - .collect(), - })) + let mut bids = Vec::new(); + let mut asks = Vec::new(); + for lane in self.lanes.read().await.values().cloned() { + let (tx, rx) = oneshot::channel(); + lane.send(WorkerCommand::Snapshot { depth, response: tx }) + .await + .map_err(|_| Status::internal("Lane worker unavailable"))?; + let snapshot = rx + .await + .map_err(|_| Status::internal("Lane worker response dropped"))??; + bids.extend(snapshot.bids); + asks.extend(snapshot.asks); + } + Ok(Response::new(proto::OrderBookResponse { bids, asks })) } async fn get_order_status( @@ -138,12 +392,21 @@ impl GrpcService for OrderBookService { ) -> Result, Status> { let order_id = request.into_inner().order_id; - let order = self.engine - .lock() - .map_err(|_| Status::internal("Lock error"))? - .get_order_status(order_id) - .ok_or_else(|| Status::not_found("Order not found"))? - .clone(); + let mut order = None; + for lane in self.lanes.read().await.values().cloned() { + let (tx, rx) = oneshot::channel(); + lane.send(WorkerCommand::Status { + order_id, + response: tx, + }) + .await + .map_err(|_| Status::internal("Lane worker unavailable"))?; + if let Ok(Ok(found)) = rx.await { + order = Some(found); + break; + } + } + let order = order.ok_or_else(|| Status::not_found("Order not found"))?; Ok(Response::new(OrderResponse { id: order.id, @@ -153,10 +416,14 @@ impl GrpcService for OrderBookService { side: order.side as i32, order_type: order.order_type as i32, status: order.status as i32, - timestamp: Some(prost_types::Timestamp { - seconds: order.timestamp.timestamp(), - nanos: order.timestamp.timestamp_subsec_nanos() as i32, + timestamp: order.timestamp.map(|ts| prost_types::Timestamp { + seconds: ts.timestamp(), + nanos: ts.timestamp_subsec_nanos() as i32, }), + instrument_id: order.instrument_id, + sequence_number: order.sequence, + ingress_timestamp_ns: order.ingress_timestamp_ns, + idempotency_key: order.idempotency_key, })) } @@ -164,23 +431,20 @@ impl GrpcService for OrderBookService { async fn get_trade_history(&self, request: Request) -> Result, Status> { let limit = request.into_inner().limit as usize; - let trades = self.engine - .lock() - .map_err(|_| Status::internal("lock error"))? - .get_trade_history(Some(limit)) - .into_iter() - .map(|trade| ProtoTrade{ - maker_order_id: trade.maker_order_id, - taker_order_id: trade.taker_order_id, - price: trade.price.to_string(), - quantity: trade.quantity.to_string(), - side: trade.side as i32, - timestamp: Some(Timestamp { - seconds: trade.timestamp.timestamp(), - nanos: trade.timestamp.timestamp_subsec_nanos() as i32, - }), - }) - .collect(); + let mut trades = Vec::new(); + for lane in self.lanes.read().await.values().cloned() { + let (tx, rx) = oneshot::channel(); + lane.send(WorkerCommand::Trades { + limit, + response: tx, + }) + .await + .map_err(|_| Status::internal("Lane worker unavailable"))?; + let mut lane_trades = rx + .await + .map_err(|_| Status::internal("Lane worker response dropped"))??; + trades.append(&mut lane_trades); + } Ok(Response::new(TradeHistoryResponse { trades })) } diff --git a/atra-ob/src/bin/server.rs b/atra-ob/src/bin/server.rs index dc27509..fab7d84 100644 --- a/atra-ob/src/bin/server.rs +++ b/atra-ob/src/bin/server.rs @@ -1,10 +1,13 @@ use atra_ob::api::service::OrderBookService; -use atra_ob::core::MatchingEngine; +use atra_ob::api::service::SequencerConfig; #[tokio::main] async fn main() -> Result<(), Box> { - let meng = MatchingEngine::new(); - let service = OrderBookService::new(meng); + let lane_count = std::env::var("ATRA_LANE_COUNT") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(4); + let service = OrderBookService::new(lane_count, SequencerConfig::from_env()); println!("Starting order book server on 0.0.0.0:50051"); service.serve("0.0.0.0:50051").await diff --git a/atra-ob/src/core/matchingengine.rs b/atra-ob/src/core/matchingengine.rs index 9a3cd68..afdc736 100644 --- a/atra-ob/src/core/matchingengine.rs +++ b/atra-ob/src/core/matchingengine.rs @@ -97,9 +97,12 @@ impl MatchingEngine { trades.push_back(Trade::new( resting_order.id, matched_order.id, + resting_order.sequence, + matched_order.sequence, price, fill_quantity, matched_order.side, + matched_order.ingress_timestamp_ns, )); // ... reflect the qty changes... diff --git a/atra-ob/src/core/trade_history.rs b/atra-ob/src/core/trade_history.rs index 0ec2767..c27dc65 100644 --- a/atra-ob/src/core/trade_history.rs +++ b/atra-ob/src/core/trade_history.rs @@ -8,21 +8,36 @@ use crate::core::Side; pub struct Trade { pub maker_order_id: u64, pub taker_order_id: u64, + pub maker_sequence: u64, + pub taker_sequence: u64, pub price: Decimal, pub quantity: Decimal, pub side: Side, - pub timestamp: DateTime, + pub timestamp: Option>, + pub ingress_timestamp_ns: Option, } impl Trade { - pub fn new(maker_order_id: u64, taker_order_id: u64, price: Decimal, quantity: Decimal, side: Side) -> Self { + pub fn new( + maker_order_id: u64, + taker_order_id: u64, + maker_sequence: u64, + taker_sequence: u64, + price: Decimal, + quantity: Decimal, + side: Side, + ingress_timestamp_ns: Option, + ) -> Self { Self { maker_order_id, taker_order_id, + maker_sequence, + taker_sequence, price, quantity, side, - timestamp: Utc::now(), + timestamp: None, + ingress_timestamp_ns, } } } diff --git a/atra-ob/src/core/types.rs b/atra-ob/src/core/types.rs index 090593f..b1f34b6 100644 --- a/atra-ob/src/core/types.rs +++ b/atra-ob/src/core/types.rs @@ -1,5 +1,5 @@ -use rust_decimal::Decimal; use chrono::{DateTime, Utc}; +use rust_decimal::Decimal; use std::cmp::Ordering; // @@ -33,24 +33,43 @@ pub enum OrderStatus { #[derive(Debug, Clone, PartialEq, Eq)] // def dont want copy pub struct Order { pub id: u64, + pub instrument_id: u32, + pub sequence: u64, pub price: Decimal, pub quantity: Decimal, pub remaining_quantity: Decimal, pub side: Side, pub order_type: OrderType, pub status: OrderStatus, - pub timestamp: DateTime, + pub timestamp: Option>, + pub ingress_timestamp_ns: Option, + pub idempotency_key: Option, } impl Order { - pub fn new(id: u64, price: Decimal, quantity: Decimal, side: Side, order_type: OrderType) -> Self { - Self { - id, price, quantity, - remaining_quantity: quantity, - side, order_type, - status: OrderStatus::Pending, - timestamp: Utc::now(), - } + pub fn new( + id: u64, + instrument_id: u32, + sequence: u64, + price: Decimal, + quantity: Decimal, + side: Side, + order_type: OrderType, + ) -> Self { + Self { + id, + instrument_id, + sequence, + price, + quantity, + remaining_quantity: quantity, + side, + order_type, + status: OrderStatus::Pending, + timestamp: None, + ingress_timestamp_ns: None, + idempotency_key: None, + } } } @@ -58,7 +77,7 @@ impl Order { impl Ord for Order { fn cmp(&self, other: &Self) -> Ordering { match self.price.cmp(&other.price) { - Ordering::Equal => self.timestamp.cmp(&other.timestamp), + Ordering::Equal => self.sequence.cmp(&other.sequence), ord => match self.side { Side::Bid => ord.reverse(), Side::Ask => ord,