foxhunt/testing/service-load/src/clients/trading_client.rs

//! Trading service gRPC client for load testing

use anyhow::{Context, Result};
use jsonwebtoken::{encode, Algorithm, EncodingKey, Header};
use serde_json::json;
use std::collections::HashMap;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use tonic::metadata::MetadataValue;
use tonic::transport::Channel;
use tonic::Request;

use crate::metrics::LoadTestMetrics;

// Generated proto code
pub mod trading {
    tonic::include_proto!("trading");
}

use trading::{
    trading_service_client::TradingServiceClient, OrderSide, OrderType, StreamMarketDataRequest,
    SubmitOrderRequest,
};

const TEST_ACCOUNT_ID: &str = "load-test-account";
const JWT_SECRET: &str =
    "YZg5/mpqzH0NehGJXiR1yUgUg74HqdOUj/q9tnVSX+gqZvuzHKI1n0NhL4yP8CkUx7WyrVs3X86OSSxIUA6sxQ==";

pub struct TradingClient {
    client: TradingServiceClient<Channel>,
    jwt_token: String,
}

impl TradingClient {
    pub async fn connect(url: &str) -> Result<Self> {
        let channel = Channel::from_shared(url.to_string())
            .context("Invalid URL")?
            .connect()
            .await
            .context("Failed to connect to trading service")?;

        // Generate JWT token for authentication
        let jwt_token = Self::generate_jwt()?;

        Ok(Self {
            client: TradingServiceClient::new(channel),
            jwt_token,
        })
    }

    fn generate_jwt() -> Result<String> {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs();

        let uuid = uuid::Uuid::new_v4();
        let claims = json!({
            "jti": format!("load-test-{uuid}"),
            "sub": "load_test_user",
            "iat": now,
            "exp": now + 3600_u64,
            "iss": "foxhunt-trading",
            "aud": "trading-api",
            "roles": ["trader"],
            "permissions": ["trading.submit_order", "trading.stream_market_data"],
            "token_type": "access",
            "session_id": "load-test-session"
        });

        let header = Header::new(Algorithm::HS256);
        let key = EncodingKey::from_secret(JWT_SECRET.as_ref());

        encode(&header, &claims, &key).context("Failed to encode JWT")
    }

    pub async fn submit_test_order(
        &mut self,
        client_id: usize,
        _order_id: usize,
    ) -> Result<Duration> {
        let start = Instant::now();

        let test_id = client_id % 100;
        let order_request = SubmitOrderRequest {
            symbol: format!("TEST{test_id:04}"),
            side: (OrderSide::Buy as i32),
            quantity: 100.0,
            order_type: (OrderType::Market as i32),
            price: None,
            stop_price: None,
            account_id: TEST_ACCOUNT_ID.to_string(),
            metadata: HashMap::new(),
        };

        let mut request = Request::new(order_request);

        // Add JWT token to metadata
        let jwt_token = &self.jwt_token;
        let token_value =
            MetadataValue::try_from(format!("Bearer {jwt_token}")).context("Invalid JWT token")?;
        request.metadata_mut().insert("authorization", token_value);

        let _ = self.client.submit_order(request).await?;
        Ok(start.elapsed())
    }

    pub async fn run_order_workload(
        &mut self,
        client_id: usize,
        duration_secs: u64,
        metrics: &LoadTestMetrics,
    ) -> Result<()> {
        let start = Instant::now();
        let mut order_count = 0;

        while start.elapsed().as_secs() < duration_secs {
            match self.submit_test_order(client_id, order_count).await {
                Ok(duration) => metrics.record_request(duration, true),
                Err(e) => {
                    tracing::warn!("Order failed: {:?}", e);
                    metrics.record_request(Duration::from_micros(0), false);
                },
            }

            order_count += 1;

            // Rate limiting: ~100 orders/sec per client
            tokio::time::sleep(Duration::from_micros(10_000)).await;
        }

        Ok(())
    }

    pub async fn run_burst_workload(
        &mut self,
        client_id: usize,
        duration_secs: u64,
        metrics: &LoadTestMetrics,
    ) -> Result<()> {
        let start = Instant::now();
        let mut order_count = 0;

        while start.elapsed().as_secs() < duration_secs {
            match self.submit_test_order(client_id, order_count).await {
                Ok(duration) => metrics.record_request(duration, true),
                Err(e) => {
                    tracing::warn!("Burst order failed: {:?}", e);
                    metrics.record_request(Duration::from_micros(0), false);
                },
            }

            order_count += 1;

            // Minimal delay for burst (~100 orders/sec)
            tokio::time::sleep(Duration::from_micros(10_000)).await;
        }

        Ok(())
    }

    pub async fn run_streaming_workload(
        &mut self,
        stream_id: usize,
        duration_secs: u64,
        max_updates: usize,
        metrics: &LoadTestMetrics,
        update_count: &AtomicUsize,
    ) -> Result<()> {
        let stream_sym_id = stream_id % 100;
        let symbols = vec![format!("STREAM{stream_sym_id:04}")];
        let stream_request = StreamMarketDataRequest {
            symbols,
            data_types: vec![], // Empty means all data types
        };

        let mut request = Request::new(stream_request);

        // Add JWT token to metadata
        let jwt_token = &self.jwt_token;
        let token_value =
            MetadataValue::try_from(format!("Bearer {jwt_token}")).context("Invalid JWT token")?;
        request.metadata_mut().insert("authorization", token_value);

        let start = Instant::now();

        match self.client.stream_market_data(request).await {
            Ok(response) => {
                let mut stream = response.into_inner();
                let mut count = 0;

                while let Ok(Some(_event)) = stream.message().await {
                    count += 1;
                    update_count.fetch_add(1_usize, Ordering::Relaxed);

                    if count >= max_updates || start.elapsed().as_secs() >= duration_secs {
                        break;
                    }
                }

                metrics.record_request(start.elapsed(), true);
            },
            Err(e) => {
                tracing::warn!("Stream {} failed: {:?}", stream_id, e);
                metrics.record_request(start.elapsed(), false);
            },
        }

        Ok(())
    }
}