1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318
use std::collections::HashMap; /// Payload for a publisher confirmation message (either an [ack](enum.Confirm.html#variant.Ack) or /// a [nack](enum.Confirm.html#variant.Nack)) from the server. #[derive(Debug, Clone, Copy, PartialEq)] pub struct ConfirmPayload { /// The tag from the server. Tags are sequentially increasing integers beginning with /// 1 (once publisher confirms [are /// enabled](struct.Channel.html#method.enable_publisher_confirms) on the channel. pub delivery_tag: u64, /// If true, the confirmation applies to all previously-unconfirmed messages with delivery tags /// less than or equal to this payload's [`delivery_tag`](#structfield.delivery_tag). pub multiple: bool, } /// A publisher confirmation message from the server. #[derive(Debug, Clone, Copy, PartialEq)] pub enum Confirm { /// Acknowledgment that the server has received the message(s) described by the associated /// payload. Note that acks do not necessarily imply that the messages have been handled by a /// consumer, merely that they have been received by the server. Ack(ConfirmPayload), /// Notification that the message(s) described by the associated payload have been rejected. Nack(ConfirmPayload), } /// Helper to smooth out of order and/or `multiple: true` publisher confirmation messages. /// /// If publisher confirms are enabled, the server may confirm messages out of order and/or may /// confirm multiple messages with a single [`Confirm`](enum.Confirm.html). `ConfirmSmoother` /// exists to "smooth" server messages out into an always-increasing-by-one sequence of /// confirmation messages. /// /// # Example /// /// ```rust /// use amiquip::{Confirm, ConfirmSmoother}; /// use crossbeam_channel::Receiver; /// /// // assume we've published n messages and want to wait for them to be confirmed /// fn wait_for_publisher_confirms(n: usize, receiver: &Receiver<Confirm>) { /// // NOTE: a new smoother assumes we will be receiving messages starting with /// // delivery_tag = 1, so this method is only valid if called after the first n /// // publishes on a channel. We could take a &mut ConfirmSmoother to be called /// // multiple times in succession on the same channel. /// let mut smoother = ConfirmSmoother::new(); /// let mut acked = 0; /// let mut nacked = 0; /// while acked + nacked < n { /// // get a confirmation from the server; this may be out of order or a confirm /// // for multiple messages in one payload /// let raw_confirm = match receiver.recv() { /// Ok(raw_confirm) => raw_confirm, /// Err(_) => { /// // the I/O thread dropped the sending side; either an error has occurred /// // or another thread of ours closed the connection; either way we'll /// // stop waiting /// return; /// } /// }; /// /// // feed the raw confirm into the smoother. Two notes: /// // 1. We must run the returned iterator to the end or risk missing confirms. /// // 2. The iterator may produce 0, 1, or multiple independent confirmations. /// // They will all have multiple: false. /// for confirm in smoother.process(raw_confirm) { /// match confirm { /// Confirm::Ack(_) => { /// acked += 1; /// } /// Confirm::Nack(_) => { /// // server rejected message; need to do something else to /// // track which messages were rejected /// nacked += 1; /// } /// } /// } /// } /// } /// ``` #[derive(Debug, Clone)] pub struct ConfirmSmoother { expected: u64, out_of_order: HashMap<u64, Confirm>, } impl Default for ConfirmSmoother { fn default() -> ConfirmSmoother { ConfirmSmoother::new() } } impl ConfirmSmoother { /// Create a new `ConfirmSmoother`. It expects the next (in absolute order) delivery tag /// received from the server to be `1`. pub fn new() -> ConfirmSmoother { ConfirmSmoother::with_expected_delivery_tag(1) } /// Create a new `ConfirmSmoother`. It expects the next (in absolute order) delivery tag /// received from the server to be `expected`. pub fn with_expected_delivery_tag(expected: u64) -> ConfirmSmoother { ConfirmSmoother { expected, out_of_order: HashMap::new(), } } /// Process a confirmation message from the server. Returns an iterator; each item returned by /// the iterator will be a single (i.e., `multiple: false`) [`Confirm`](enum.Confirm.html). You /// _must_ run the iterator to its completion or risk missing confirmations; future calls to /// `process` will not return an iterator that will repeat confirms that would have been /// returned by a previously returned iterator (even if that earlier iterator was dropped /// before it ran to completion). /// /// The returned iterator may have 0 items (if `confirm` is a non-multiple confirmation that is /// later than the next expected delivery tag), 1 item (if `confirm` exactly matches our next /// expected delivery tag and we had not previously seen the next tag), or multiple items (if /// `confirm` is a `multiple: true` confirmation or we've previously seen out-of-order tags /// that are next sequentially after `confirm`'s tag). pub fn process<'a>(&'a mut self, confirm: Confirm) -> impl Iterator<Item = Confirm> + 'a { match confirm { Confirm::Ack(inner) => self.new_iter(inner, Confirm::Ack), Confirm::Nack(inner) => self.new_iter(inner, Confirm::Nack), } } fn new_iter<'a>( &'a mut self, payload: ConfirmPayload, to_confirm: fn(ConfirmPayload) -> Confirm, ) -> impl Iterator<Item = Confirm> + 'a { Iter { parent: self, payload, next: None, to_confirm: move |tag| { to_confirm(ConfirmPayload { delivery_tag: tag, multiple: false, }) }, done: false, } } } struct Iter<'a, F: Fn(u64) -> Confirm> { parent: &'a mut ConfirmSmoother, payload: ConfirmPayload, next: Option<Confirm>, to_confirm: F, done: bool, } impl<'a, F> Drop for Iter<'a, F> where F: Fn(u64) -> Confirm, { fn drop(&mut self) { while !self.done { let _ = self.next(); } } } impl<'a, F> Iterator for Iter<'a, F> where F: Fn(u64) -> Confirm, { type Item = Confirm; fn next(&mut self) -> Option<Confirm> { if self.done { return None; } let payload = self.payload; if payload.delivery_tag == self.parent.expected { // exact match - we'll return this tag, and set next to an out-of-order // entry for the next tag if we had one self.parent.expected += 1; self.next = self.parent.out_of_order.remove(&self.parent.expected); return Some((self.to_confirm)(payload.delivery_tag)); } if payload.delivery_tag > self.parent.expected { // tag is in the future; if it's "multiple", keep sending all tags in // between where we are now and payload.delivery_tag if payload.multiple { let ret = (self.to_confirm)(self.parent.expected); self.parent.expected += 1; return Some(ret); } else { // if it's _not_ multiple, stash it away in out_of_order self.parent.out_of_order.insert( payload.delivery_tag, (self.to_confirm)(payload.delivery_tag), ); self.done = true; return None; } } match self.next.take() { Some(next) => { // self.next is Some() only if a previous call to next() hit the tag==expected // case _and_ out_of_order held a confirm for the next expected tag self.parent.expected += 1; self.next = self.parent.out_of_order.remove(&self.parent.expected); Some(next) } None => { self.done = true; None } } } } #[cfg(test)] mod tests { use super::*; fn single(delivery_tag: u64, f: fn(ConfirmPayload) -> Confirm) -> Confirm { f(ConfirmPayload { delivery_tag, multiple: false, }) } fn multiple(delivery_tag: u64, f: fn(ConfirmPayload) -> Confirm) -> Confirm { f(ConfirmPayload { delivery_tag, multiple: true, }) } #[test] fn simple_single() { let mut flat = ConfirmSmoother::new(); let one = flat.process(single(1, Confirm::Ack)); let expected = vec![single(1, Confirm::Ack)]; assert_eq!(expected, one.collect::<Vec<_>>()); } #[test] fn simple_multiple() { let mut flat = ConfirmSmoother::new(); let three = flat.process(multiple(3, Confirm::Ack)); let expected = (1..=3).map(|i| single(i, Confirm::Ack)).collect::<Vec<_>>(); assert_eq!(expected, three.collect::<Vec<_>>()); } #[test] fn single_then_single() { let mut flat = ConfirmSmoother::new(); let empty = flat.process(single(3, Confirm::Ack)); assert_eq!(empty.count(), 0); let empty = flat.process(single(2, Confirm::Ack)); assert_eq!(empty.count(), 0); let two = flat.process(single(1, Confirm::Nack)); let expected = vec![ single(1, Confirm::Nack), single(2, Confirm::Ack), single(3, Confirm::Ack), ]; assert_eq!(expected, two.collect::<Vec<_>>()); } #[test] fn redelivery() { let mut flat = ConfirmSmoother::new(); let two = flat.process(multiple(2, Confirm::Ack)); assert_eq!(two.count(), 2); // getting another confirm for 1 or 2 should do nothing since flattener already // dispatched a confirm for 1 and 2 let empty = flat.process(single(2, Confirm::Ack)); assert_eq!(empty.count(), 0); let empty = flat.process(single(1, Confirm::Nack)); assert_eq!(empty.count(), 0); } #[test] fn single_single_multiple_single() { let mut flat = ConfirmSmoother::new(); let empty = flat.process(single(5, Confirm::Nack)); assert_eq!(empty.count(), 0); let empty = flat.process(single(3, Confirm::Nack)); assert_eq!(empty.count(), 0); let three = flat.process(multiple(2, Confirm::Ack)); assert_eq!( vec![ single(1, Confirm::Ack), single(2, Confirm::Ack), single(3, Confirm::Nack) ], three.collect::<Vec<_>>() ); let two = flat.process(single(4, Confirm::Ack)); assert_eq!( vec![single(4, Confirm::Ack), single(5, Confirm::Nack),], two.collect::<Vec<_>>() ); } #[test] fn drop_without_running_iter_to_completion() { let mut flat = ConfirmSmoother::new(); let _ = flat.process(multiple(2, Confirm::Ack)); let one = flat.process(single(3, Confirm::Ack)); let expected = vec![single(3, Confirm::Ack)]; assert_eq!(expected, one.collect::<Vec<_>>()); } }