1#[cfg(test)]
2mod agent_test;
3
4use shared::error::*;
5use std::collections::{HashMap, VecDeque};
6use std::time::Instant;
7
8use crate::message::*;
9
10#[derive(Default)]
13pub struct Agent {
14 transactions: HashMap<TransactionId, AgentTransaction>,
20 closed: bool,
22 events_queue: VecDeque<Event>,
24}
25
26#[derive(Debug)] pub struct Event {
30 pub id: TransactionId,
31 pub evt: StunEvent,
32}
33
34#[derive(Debug)] pub enum StunEvent {
36 AgentClosed,
37 TransactionStopped,
38 TransactionTimeOut,
39 Message(Message),
40}
41
42pub(crate) struct AgentTransaction {
45 id: TransactionId,
46 deadline: Instant,
47}
48
49const AGENT_COLLECT_CAP: usize = 100;
52
53#[derive(Debug)]
56pub enum ClientAgent {
57 Process(Message),
58 Collect(Instant),
59 Start(TransactionId, Instant),
60 Stop(TransactionId),
61 Close,
62}
63
64impl Agent {
65 pub fn new() -> Self {
67 Agent {
68 transactions: HashMap::new(),
69 closed: false,
70 events_queue: VecDeque::new(),
71 }
72 }
73
74 pub fn handle_event(&mut self, client_agent: ClientAgent) -> Result<()> {
75 match client_agent {
76 ClientAgent::Process(message) => self.process(message),
77 ClientAgent::Collect(deadline) => self.collect(deadline),
78 ClientAgent::Start(tid, deadline) => self.start(tid, deadline),
79 ClientAgent::Stop(tid) => self.stop(tid),
80 ClientAgent::Close => self.close(),
81 }
82 }
83
84 pub fn poll_timeout(&mut self) -> Option<Instant> {
85 let mut deadline = None;
86 for transaction in self.transactions.values() {
87 if deadline.is_none() || transaction.deadline < *deadline.as_ref().unwrap() {
88 deadline = Some(transaction.deadline);
89 }
90 }
91 deadline
92 }
93
94 pub fn poll_event(&mut self) -> Option<Event> {
95 self.events_queue.pop_front()
96 }
97
98 fn process(&mut self, message: Message) -> Result<()> {
100 if self.closed {
101 return Err(Error::ErrAgentClosed);
102 }
103
104 self.transactions.remove(&message.transaction_id);
105
106 self.events_queue.push_back(Event {
107 id: message.transaction_id,
108 evt: StunEvent::Message(message),
109 });
110
111 Ok(())
112 }
113
114 fn close(&mut self) -> Result<()> {
117 if self.closed {
118 return Err(Error::ErrAgentClosed);
119 }
120
121 for id in self.transactions.keys() {
122 self.events_queue.push_back(Event {
123 id: *id,
124 evt: StunEvent::AgentClosed,
125 });
126 }
127 self.transactions.clear();
128 self.closed = true;
129
130 Ok(())
131 }
132
133 fn start(&mut self, id: TransactionId, deadline: Instant) -> Result<()> {
138 if self.closed {
139 return Err(Error::ErrAgentClosed);
140 }
141 if self.transactions.contains_key(&id) {
142 return Err(Error::ErrTransactionExists);
143 }
144
145 self.transactions
146 .insert(id, AgentTransaction { id, deadline });
147
148 Ok(())
149 }
150
151 fn stop(&mut self, id: TransactionId) -> Result<()> {
154 if self.closed {
155 return Err(Error::ErrAgentClosed);
156 }
157
158 let v = self.transactions.remove(&id);
159 if let Some(t) = v {
160 self.events_queue.push_back(Event {
161 id: t.id,
162 evt: StunEvent::TransactionStopped,
163 });
164 Ok(())
165 } else {
166 Err(Error::ErrTransactionNotExists)
167 }
168 }
169
170 fn collect(&mut self, deadline: Instant) -> Result<()> {
176 if self.closed {
177 return Err(Error::ErrAgentClosed);
181 }
182
183 let mut to_remove: Vec<TransactionId> = Vec::with_capacity(AGENT_COLLECT_CAP);
184
185 for (id, t) in &self.transactions {
190 if t.deadline < deadline {
191 to_remove.push(*id);
192 }
193 }
194 for id in &to_remove {
196 self.transactions.remove(id);
197 }
198
199 for id in to_remove {
200 self.events_queue.push_back(Event {
201 id,
202 evt: StunEvent::TransactionTimeOut,
203 });
204 }
205
206 Ok(())
207 }
208}