tradingview/live/handler/
command.rs

1use core::fmt;
2use serde::{Deserialize, Serialize};
3use std::{collections::VecDeque, sync::Arc};
4use tokio::{
5    select,
6    task::JoinHandle,
7    time::{Duration, Instant, interval, sleep, timeout},
8};
9use tokio_util::sync::CancellationToken;
10use tracing::{debug, error, info, instrument, warn};
11
12use crate::{
13    Error, Result,
14    error::TradingViewError,
15    live::handler::{CommandRx, Handler, message::*},
16    websocket::WebSocketClient,
17};
18
19/// Priority level for queued commands.
20///
21/// Higher-priority commands are dispatched before lower-priority ones.
22#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
23pub enum CommandPriority {
24    Critical = 3,
25    High = 2,
26    Normal = 1,
27    Low = 0,
28}
29
30/// A command to be dispatched to the TradingView WebSocket session.
31///
32/// Covers all protocol operations: chart sessions, quote subscriptions,
33/// replay mode, Pine Script studies, and connection management.
34#[derive(Debug, Clone, Serialize, Deserialize)]
35pub enum Command {
36    Close,
37    Ping,
38    SendRawMessage(CommandMsg),
39    SetAuthToken(CommandMsg),
40    SetLocale(SetLocaleCommandMsg),
41    SetDataQuality(CommandMsg),
42    SetTimeZone(SetTimeZoneCommandMsg),
43
44    /// Quote Session Commands
45    CreateQuoteSession(CommandMsg),
46    DeleteQuoteSession(CommandMsg),
47    FastSymbols(QuoteCommandMsg),
48    SetQuoteFields(CommandMsg),
49    AddQuoteSymbols(QuoteCommandMsg),
50    RemoveQuoteSymbols(QuoteCommandMsg),
51
52    /// Chart Session Commands
53    CreateChartSession(CommandMsg),
54    DeleteChartSession(CommandMsg),
55    RequestMoreData(ChartDataRequestMsg),
56    RequestMoreTickmarks(ChartDataRequestMsg),
57    CreateChartSeries(ChartSeriesCommandMsg),
58    ModifyChartSeries(ChartSeriesCommandMsg),
59    RemoveSeries(SessionTerminationCommandMsg),
60    ResolveSymbol(ResolveSymbolCommandMsg),
61
62    /// Replay Session Commands
63    CreateReplaySession(CommandMsg),
64    DeleteReplaySession(CommandMsg),
65    AddReplaySeries(AddReplaySeriesCommandMsg),
66    ReplayStep(ReplayStepCommandMsg),
67    ReplayStart(ReplayStartCommandMsg),
68    ReplayStop(ReplayStopCommandMsg),
69    ReplayReset(ReplayResetCommandMsg),
70
71    /// Study Commands
72    CreateStudy(StudyCommandMsg),
73    ModifyStudy(StudyCommandMsg),
74    RemoveStudy(SessionTerminationCommandMsg),
75
76    /// Batch Commands
77    BatchCommands(Vec<Command>),
78
79    /// Conditional Commands
80    ConditionalCommand {
81        condition: CommandCondition,
82        command: Box<Command>,
83        fallback: Option<Box<Command>>,
84    },
85}
86
87impl Command {
88    /// Validate command parameters
89    pub fn validate(&self) -> Result<()> {
90        match self {
91            Command::CreateQuoteSession(msg)
92            | Command::CreateChartSession(msg)
93            | Command::CreateReplaySession(msg) => {
94                if msg.inner.is_empty() {
95                    return Err(Error::Internal("Session name cannot be empty".into()));
96                }
97            }
98            Command::AddQuoteSymbols(msg)
99            | Command::RemoveQuoteSymbols(msg)
100            | Command::FastSymbols(msg) => {
101                if msg.symbols.is_empty() {
102                    return Err(Error::Internal("Symbol list cannot be empty".into()));
103                }
104                if msg.quote_session.is_empty() {
105                    return Err(Error::Internal("Quote session cannot be empty".into()));
106                }
107            }
108            Command::CreateChartSeries(msg) | Command::ModifyChartSeries(msg) => {
109                if msg.chart_session.is_empty()
110                    || msg.series_id.is_empty()
111                    || msg.symbol_series_id.is_empty()
112                {
113                    return Err(Error::Internal(
114                        "Chart series parameters cannot be empty".into(),
115                    ));
116                }
117                // bar_count == 0 is valid in range mode (7-arg form).
118                // In count mode (6-arg form), bar_count must be > 0.
119                if msg.range.is_none() && msg.bar_count == 0 {
120                    return Err(Error::Internal(
121                        "Bar count must be greater than 0 in count mode (range mode allows 0)"
122                            .into(),
123                    ));
124                }
125            }
126            Command::BatchCommands(commands) => {
127                if commands.is_empty() {
128                    return Err(Error::Internal("Batch commands cannot be empty".into()));
129                }
130                if commands.len() > 100 {
131                    return Err(Error::Internal("Batch size too large (max 100)".into()));
132                }
133                for cmd in commands {
134                    cmd.validate()?;
135                }
136            }
137            Command::AddReplaySeries(msg) => {
138                if msg.replay_session.is_empty() || msg.request_id.is_empty() {
139                    return Err(Error::Internal("Replay parameters cannot be empty".into()));
140                }
141            }
142            Command::ReplayStep(msg) => {
143                if msg.replay_session.is_empty() || msg.request_id.is_empty() {
144                    return Err(Error::Internal("Replay parameters cannot be empty".into()));
145                }
146            }
147            Command::ReplayStart(msg) => {
148                if msg.replay_session.is_empty() || msg.request_id.is_empty() {
149                    return Err(Error::Internal("Replay parameters cannot be empty".into()));
150                }
151            }
152            Command::ReplayStop(msg) => {
153                if msg.replay_session.is_empty() || msg.request_id.is_empty() {
154                    return Err(Error::Internal("Replay parameters cannot be empty".into()));
155                }
156            }
157            Command::ReplayReset(msg)
158                if msg.replay_session.is_empty() || msg.request_id.is_empty() =>
159            {
160                return Err(Error::Internal("Replay parameters cannot be empty".into()));
161            }
162            _ => {}
163        }
164        Ok(())
165    }
166
167    /// Get command priority for queue management
168    pub fn priority(&self) -> CommandPriority {
169        match self {
170            Command::Close | Command::SetAuthToken(_) => CommandPriority::Critical,
171            Command::Ping => CommandPriority::High,
172            Command::CreateQuoteSession(_)
173            | Command::CreateChartSession(_)
174            | Command::CreateReplaySession(_) => CommandPriority::High,
175            Command::BatchCommands(_) => CommandPriority::Normal,
176            Command::ConditionalCommand { .. } => CommandPriority::Normal,
177            _ => CommandPriority::Normal,
178        }
179    }
180
181    /// Get estimated execution time for timeout management
182    pub fn estimated_duration(&self) -> Duration {
183        match self {
184            Command::Ping => Duration::from_secs(1),
185            Command::CreateChartSeries(_)
186            | Command::ModifyChartSeries(_)
187            | Command::ResolveSymbol(_) => Duration::from_secs(5),
188            Command::BatchCommands(commands) => Duration::from_millis(commands.len() as u64 * 100),
189            _ => Duration::from_secs(3),
190        }
191    }
192
193    /// Check if command requires session to exist
194    pub fn requires_session(&self) -> Option<&str> {
195        match self {
196            Command::DeleteQuoteSession(msg) | Command::SetQuoteFields(msg) => Some(&msg.inner),
197            Command::AddQuoteSymbols(msg)
198            | Command::RemoveQuoteSymbols(msg)
199            | Command::FastSymbols(msg) => Some(&msg.quote_session),
200            Command::DeleteChartSession(msg) => Some(&msg.inner),
201            Command::CreateChartSeries(msg) | Command::ModifyChartSeries(msg) => {
202                Some(&msg.chart_session)
203            }
204            Command::RemoveSeries(msg) | Command::RemoveStudy(msg) => Some(&msg.chart_session),
205            Command::DeleteReplaySession(msg) => Some(&msg.inner),
206            Command::AddReplaySeries(msg) => Some(&msg.replay_session),
207            Command::ReplayStep(msg) => Some(&msg.replay_session),
208            Command::ReplayStart(msg) => Some(&msg.replay_session),
209            Command::ReplayStop(msg) => Some(&msg.replay_session),
210            Command::ReplayReset(msg) => Some(&msg.replay_session),
211            _ => None,
212        }
213    }
214}
215
216/// Connection state tracking with timestamps for better monitoring
217#[derive(Debug, Clone, PartialEq, Copy, Eq)]
218pub struct ConnectionState {
219    pub status: ConnectionStatus,
220    pub last_change: Instant,
221    pub last_successful_operation: Option<Instant>,
222}
223
224/// Current state of the WebSocket connection.
225#[derive(Debug, Clone, PartialEq, Copy, Eq)]
226pub enum ConnectionStatus {
227    Connected,
228    Disconnected,
229    Reconnecting,
230    Shutdown,
231}
232
233impl ConnectionState {
234    fn new(status: ConnectionStatus) -> Self {
235        Self {
236            status,
237            last_change: Instant::now(),
238            last_successful_operation: None,
239        }
240    }
241
242    fn transition_to(&mut self, new_status: ConnectionStatus) {
243        if self.status != new_status {
244            self.status = new_status;
245            self.last_change = Instant::now();
246        }
247    }
248
249    fn mark_successful_operation(&mut self) {
250        self.last_successful_operation = Some(Instant::now());
251    }
252
253    fn time_since_last_success(&self) -> Option<Duration> {
254        self.last_successful_operation.map(|t| t.elapsed())
255    }
256}
257
258/// Error classification for better handling
259#[derive(Debug, Clone, PartialEq, Copy, Eq)]
260enum ErrorSeverity {
261    /// Temporary errors that should trigger reconnection
262    Recoverable,
263    /// Permanent errors that should stop the runner
264    Fatal,
265    /// Command-specific errors that don't affect connection
266    CommandOnly,
267}
268
269#[derive(Debug, Clone, Copy)]
270struct ExponentialBackoff {
271    config: BackoffConfig,
272    current_delay: Duration,
273    attempts: usize,
274}
275
276#[derive(Debug, Clone, Copy)]
277pub struct BackoffConfig {
278    initial_delay: Duration,
279    max_delay: Duration,
280    max_attempts: usize,
281    multiplier: f64,
282    jitter_percent: f64,
283}
284
285impl Default for BackoffConfig {
286    fn default() -> Self {
287        Self {
288            initial_delay: Duration::from_millis(1000),
289            max_delay: Duration::from_secs(60),
290            max_attempts: 10,
291            multiplier: 2.0,
292            jitter_percent: 0.1,
293        }
294    }
295}
296
297impl ExponentialBackoff {
298    fn new(config: BackoffConfig) -> Self {
299        Self {
300            current_delay: config.initial_delay,
301            config,
302            attempts: 0,
303        }
304    }
305
306    fn next_backoff(&mut self) -> Option<Duration> {
307        if self.attempts >= self.config.max_attempts {
308            return None;
309        }
310
311        let mut delay = self.current_delay;
312        self.attempts += 1;
313
314        // Add jitter to prevent thundering herd problem
315        let jitter_range = delay.as_millis() as f64 * self.config.jitter_percent;
316        let jitter = (rand::random::<f64>() - 0.5) * 2.0 * jitter_range;
317        delay = Duration::from_millis((delay.as_millis() as f64 + jitter).max(0.0) as u64);
318
319        // Exponential backoff with cap
320        self.current_delay = std::cmp::min(
321            Duration::from_millis(
322                (self.current_delay.as_millis() as f64 * self.config.multiplier) as u64,
323            ),
324            self.config.max_delay,
325        );
326
327        Some(delay)
328    }
329
330    fn reset(&mut self) {
331        self.current_delay = self.config.initial_delay;
332        self.attempts = 0;
333    }
334
335    fn remaining_attempts(&self) -> usize {
336        self.config.max_attempts.saturating_sub(self.attempts)
337    }
338}
339
340#[derive(Debug, Clone)]
341pub struct CommandQueue {
342    critical_queue: VecDeque<Command>,
343    high_queue: VecDeque<Command>,
344    normal_queue: VecDeque<Command>,
345    low_queue: VecDeque<Command>,
346    max_size: usize,
347    dropped_count: u64,
348    session_tracker: std::collections::HashSet<String>,
349}
350
351impl CommandQueue {
352    fn new(max_size: usize) -> Self {
353        Self {
354            critical_queue: VecDeque::new(),
355            high_queue: VecDeque::new(),
356            normal_queue: VecDeque::new(),
357            low_queue: VecDeque::new(),
358            max_size,
359            dropped_count: 0,
360            session_tracker: std::collections::HashSet::new(),
361        }
362    }
363
364    fn enqueue(&mut self, cmd: Command) -> Result<()> {
365        // Validate command first
366        cmd.validate()?;
367
368        // Check if command requires existing session
369        if let Some(session) = cmd.requires_session()
370            && !self.session_tracker.contains(session)
371        {
372            return Err(Error::Internal(
373                format!("Session '{}' does not exist", session).into(),
374            ));
375        }
376
377        // Track session creation/deletion
378        match &cmd {
379            Command::CreateQuoteSession(msg)
380            | Command::CreateChartSession(msg)
381            | Command::CreateReplaySession(msg) => {
382                self.session_tracker.insert(msg.inner.to_string());
383            }
384            Command::DeleteQuoteSession(msg)
385            | Command::DeleteChartSession(msg)
386            | Command::DeleteReplaySession(msg) => {
387                self.session_tracker.remove(&msg.inner.to_string());
388            }
389            _ => {}
390        }
391
392        // Check capacity and drop if necessary
393        if self.total_len() >= self.max_size {
394            // Try to drop from lower priority queues first
395            if self.drop_lowest_priority() {
396                self.dropped_count += 1;
397                warn!(
398                    "Dropped command due to queue overflow (total dropped: {})",
399                    self.dropped_count
400                );
401            } else {
402                return Err(Error::Internal("Command queue is full".into()));
403            }
404        }
405
406        let queue = match cmd.priority() {
407            CommandPriority::Critical => &mut self.critical_queue,
408            CommandPriority::High => &mut self.high_queue,
409            CommandPriority::Normal => &mut self.normal_queue,
410            CommandPriority::Low => &mut self.low_queue,
411        };
412
413        queue.push_back(cmd);
414        Ok(())
415    }
416
417    fn drop_lowest_priority(&mut self) -> bool {
418        if self.low_queue.pop_front().is_some() {
419            return true;
420        }
421        if self.normal_queue.pop_front().is_some() {
422            return true;
423        }
424        if self.high_queue.pop_front().is_some() {
425            return true;
426        }
427        false
428    }
429
430    fn dequeue(&mut self) -> Option<Command> {
431        self.critical_queue
432            .pop_front()
433            .or_else(|| self.high_queue.pop_front())
434            .or_else(|| self.normal_queue.pop_front())
435            .or_else(|| self.low_queue.pop_front())
436    }
437
438    fn drain(&mut self) -> Vec<Command> {
439        let mut commands = Vec::with_capacity(self.total_len());
440
441        while let Some(cmd) = self.dequeue() {
442            commands.push(cmd);
443        }
444
445        commands
446    }
447
448    fn total_len(&self) -> usize {
449        self.critical_queue.len()
450            + self.high_queue.len()
451            + self.normal_queue.len()
452            + self.low_queue.len()
453    }
454
455    fn is_empty(&self) -> bool {
456        self.total_len() == 0
457    }
458
459    fn clear(&mut self) {
460        self.critical_queue.clear();
461        self.high_queue.clear();
462        self.normal_queue.clear();
463        self.low_queue.clear();
464        self.session_tracker.clear();
465    }
466
467    fn detailed_stats(&self) -> CommandQueueStats {
468        CommandQueueStats {
469            critical_queue_len: self.critical_queue.len(),
470            high_queue_len: self.high_queue.len(),
471            normal_queue_len: self.normal_queue.len(),
472            low_queue_len: self.low_queue.len(),
473            total_len: self.total_len(),
474            max_capacity: self.max_size,
475            dropped_count: self.dropped_count,
476            active_sessions: self.session_tracker.len(),
477        }
478    }
479}
480
481#[derive(Debug, Clone, Copy)]
482pub struct CommandQueueStats {
483    pub critical_queue_len: usize,
484    pub high_queue_len: usize,
485    pub normal_queue_len: usize,
486    pub low_queue_len: usize,
487    pub total_len: usize,
488    pub max_capacity: usize,
489    pub dropped_count: u64,
490    pub active_sessions: usize,
491}
492
493impl fmt::Display for CommandQueueStats {
494    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
495        write!(
496            f,
497            "Queue: C:{} H:{} N:{} L:{} | Total:{}/{} | Sessions:{} | Dropped:{}",
498            self.critical_queue_len,
499            self.high_queue_len,
500            self.normal_queue_len,
501            self.low_queue_len,
502            self.total_len,
503            self.max_capacity,
504            self.active_sessions,
505            self.dropped_count
506        )
507    }
508}
509
510#[derive(Debug, Default, Clone)]
511pub struct ConnectionStats {
512    pub reconnect_attempts: u64,
513    pub successful_reconnects: u64,
514    pub commands_processed: u64,
515    pub commands_failed: u64,
516    pub errors_handled: u64,
517    pub total_uptime: Duration,
518    pub last_reconnect_duration: Option<Duration>,
519    pub average_command_latency: Option<Duration>,
520}
521
522impl ConnectionStats {
523    fn record_command_success(&mut self, duration: Duration) {
524        self.commands_processed += 1;
525        self.average_command_latency = Some(
526            self.average_command_latency
527                .map(|avg| (avg + duration) / 2)
528                .unwrap_or(duration),
529        );
530    }
531
532    fn record_command_failure(&mut self) {
533        self.commands_failed += 1;
534    }
535
536    fn record_reconnect_attempt(&mut self) {
537        self.reconnect_attempts += 1;
538    }
539
540    fn record_successful_reconnect(&mut self, duration: Duration) {
541        self.successful_reconnects += 1;
542        self.last_reconnect_duration = Some(duration);
543    }
544
545    fn success_rate(&self) -> f64 {
546        let total = self.commands_processed + self.commands_failed;
547        if total == 0 {
548            return 1.0;
549        }
550        self.commands_processed as f64 / total as f64
551    }
552}
553
554/// Enhanced configuration for CommandRunner
555#[derive(Debug, Clone)]
556pub struct CommandRunnerConfig {
557    pub heartbeat_interval: Duration,
558    pub health_check_interval: Duration,
559    pub command_timeout: Duration,
560    pub reconnect_timeout: Duration,
561    pub max_queue_size: usize,
562    pub backoff_config: BackoffConfig,
563    pub health_check_timeout: Duration,
564}
565
566impl Default for CommandRunnerConfig {
567    fn default() -> Self {
568        Self {
569            heartbeat_interval: Duration::from_secs(30),
570            health_check_interval: Duration::from_secs(5),
571            command_timeout: Duration::from_secs(10),
572            reconnect_timeout: Duration::from_secs(30),
573            max_queue_size: 100,
574            backoff_config: BackoffConfig::default(),
575            health_check_timeout: Duration::from_secs(60),
576        }
577    }
578}
579
580/// Central command dispatcher for the WebSocket session.
581///
582/// Reads [`Command`]s from an MPSC channel and dispatches them to the
583/// [`WebSocketClient`]. Manages connection lifecycle, reconnection, and
584/// error recovery.
585///
586/// Spawn via [`CommandRunner::run()`] — it blocks until shutdown is triggered.
587///
588/// [`Command`]: crate::live::handler::command::Command
589/// [`WebSocketClient`]: crate::websocket::WebSocketClient
590/// [`CommandRunner::run()`]: CommandRunner::run()
591pub struct CommandRunner<T: Handler> {
592    rx: CommandRx,
593    ws: Arc<WebSocketClient<T>>,
594    shutdown: CancellationToken,
595    state: ConnectionState,
596    command_queue: CommandQueue,
597    stats: ConnectionStats,
598    reader_handle: Option<JoinHandle<()>>,
599    config: CommandRunnerConfig,
600    start_time: Instant,
601}
602
603impl<T: Handler> CommandRunner<T> {
604    pub fn new(rx: CommandRx, ws: Arc<WebSocketClient<T>>) -> Self {
605        Self::with_config(rx, ws, CommandRunnerConfig::default())
606    }
607
608    pub fn with_config(
609        rx: CommandRx,
610        ws: Arc<WebSocketClient<T>>,
611        config: CommandRunnerConfig,
612    ) -> Self {
613        Self {
614            rx,
615            ws,
616            shutdown: CancellationToken::new(),
617            state: ConnectionState::new(ConnectionStatus::Connected),
618            command_queue: CommandQueue::new(config.max_queue_size),
619            stats: ConnectionStats::default(),
620            reader_handle: None,
621            config,
622            start_time: Instant::now(),
623        }
624    }
625
626    // #[instrument(skip(self), fields(runner_id = %std::ptr::addr_of!(*self) as usize))]
627    pub async fn run(mut self) -> Result<()> {
628        let mut hb = interval(self.config.heartbeat_interval);
629        let mut backoff = ExponentialBackoff::new(self.config.backoff_config);
630        let mut health_check = interval(self.config.health_check_interval);
631        let mut stats_timer = interval(Duration::from_secs(60)); // Log stats every minute
632
633        info!("CommandRunner started with config: {:?}", self.config);
634
635        // Initialize connection
636        if let Err(e) = self.initialize_connection().await {
637            error!("Failed to initialize connection: {}", e);
638            self.state.transition_to(ConnectionStatus::Disconnected);
639        }
640
641        loop {
642            select! {
643                biased;
644
645                _ = self.shutdown.cancelled() => {
646                    info!("Shutdown signal received");
647                    self.state.transition_to(ConnectionStatus::Shutdown);
648                    break;
649                },
650
651                cmd = self.rx.recv() => match cmd {
652                    Some(cmd) => {
653                        let start = Instant::now();
654                        match self.handle_command(cmd, &mut backoff).await {
655                            Ok(_) => {
656                                self.stats.record_command_success(start.elapsed());
657                                self.state.mark_successful_operation();
658                            },
659                            Err(e) => {
660                                self.stats.record_command_failure();
661                                self.stats.errors_handled += 1;
662                                self.handle_command_error(e).await;
663                            }
664                        }
665                    },
666                    None => {
667                        info!("All command senders dropped, shutting down");
668                        break;
669                    }
670                },
671
672                _ = health_check.tick() => {
673                    // Only perform health check if we're supposed to be connected
674                    if matches!(self.state.status, ConnectionStatus::Connected | ConnectionStatus::Reconnecting) {
675                        if let Err(e) = self.perform_health_check().await {
676                            warn!("Health check failed: {}", e);
677                            self.state.transition_to(ConnectionStatus::Disconnected);
678                        } else {
679                            self.state.mark_successful_operation();
680                        }
681                    }
682                },
683
684                _ = hb.tick() => {
685                    if self.state.status == ConnectionStatus::Connected {
686                        if let Err(e) = self.send_heartbeat().await {
687                            warn!("Heartbeat failed: {}", e);
688                            self.state.transition_to(ConnectionStatus::Disconnected);
689                        } else {
690                            self.state.mark_successful_operation();
691                        }
692                    }
693                },
694
695                _ = stats_timer.tick() => {
696                    self.log_stats();
697                },
698            }
699
700            // Handle disconnection state
701            if self.state.status == ConnectionStatus::Disconnected
702                && let Err(e) = self.handle_reconnection(&mut backoff).await
703            {
704                error!("Reconnection failed: {}", e);
705                break;
706            }
707        }
708
709        self.cleanup().await;
710
711        Ok(())
712    }
713
714    /// Comprehensive health check that actively tests the connection
715    async fn perform_health_check(&self) -> Result<()> {
716        // First check basic connection state
717        if self.ws.is_closed() {
718            return Err(Error::Internal("WebSocket is closed".into()));
719        }
720
721        // Check if connection has been unhealthy for too long
722        if let Some(time_since_success) = self.state.time_since_last_success()
723            && time_since_success > self.config.health_check_timeout
724        {
725            return Err(Error::Internal(
726                format!("No successful operations for {time_since_success:?}").into(),
727            ));
728        }
729
730        // Actively test the connection with a ping
731        // Use a shorter timeout for health check pings
732        let health_check_timeout = Duration::from_secs(5);
733        match timeout(health_check_timeout, self.ws.try_ping()).await {
734            Ok(Ok(_)) => {
735                debug!("Health check ping successful");
736                Ok(())
737            }
738            Ok(Err(e)) => {
739                warn!("Health check ping failed: {}", e);
740                Err(Error::Internal(
741                    format!("Health check ping failed: {e}").into(),
742                ))
743            }
744            Err(_) => {
745                warn!(
746                    "Health check ping timed out after {:?}",
747                    health_check_timeout
748                );
749                Err(Error::Internal("Health check ping timeout".into()))
750            }
751        }
752    }
753
754    #[instrument(skip(self))]
755    async fn initialize_connection(&mut self) -> Result<()> {
756        // Start the WebSocket reader task
757        self.start_reader_task();
758
759        // Send initial authentication if we have a token
760        let auth_token = *self.ws.auth_token.read().await;
761        if auth_token != "unauthorized_user_token" {
762            info!("Sending initial authentication token");
763            self.ws.set_auth_token(&auth_token).await?;
764        }
765
766        self.state.mark_successful_operation();
767        Ok(())
768    }
769
770    async fn handle_command_error(&mut self, error: Error) {
771        match self.classify_error(&error) {
772            ErrorSeverity::Fatal => {
773                error!("Fatal error, shutting down: {}", error);
774                self.state.transition_to(ConnectionStatus::Shutdown);
775            }
776            ErrorSeverity::Recoverable => {
777                warn!("Recoverable error, will attempt reconnection: {}", error);
778                self.state.transition_to(ConnectionStatus::Disconnected);
779            }
780            ErrorSeverity::CommandOnly => {
781                warn!("Command error (continuing): {}", error);
782            }
783        }
784    }
785
786    fn start_reader_task(&mut self) {
787        if self.reader_handle.is_none() {
788            let ws = Arc::clone(&self.ws);
789            let shutdown = self.shutdown.clone();
790
791            self.reader_handle = Some(tokio::spawn(async move {
792                info!("Starting WebSocket reader task");
793
794                loop {
795                    select! {
796                        _ = shutdown.cancelled() => {
797                            info!("WebSocket reader task shutting down");
798                            break;
799                        },
800                        result = ws.subscribe() => {
801                            match result {
802                                Ok(_) => {
803                                    info!("WebSocket reader task completed successfully");
804                                    break;
805                                },
806                                Err(e) => {
807                                    error!("WebSocket reader task failed: {}", e);
808                                    // Small delay before retrying
809                                    sleep(Duration::from_secs(1)).await;
810                                }
811                            }
812                        }
813                    }
814                }
815            }));
816        }
817    }
818
819    #[instrument(skip(self, cmd, backoff))]
820    async fn handle_command(
821        &mut self,
822        cmd: Command,
823        backoff: &mut ExponentialBackoff,
824    ) -> Result<()> {
825        match self.state.status {
826            ConnectionStatus::Connected => self.process_command(cmd).await,
827            ConnectionStatus::Reconnecting => {
828                if let Ok(()) = self.command_queue.enqueue(cmd) {
829                    let stats = self.command_queue.detailed_stats();
830                    info!("Command queued during reconnection: {}", stats);
831                }
832                Ok(())
833            }
834            ConnectionStatus::Disconnected => {
835                if self.is_critical_command(&cmd) {
836                    warn!(
837                        "Critical command received while disconnected, attempting immediate reconnection"
838                    );
839                    self.handle_reconnection(backoff).await?;
840                    self.process_command(cmd).await
841                } else {
842                    self.command_queue.enqueue(cmd)?;
843                    Ok(())
844                }
845            }
846            ConnectionStatus::Shutdown => {
847                warn!("Ignoring command after shutdown");
848                Ok(())
849            }
850        }
851    }
852
853    #[instrument(skip(self, cmd), fields(command_type = ?std::mem::discriminant(&cmd)))]
854    async fn process_command(&mut self, cmd: Command) -> Result<()> {
855        // Validate command before processing
856        cmd.validate()?;
857
858        // Use an iterative approach for nested commands
859        let mut command_stack = VecDeque::new();
860        command_stack.push_back(cmd);
861
862        while let Some(current_cmd) = command_stack.pop_front() {
863            match current_cmd {
864                Command::BatchCommands(commands) => {
865                    // Add all batch commands to the front of the stack in reverse order
866                    // so they're processed in the correct order
867                    for cmd in commands.into_iter().rev() {
868                        command_stack.push_front(cmd);
869                    }
870                    continue;
871                }
872                Command::ConditionalCommand {
873                    condition,
874                    command,
875                    fallback,
876                } => {
877                    let condition_met = self.evaluate_condition(&condition).await;
878
879                    if condition_met {
880                        debug!("Condition met, executing primary command");
881                        command_stack.push_front(*command);
882                    } else if let Some(fallback_cmd) = fallback {
883                        debug!("Condition not met, executing fallback command");
884                        command_stack.push_front(*fallback_cmd);
885                    } else {
886                        debug!("Condition not met and no fallback provided");
887                    }
888                    continue;
889                }
890                _ => {
891                    // Process regular command
892                    self.execute_single_command(current_cmd).await?;
893                }
894            }
895        }
896
897        Ok(())
898    }
899
900    async fn execute_single_command(&mut self, cmd: Command) -> Result<()> {
901        // Use dynamic timeout based on command type
902        let timeout_duration = cmd.estimated_duration().mul_f32(1.5); // 50% buffer
903
904        let result = timeout(timeout_duration, async {
905            match cmd {
906                Command::Close => self.ws.close().await,
907                Command::Ping => self.ws.try_ping().await,
908                Command::SetAuthToken(auth_token) => {
909                    self.ws.set_auth_token(&auth_token.inner).await
910                }
911                Command::CreateQuoteSession(session) => {
912                    self.ws.create_quote_session(&session.inner).await
913                }
914                Command::SetLocale(locale) => {
915                    self.ws.set_locale(&locale.language, &locale.country).await
916                }
917                Command::SetDataQuality(quality) => self.ws.set_data_quality(&quality.inner).await,
918                Command::SetTimeZone(timezone) => {
919                    self.ws
920                        .set_timezone(&timezone.chart_session, timezone.timezone)
921                        .await
922                }
923                Command::CreateChartSession(session) => {
924                    self.ws.create_chart_session(&session.inner).await
925                }
926                Command::DeleteChartSession(session) => {
927                    self.ws.delete_chart_session(&session.inner).await
928                }
929                Command::RequestMoreData(request) => {
930                    self.ws
931                        .request_more_data(&request.chart_session, &request.series_id, request.num)
932                        .await
933                }
934                Command::RequestMoreTickmarks(request) => {
935                    self.ws
936                        .request_more_tickmarks(
937                            &request.chart_session,
938                            &request.series_id,
939                            request.num,
940                        )
941                        .await
942                }
943                Command::SendRawMessage(command_msg) => {
944                    self.ws.send_raw_message(&command_msg.inner).await
945                }
946                Command::DeleteQuoteSession(session) => {
947                    self.ws.delete_quote_session(&session.inner).await
948                }
949                Command::FastSymbols(quote_command_msg) => {
950                    self.ws
951                        .fast_symbols(
952                            &quote_command_msg.quote_session,
953                            &quote_command_msg
954                                .symbols
955                                .iter()
956                                .map(|s| s.as_str())
957                                .collect::<Vec<_>>(),
958                        )
959                        .await
960                }
961                Command::SetQuoteFields(command_msg) => {
962                    self.ws.set_fields(&command_msg.inner).await
963                }
964                Command::AddQuoteSymbols(quote_command_msg) => {
965                    self.ws
966                        .add_symbols(
967                            &quote_command_msg.quote_session,
968                            &quote_command_msg
969                                .symbols
970                                .iter()
971                                .map(|s| s.as_str())
972                                .collect::<Vec<_>>(),
973                        )
974                        .await
975                }
976                Command::RemoveQuoteSymbols(quote_command_msg) => {
977                    self.ws
978                        .remove_symbols(
979                            &quote_command_msg.quote_session,
980                            &quote_command_msg
981                                .symbols
982                                .iter()
983                                .map(|s| s.as_str())
984                                .collect::<Vec<_>>(),
985                        )
986                        .await
987                }
988                Command::CreateChartSeries(chart_series_command_msg) => {
989                    self.ws
990                        .create_series()
991                        .chart_session(&chart_series_command_msg.chart_session)
992                        .series_identifier(&chart_series_command_msg.series_identifier)
993                        .series_id(&chart_series_command_msg.series_id)
994                        .symbol_series_id(&chart_series_command_msg.symbol_series_id)
995                        .interval(chart_series_command_msg.interval)
996                        .bar_count(chart_series_command_msg.bar_count)
997                        .maybe_range(chart_series_command_msg.range)
998                        .call()
999                        .await
1000                }
1001                Command::ModifyChartSeries(command) => {
1002                    self.ws
1003                        .modify_series()
1004                        .chart_session(&command.chart_session)
1005                        .series_identifier(&command.series_identifier)
1006                        .series_id(&command.series_id)
1007                        .symbol_series_id(&command.symbol_series_id)
1008                        .interval(command.interval)
1009                        .bar_count(command.bar_count)
1010                        .maybe_range(command.range)
1011                        .call()
1012                        .await
1013                }
1014                Command::RemoveSeries(command) => {
1015                    self.ws
1016                        .remove_series(&command.chart_session, &command.id)
1017                        .await
1018                }
1019                Command::ResolveSymbol(command) => {
1020                    self.ws
1021                        .resolve_symbol()
1022                        .session(&command.session)
1023                        .symbol_series_id(&command.symbol_series_id)
1024                        .maybe_adjustment(command.adjustment)
1025                        .maybe_currency(command.currency)
1026                        .maybe_session_type(command.session_type)
1027                        .maybe_replay_session(command.replay_session.as_deref())
1028                        .instrument(&command.instrument)
1029                        .call()
1030                        .await
1031                }
1032                Command::CreateReplaySession(command_msg) => {
1033                    self.ws.create_replay_session(&command_msg.inner).await
1034                }
1035                Command::DeleteReplaySession(command) => {
1036                    self.ws.delete_replay_session(&command.inner).await
1037                }
1038                Command::AddReplaySeries(command) => {
1039                    self.ws
1040                        .add_replay_series()
1041                        .replay_session(&command.replay_session)
1042                        .request_id(&command.request_id)
1043                        .instrument(&command.instrument)
1044                        .maybe_adjustment(command.adjustment)
1045                        .maybe_session_type(command.session_type)
1046                        .maybe_currency(command.currency)
1047                        .interval(command.interval)
1048                        .call()
1049                        .await
1050                }
1051                Command::ReplayStep(command) => {
1052                    self.ws
1053                        .replay_step(
1054                            &command.replay_session,
1055                            &command.request_id,
1056                            command.step as u64,
1057                        )
1058                        .await
1059                }
1060                Command::ReplayStart(command) => {
1061                    self.ws
1062                        .replay_start(
1063                            &command.replay_session,
1064                            &command.request_id,
1065                            command.interval,
1066                        )
1067                        .await
1068                }
1069                Command::ReplayStop(command) => {
1070                    self.ws
1071                        .replay_stop(&command.replay_session, &command.request_id)
1072                        .await
1073                }
1074                Command::ReplayReset(command) => {
1075                    self.ws
1076                        .replay_reset(
1077                            &command.replay_session,
1078                            &command.request_id,
1079                            command.timestamp,
1080                        )
1081                        .await
1082                }
1083                Command::CreateStudy(command) => {
1084                    self.ws
1085                        .create_study()
1086                        .chart_session(&command.chart_session)
1087                        .study_ids(
1088                            &command
1089                                .study_ids
1090                                .iter()
1091                                .map(|s| s.as_str())
1092                                .collect::<Vec<_>>()
1093                                .try_into()
1094                                .unwrap(),
1095                        )
1096                        .chart_series_id(&command.chart_series_id)
1097                        .study(command.study.clone())
1098                        .call()
1099                        .await
1100                }
1101                Command::ModifyStudy(command) => {
1102                    self.ws
1103                        .modify_study()
1104                        .chart_session(&command.chart_session)
1105                        .study_ids(
1106                            &command
1107                                .study_ids
1108                                .iter()
1109                                .map(|s| s.as_str())
1110                                .collect::<Vec<_>>()
1111                                .try_into()
1112                                .expect("Study IDs must be exactly 2"),
1113                        )
1114                        .study(command.study.clone())
1115                        .call()
1116                        .await
1117                }
1118                Command::RemoveStudy(session_termination_command_msg) => {
1119                    self.ws
1120                        .remove_study(
1121                            &session_termination_command_msg.chart_session,
1122                            &session_termination_command_msg.id,
1123                        )
1124                        .await
1125                }
1126                _ => Ok(()),
1127            }
1128        })
1129        .await;
1130
1131        result.unwrap_or_else(|_| Err(Error::Internal("Command timeout".into())))
1132    }
1133
1134    async fn evaluate_condition(&self, condition: &CommandCondition) -> bool {
1135        match condition {
1136            CommandCondition::SessionExists(session) => self
1137                .command_queue
1138                .session_tracker
1139                .contains(session.as_str()),
1140            CommandCondition::SymbolResolved(_symbol) => {
1141                // This would require maintaining state of resolved symbols
1142                true // Placeholder
1143            }
1144            CommandCondition::ConnectionHealthy => {
1145                !self.ws.is_closed() && matches!(self.state.status, ConnectionStatus::Connected)
1146            }
1147            CommandCondition::QueueEmpty => self.command_queue.is_empty(),
1148        }
1149    }
1150
1151    fn is_critical_command(&self, cmd: &Command) -> bool {
1152        matches!(
1153            cmd,
1154            Command::SetAuthToken { .. }
1155                | Command::Close
1156                | Command::Ping
1157                | Command::CreateQuoteSession { .. }
1158                | Command::CreateChartSession { .. }
1159                | Command::CreateReplaySession { .. }
1160        )
1161    }
1162
1163    #[instrument(skip(self, backoff))]
1164    async fn handle_reconnection(&mut self, backoff: &mut ExponentialBackoff) -> Result<()> {
1165        if self.state.status == ConnectionStatus::Shutdown {
1166            return Ok(());
1167        }
1168
1169        self.state.transition_to(ConnectionStatus::Reconnecting);
1170        let reconnect_start = Instant::now();
1171
1172        // Stop the old reader task gracefully
1173        self.stop_reader_task().await;
1174
1175        while let Some(delay) = backoff.next_backoff() {
1176            self.stats.record_reconnect_attempt();
1177
1178            info!(
1179                "Attempting reconnection in {:?} (attempt {}/{}, remaining: {})",
1180                delay,
1181                backoff.attempts,
1182                backoff.config.max_attempts,
1183                backoff.remaining_attempts()
1184            );
1185
1186            sleep(delay).await;
1187
1188            if self.shutdown.is_cancelled() {
1189                self.state.transition_to(ConnectionStatus::Shutdown);
1190                return Ok(());
1191            }
1192
1193            match timeout(self.config.reconnect_timeout, self.ws.reconnect()).await {
1194                Ok(Ok(_)) => {
1195                    let reconnect_duration = reconnect_start.elapsed();
1196                    info!(
1197                        "WebSocket reconnected successfully in {:?}",
1198                        reconnect_duration
1199                    );
1200
1201                    self.state.transition_to(ConnectionStatus::Connected);
1202                    self.state.mark_successful_operation();
1203                    self.stats.record_successful_reconnect(reconnect_duration);
1204                    backoff.reset();
1205
1206                    // Restart reader task and process queued commands
1207                    self.start_reader_task();
1208                    self.process_queued_commands().await;
1209                    return Ok(());
1210                }
1211                Ok(Err(e)) => {
1212                    warn!("Reconnection attempt {} failed: {}", backoff.attempts, e);
1213                }
1214                Err(_) => {
1215                    warn!(
1216                        "Reconnection attempt {} timed out after {:?}",
1217                        backoff.attempts, self.config.reconnect_timeout
1218                    );
1219                }
1220            }
1221        }
1222
1223        error!(
1224            "Reconnection backoff exhausted after {} attempts",
1225            backoff.config.max_attempts
1226        );
1227        self.state.transition_to(ConnectionStatus::Shutdown);
1228        Err(Error::Internal(
1229            "Reconnection failed after maximum attempts".into(),
1230        ))
1231    }
1232
1233    async fn stop_reader_task(&mut self) {
1234        if let Some(handle) = self.reader_handle.take() {
1235            handle.abort();
1236            let _ = timeout(Duration::from_secs(2), handle).await;
1237        }
1238    }
1239
1240    #[instrument(skip(self))]
1241    async fn process_queued_commands(&mut self) {
1242        let commands = self.command_queue.drain();
1243        if !commands.is_empty() {
1244            info!("Processing {} queued commands", commands.len());
1245
1246            let mut successful = 0;
1247            let mut failed = 0;
1248
1249            for cmd in commands {
1250                match timeout(self.config.command_timeout, self.process_command(cmd)).await {
1251                    Ok(Ok(_)) => {
1252                        successful += 1;
1253                        self.state.mark_successful_operation();
1254                    }
1255                    Ok(Err(e)) => {
1256                        failed += 1;
1257                        error!("Failed to process queued command: {}", e);
1258                    }
1259                    Err(_) => {
1260                        failed += 1;
1261                        error!("Queued command timed out");
1262                    }
1263                }
1264            }
1265
1266            info!(
1267                "Processed queued commands: {} successful, {} failed",
1268                successful, failed
1269            );
1270        }
1271    }
1272
1273    /// Send heartbeat ping (separate from health check)
1274    async fn send_heartbeat(&self) -> Result<()> {
1275        let heartbeat_timeout = Duration::from_secs(10);
1276        match timeout(heartbeat_timeout, self.ws.try_ping()).await {
1277            Ok(Ok(_)) => {
1278                debug!("Heartbeat ping successful");
1279                Ok(())
1280            }
1281            Ok(Err(e)) => {
1282                warn!("Heartbeat ping failed: {}", e);
1283                Err(Error::Internal(format!("Heartbeat failed: {e}").into()))
1284            }
1285            Err(_) => {
1286                warn!("Heartbeat ping timed out after {:?}", heartbeat_timeout);
1287                Err(Error::Internal("Heartbeat timeout".into()))
1288            }
1289        }
1290    }
1291
1292    fn classify_error(&self, error: &Error) -> ErrorSeverity {
1293        use Error::*;
1294
1295        match error {
1296            WebSocket(_) => ErrorSeverity::Recoverable,
1297            TradingView {
1298                source: TradingViewError::ProtocolError | TradingViewError::CriticalError,
1299            } => ErrorSeverity::Recoverable,
1300            JsonParse(_) | UrlParse(_) => ErrorSeverity::CommandOnly,
1301            Internal(msg) => {
1302                let msg_lower = msg.to_lowercase();
1303                if msg_lower.contains("timeout")
1304                    || msg_lower.contains("connection")
1305                    || msg_lower.contains("network")
1306                {
1307                    ErrorSeverity::Recoverable
1308                } else if msg_lower.contains("shutdown") || msg_lower.contains("cancelled") {
1309                    ErrorSeverity::Fatal
1310                } else {
1311                    ErrorSeverity::CommandOnly
1312                }
1313            }
1314            _ => ErrorSeverity::CommandOnly,
1315        }
1316    }
1317
1318    fn log_stats(&mut self) {
1319        self.stats.total_uptime = self.start_time.elapsed();
1320        info!(
1321            "Connection Stats: {:?}, Success Rate: {:.2}%, Uptime: {:?}",
1322            self.stats,
1323            self.stats.success_rate() * 100.0,
1324            self.stats.total_uptime
1325        );
1326    }
1327
1328    #[instrument(skip(self))]
1329    async fn cleanup(&mut self) {
1330        info!("Cleaning up CommandRunner");
1331
1332        // Cancel shutdown token to stop all tasks
1333        self.shutdown.cancel();
1334
1335        // Stop reader task
1336        self.stop_reader_task().await;
1337
1338        let remaining_commands = self.command_queue.total_len();
1339        if remaining_commands > 0 {
1340            warn!(
1341                "Dropping {} queued commands during cleanup",
1342                remaining_commands
1343            );
1344            self.command_queue.clear();
1345        }
1346
1347        // Cleanup WebSocket with timeout
1348        if let Err(e) = timeout(Duration::from_secs(5), self.ws.delete()).await {
1349            warn!("WebSocket cleanup timeout: {:?}", e);
1350        }
1351    }
1352
1353    // Public API methods
1354    pub fn shutdown_token(&self) -> CancellationToken {
1355        self.shutdown.clone()
1356    }
1357
1358    pub fn stats(&self) -> &ConnectionStats {
1359        &self.stats
1360    }
1361
1362    pub fn state(&self) -> &ConnectionState {
1363        &self.state
1364    }
1365
1366    pub fn config(&self) -> &CommandRunnerConfig {
1367        &self.config
1368    }
1369}
1370
1371#[cfg(test)]
1372mod tests {
1373    use super::*;
1374    use ustr::ustr;
1375
1376    #[test]
1377    fn test_replay_commands_validation() {
1378        let valid_start = Command::ReplayStart(ReplayStartCommandMsg {
1379            replay_session: ustr("rs_123"),
1380            request_id: ustr("req_1"),
1381            interval: 500,
1382        });
1383        assert!(valid_start.validate().is_ok());
1384
1385        let invalid_start = Command::ReplayStart(ReplayStartCommandMsg {
1386            replay_session: ustr(""),
1387            request_id: ustr("req_1"),
1388            interval: 500,
1389        });
1390        assert!(invalid_start.validate().is_err());
1391
1392        let invalid_start_id = Command::ReplayStart(ReplayStartCommandMsg {
1393            replay_session: ustr("rs_123"),
1394            request_id: ustr(""),
1395            interval: 500,
1396        });
1397        assert!(invalid_start_id.validate().is_err());
1398
1399        let valid_step = Command::ReplayStep(ReplayStepCommandMsg {
1400            replay_session: ustr("rs_123"),
1401            request_id: ustr("req_2"),
1402            step: 5,
1403        });
1404        assert!(valid_step.validate().is_ok());
1405
1406        let valid_stop = Command::ReplayStop(ReplayStopCommandMsg {
1407            replay_session: ustr("rs_123"),
1408            request_id: ustr("req_3"),
1409        });
1410        assert!(valid_stop.validate().is_ok());
1411
1412        let valid_reset = Command::ReplayReset(ReplayResetCommandMsg {
1413            replay_session: ustr("rs_123"),
1414            request_id: ustr("req_4"),
1415            timestamp: 1_700_000_000,
1416        });
1417        assert!(valid_reset.validate().is_ok());
1418    }
1419
1420    #[test]
1421    fn test_replay_commands_requires_session() {
1422        let cmd_start = Command::ReplayStart(ReplayStartCommandMsg {
1423            replay_session: ustr("rs_my_session"),
1424            request_id: ustr("req_1"),
1425            interval: 1000,
1426        });
1427        assert_eq!(cmd_start.requires_session(), Some("rs_my_session"));
1428
1429        let cmd_step = Command::ReplayStep(ReplayStepCommandMsg {
1430            replay_session: ustr("rs_my_session"),
1431            request_id: ustr("req_2"),
1432            step: 1,
1433        });
1434        assert_eq!(cmd_step.requires_session(), Some("rs_my_session"));
1435
1436        let cmd_stop = Command::ReplayStop(ReplayStopCommandMsg {
1437            replay_session: ustr("rs_my_session"),
1438            request_id: ustr("req_3"),
1439        });
1440        assert_eq!(cmd_stop.requires_session(), Some("rs_my_session"));
1441
1442        let cmd_reset = Command::ReplayReset(ReplayResetCommandMsg {
1443            replay_session: ustr("rs_my_session"),
1444            request_id: ustr("req_4"),
1445            timestamp: 1_700_000_000,
1446        });
1447        assert_eq!(cmd_reset.requires_session(), Some("rs_my_session"));
1448    }
1449}
tradingview/live/handler/command.rs

tradingview/live/handler/
command.rs
