1use std::cmp;
28
29use std::collections::VecDeque;
30
31use crate::buffers::BufSplit;
32use crate::range_buf::RangeBuf;
33use crate::BufFactory;
34use crate::Error;
35use crate::Result;
36
37use crate::buffers::DefaultBufFactory;
38use crate::ranges;
39
40#[cfg(test)]
41const SEND_BUFFER_SIZE: usize = 5;
42
43#[cfg(not(test))]
44const SEND_BUFFER_SIZE: usize = 4096;
45
46struct SendReserve<'a, F: BufFactory> {
47 inner: &'a mut SendBuf<F>,
48 reserved: usize,
49 fin: bool,
50}
51
52impl<F: BufFactory> SendReserve<'_, F> {
53 fn append_buf(&mut self, buf: F::Buf) -> Result<()> {
54 let len = buf.as_ref().len();
55 let inner = &mut self.inner;
56
57 if len > self.reserved {
58 return Err(Error::BufferTooShort);
59 }
60
61 let fin: bool = self.reserved == len && self.fin;
62
63 let buf = RangeBuf::from_raw(buf, inner.off, fin);
64
65 inner.data.push_back(buf);
67
68 inner.off += len as u64;
69 inner.buffered_bytes += len as u64;
70 self.reserved -= len;
71
72 Ok(())
73 }
74}
75
76impl<F: BufFactory> Drop for SendReserve<'_, F> {
77 fn drop(&mut self) {
78 assert_eq!(self.reserved, 0)
79 }
80}
81
82#[derive(Debug, Default)]
92pub struct SendBuf<F = DefaultBufFactory>
93where
94 F: BufFactory,
95{
96 data: VecDeque<RangeBuf<F>>,
98
99 pos: usize,
101
102 off: u64,
104
105 emit_off: u64,
108
109 buffered_bytes: u64,
115
116 max_data: u64,
118
119 blocked_at: Option<u64>,
121
122 fin_off: Option<u64>,
124
125 shutdown: bool,
127
128 acked: ranges::RangeSet,
130
131 error: Option<u64>,
133}
134
135impl<F: BufFactory> SendBuf<F> {
136 pub fn new(max_data: u64) -> SendBuf<F> {
138 SendBuf {
139 max_data,
140 ..SendBuf::default()
141 }
142 }
143
144 fn reserve_for_write(
146 &mut self, mut len: usize, mut fin: bool,
147 ) -> Result<SendReserve<'_, F>> {
148 let max_off = self.off + len as u64;
149
150 if len > self.cap()? {
153 len = self.cap()?;
154 fin = false;
155 }
156
157 if let Some(fin_off) = self.fin_off {
158 if max_off > fin_off {
160 return Err(Error::FinalSize);
161 }
162
163 if max_off == fin_off && !fin {
165 return Err(Error::FinalSize);
166 }
167 }
168
169 if fin {
170 self.fin_off = Some(max_off);
171 }
172
173 if self.ack_off() >= max_off {
175 return Ok(SendReserve {
176 inner: self,
177 reserved: 0,
178 fin,
179 });
180 }
181
182 Ok(SendReserve {
183 inner: self,
184 reserved: len,
185 fin,
186 })
187 }
188
189 pub fn write(&mut self, data: &[u8], fin: bool) -> Result<usize> {
195 let mut reserve = self.reserve_for_write(data.len(), fin)?;
196
197 if reserve.reserved == 0 {
198 return Ok(0);
199 }
200
201 let ret = reserve.reserved;
202
203 for chunk in data[..reserve.reserved].chunks(SEND_BUFFER_SIZE) {
206 reserve.append_buf(F::buf_from_slice(chunk))?;
207 }
208
209 Ok(ret)
210 }
211
212 pub fn append_buf(
218 &mut self, mut data: F::Buf, cap: usize, fin: bool,
219 ) -> Result<(usize, Option<F::Buf>)>
220 where
221 F::Buf: BufSplit,
222 {
223 let len = data.as_ref().len();
224 let mut reserve = self.reserve_for_write(cap.min(len), fin)?;
225
226 if reserve.reserved == 0 {
227 return Ok((0, Some(data)));
228 }
229
230 let remainder =
231 (reserve.reserved < len).then(|| data.split_at(reserve.reserved));
232
233 let ret = reserve.reserved;
234
235 reserve.append_buf(data)?;
236
237 Ok((ret, remainder))
238 }
239
240 pub fn emit(&mut self, out: &mut [u8]) -> Result<(usize, bool)> {
242 let mut out_len = out.len();
243 let out_off = self.off_front();
244
245 let mut next_off = out_off;
246
247 while out_len > 0 {
248 let off_front = self.off_front();
249
250 if self.is_empty() ||
251 off_front >= self.off ||
252 off_front != next_off ||
253 off_front >= self.max_data
254 {
255 break;
256 }
257
258 let buf = match self.data.get_mut(self.pos) {
259 Some(v) => v,
260
261 None => break,
262 };
263
264 if buf.is_empty() {
265 self.pos += 1;
266 continue;
267 }
268
269 let buf_len = cmp::min(buf.len(), out_len);
270 let partial = buf_len < buf.len();
271
272 let out_pos = (next_off - out_off) as usize;
274 out[out_pos..out_pos + buf_len].copy_from_slice(&buf[..buf_len]);
275
276 self.buffered_bytes -= buf_len as u64;
277
278 out_len -= buf_len;
279
280 next_off = buf.off() + buf_len as u64;
281
282 buf.consume(buf_len);
283
284 if partial {
285 break;
287 }
288
289 self.pos += 1;
290 }
291
292 let fin = self.fin_off == Some(next_off);
299
300 self.emit_off = cmp::max(self.emit_off, next_off);
303
304 Ok((out.len() - out_len, fin))
305 }
306
307 pub fn update_max_data(&mut self, max_data: u64) {
309 self.max_data = cmp::max(self.max_data, max_data);
310 }
311
312 pub fn update_blocked_at(&mut self, blocked_at: Option<u64>) {
314 self.blocked_at = blocked_at;
315 }
316
317 pub fn blocked_at(&self) -> Option<u64> {
319 self.blocked_at
320 }
321
322 pub fn ack(&mut self, off: u64, len: usize) {
324 self.acked.insert(off..off + len as u64);
325 }
326
327 pub fn ack_and_drop(&mut self, off: u64, len: usize) -> usize {
328 self.ack(off, len);
329
330 let ack_off = self.ack_off();
331
332 if self.data.is_empty() {
333 return 0;
334 }
335
336 if off > ack_off {
337 return 0;
338 }
339
340 let mut drop_until = None;
341
342 for (i, buf) in self.data.iter_mut().enumerate() {
344 if buf.off >= ack_off {
347 break;
348 }
349
350 if buf.off < ack_off && ack_off < buf.max_off() {
353 break;
354 }
355
356 drop_until = Some(i);
358 }
359
360 if let Some(drop) = drop_until {
361 let dropped_len: u64 =
364 (0..=drop).map(|i| self.data[i].len() as u64).sum();
365 self.buffered_bytes = self.buffered_bytes.saturating_sub(dropped_len);
366
367 self.data.drain(..=drop);
368
369 self.pos = self.pos.saturating_sub(drop + 1);
373
374 dropped_len as usize
375 } else {
376 0
377 }
378 }
379
380 pub fn retransmit(&mut self, off: u64, len: usize) -> usize {
381 let max_off = off + len as u64;
382 let ack_off = self.ack_off();
383
384 if self.data.is_empty() {
385 return 0;
386 }
387
388 if max_off <= ack_off {
389 return 0;
390 }
391
392 let mut total_retransmitted = 0;
393
394 for i in 0..self.data.len() {
395 let buf = &mut self.data[i];
396
397 if buf.off >= max_off {
398 break;
399 }
400
401 if off > buf.max_off() {
402 continue;
403 }
404
405 let new_buf = if buf.off < max_off && max_off < buf.max_off() {
408 Some(buf.split_off((max_off - buf.off) as usize))
409 } else {
410 None
411 };
412
413 let prev_pos = buf.pos;
414
415 buf.pos = if off > buf.off && off <= buf.max_off() {
418 cmp::min(buf.pos, buf.start + (off - buf.off) as usize)
419 } else {
420 buf.start
421 };
422
423 self.pos = cmp::min(self.pos, i);
424
425 let retransmitted = (prev_pos - buf.pos) as u64;
426 self.buffered_bytes += retransmitted;
427 total_retransmitted += retransmitted;
428
429 if let Some(b) = new_buf {
430 self.data.insert(i + 1, b);
431 }
432 }
433
434 total_retransmitted as usize
435 }
436
437 pub fn reset(&mut self) -> (u64, u64) {
439 let unsent_off = cmp::max(self.off_front(), self.emit_off);
440 let unsent_len = self.off_back().saturating_sub(unsent_off);
441
442 self.fin_off = Some(unsent_off);
443
444 self.data.clear();
446
447 self.off = unsent_off;
449 self.ack(0, self.off as usize);
450
451 self.pos = 0;
452 self.buffered_bytes = 0;
453
454 (self.emit_off, unsent_len)
455 }
456
457 pub fn stop(&mut self, error_code: u64) -> Result<(u64, u64)> {
461 if self.error.is_some() {
462 return Err(Error::Done);
463 }
464
465 let (max_off, unsent) = self.reset();
466
467 self.error = Some(error_code);
468
469 Ok((max_off, unsent))
470 }
471
472 pub fn shutdown(&mut self) -> Result<(u64, u64)> {
474 if self.shutdown {
475 return Err(Error::Done);
476 }
477
478 self.shutdown = true;
479
480 Ok(self.reset())
481 }
482
483 pub fn off_back(&self) -> u64 {
485 self.off
486 }
487
488 pub fn off_front(&self) -> u64 {
490 let mut pos = self.pos;
491
492 while let Some(b) = self.data.get(pos) {
494 if !b.is_empty() {
495 return b.off();
496 }
497
498 pos += 1;
499 }
500
501 self.off
502 }
503
504 pub fn max_off(&self) -> u64 {
506 self.max_data
507 }
508
509 pub fn is_fin(&self) -> bool {
514 if self.fin_off == Some(self.off) {
515 return true;
516 }
517
518 false
519 }
520
521 pub fn is_complete(&self) -> bool {
526 if let Some(fin_off) = self.fin_off {
527 if self.acked == (0..fin_off) {
528 return true;
529 }
530 }
531
532 false
533 }
534
535 pub fn is_stopped(&self) -> bool {
537 self.error.is_some()
538 }
539
540 pub fn is_shutdown(&self) -> bool {
542 self.shutdown
543 }
544
545 pub fn is_empty(&self) -> bool {
547 self.data.is_empty()
548 }
549
550 pub fn ack_off(&self) -> u64 {
552 match self.acked.iter().next() {
553 Some(std::ops::Range { start: 0, end }) => end,
556
557 Some(_) | None => 0,
558 }
559 }
560
561 pub fn cap(&self) -> Result<usize> {
563 if let Some(e) = self.error {
565 return Err(Error::StreamStopped(e));
566 }
567
568 Ok((self.max_data - self.off) as usize)
569 }
570
571 #[allow(dead_code)]
573 pub fn bufs_count(&self) -> usize {
574 self.data.len()
575 }
576
577 pub fn buffered_bytes(&self) -> u64 {
583 self.buffered_bytes
584 }
585}
586
587#[cfg(test)]
588mod tests {
589 use super::*;
590
591 #[test]
592 fn empty_write() {
593 let mut buf = [0; 5];
594
595 let mut send = <SendBuf>::new(u64::MAX);
596 assert_eq!(send.buffered_bytes, 0);
597
598 let (written, fin) = send.emit(&mut buf).unwrap();
599 assert_eq!(written, 0);
600 assert!(!fin);
601 }
602
603 #[test]
604 fn multi_write() {
605 let mut buf = [0; 128];
606
607 let mut send = <SendBuf>::new(u64::MAX);
608 assert_eq!(send.buffered_bytes, 0);
609
610 let first = b"something";
611 let second = b"helloworld";
612
613 assert!(send.write(first, false).is_ok());
614 assert_eq!(send.buffered_bytes, 9);
615
616 assert!(send.write(second, true).is_ok());
617 assert_eq!(send.buffered_bytes, 19);
618
619 let (written, fin) = send.emit(&mut buf[..128]).unwrap();
620 assert_eq!(written, 19);
621 assert!(fin);
622 assert_eq!(&buf[..written], b"somethinghelloworld");
623 assert_eq!(send.buffered_bytes, 0);
624 }
625
626 #[test]
627 fn split_write() {
628 let mut buf = [0; 10];
629
630 let mut send = <SendBuf>::new(u64::MAX);
631 assert_eq!(send.buffered_bytes, 0);
632
633 let first = b"something";
634 let second = b"helloworld";
635
636 assert!(send.write(first, false).is_ok());
637 assert_eq!(send.buffered_bytes, 9);
638
639 assert!(send.write(second, true).is_ok());
640 assert_eq!(send.buffered_bytes, 19);
641
642 assert_eq!(send.off_front(), 0);
643
644 let (written, fin) = send.emit(&mut buf[..10]).unwrap();
645 assert_eq!(written, 10);
646 assert!(!fin);
647 assert_eq!(&buf[..written], b"somethingh");
648 assert_eq!(send.buffered_bytes, 9);
649
650 assert_eq!(send.off_front(), 10);
651
652 let (written, fin) = send.emit(&mut buf[..5]).unwrap();
653 assert_eq!(written, 5);
654 assert!(!fin);
655 assert_eq!(&buf[..written], b"ellow");
656 assert_eq!(send.buffered_bytes, 4);
657
658 assert_eq!(send.off_front(), 15);
659
660 let (written, fin) = send.emit(&mut buf[..10]).unwrap();
661 assert_eq!(written, 4);
662 assert!(fin);
663 assert_eq!(&buf[..written], b"orld");
664 assert_eq!(send.buffered_bytes, 0);
665
666 assert_eq!(send.off_front(), 19);
667 }
668
669 #[test]
670 fn resend() {
671 let mut buf = [0; 15];
672
673 let mut send = <SendBuf>::new(u64::MAX);
674 assert_eq!(send.buffered_bytes, 0);
675 assert_eq!(send.off_front(), 0);
676
677 let first = b"something";
678 let second = b"helloworld";
679
680 assert!(send.write(first, false).is_ok());
681 assert_eq!(send.off_front(), 0);
682
683 assert!(send.write(second, true).is_ok());
684 assert_eq!(send.off_front(), 0);
685
686 assert_eq!(send.buffered_bytes, 19);
687
688 let (written, fin) = send.emit(&mut buf[..4]).unwrap();
689 assert_eq!(written, 4);
690 assert!(!fin);
691 assert_eq!(&buf[..written], b"some");
692 assert_eq!(send.buffered_bytes, 15);
693 assert_eq!(send.off_front(), 4);
694
695 let (written, fin) = send.emit(&mut buf[..5]).unwrap();
696 assert_eq!(written, 5);
697 assert!(!fin);
698 assert_eq!(&buf[..written], b"thing");
699 assert_eq!(send.buffered_bytes, 10);
700 assert_eq!(send.off_front(), 9);
701
702 let (written, fin) = send.emit(&mut buf[..5]).unwrap();
703 assert_eq!(written, 5);
704 assert!(!fin);
705 assert_eq!(&buf[..written], b"hello");
706 assert_eq!(send.buffered_bytes, 5);
707 assert_eq!(send.off_front(), 14);
708
709 send.retransmit(4, 5);
710 assert_eq!(send.buffered_bytes, 10);
711 assert_eq!(send.off_front(), 4);
712
713 send.retransmit(0, 4);
714 assert_eq!(send.buffered_bytes, 14);
715 assert_eq!(send.off_front(), 0);
716
717 let (written, fin) = send.emit(&mut buf[..11]).unwrap();
718 assert_eq!(written, 9);
719 assert!(!fin);
720 assert_eq!(&buf[..written], b"something");
721 assert_eq!(send.buffered_bytes, 5);
722 assert_eq!(send.off_front(), 14);
723
724 let (written, fin) = send.emit(&mut buf[..11]).unwrap();
725 assert_eq!(written, 5);
726 assert!(fin);
727 assert_eq!(&buf[..written], b"world");
728 assert_eq!(send.buffered_bytes, 0);
729 assert_eq!(send.off_front(), 19);
730 }
731
732 #[test]
733 fn write_blocked_by_off() {
734 let mut buf = [0; 10];
735
736 let mut send = <SendBuf>::default();
737 assert_eq!(send.buffered_bytes, 0);
738
739 let first = b"something";
740 let second = b"helloworld";
741
742 assert_eq!(send.write(first, false), Ok(0));
743 assert_eq!(send.buffered_bytes, 0);
744
745 assert_eq!(send.write(second, true), Ok(0));
746 assert_eq!(send.buffered_bytes, 0);
747
748 send.update_max_data(5);
749
750 assert_eq!(send.write(first, false), Ok(5));
751 assert_eq!(send.buffered_bytes, 5);
752
753 assert_eq!(send.write(second, true), Ok(0));
754 assert_eq!(send.buffered_bytes, 5);
755
756 assert_eq!(send.off_front(), 0);
757
758 let (written, fin) = send.emit(&mut buf[..10]).unwrap();
759 assert_eq!(written, 5);
760 assert!(!fin);
761 assert_eq!(&buf[..written], b"somet");
762 assert_eq!(send.buffered_bytes, 0);
763
764 assert_eq!(send.off_front(), 5);
765
766 let (written, fin) = send.emit(&mut buf[..10]).unwrap();
767 assert_eq!(written, 0);
768 assert!(!fin);
769 assert_eq!(&buf[..written], b"");
770 assert_eq!(send.buffered_bytes, 0);
771
772 send.update_max_data(15);
773
774 assert_eq!(send.write(&first[5..], false), Ok(4));
775 assert_eq!(send.buffered_bytes, 4);
776
777 assert_eq!(send.write(second, true), Ok(6));
778 assert_eq!(send.buffered_bytes, 10);
779
780 assert_eq!(send.off_front(), 5);
781
782 let (written, fin) = send.emit(&mut buf[..10]).unwrap();
783 assert_eq!(written, 10);
784 assert!(!fin);
785 assert_eq!(&buf[..10], b"hinghellow");
786 assert_eq!(send.buffered_bytes, 0);
787
788 send.update_max_data(25);
789
790 assert_eq!(send.write(&second[6..], true), Ok(4));
791 assert_eq!(send.buffered_bytes, 4);
792
793 assert_eq!(send.off_front(), 15);
794
795 let (written, fin) = send.emit(&mut buf[..10]).unwrap();
796 assert_eq!(written, 4);
797 assert!(fin);
798 assert_eq!(&buf[..written], b"orld");
799 assert_eq!(send.buffered_bytes, 0);
800 }
801
802 #[test]
803 fn zero_len_write() {
804 let mut buf = [0; 10];
805
806 let mut send = <SendBuf>::new(u64::MAX);
807 assert_eq!(send.buffered_bytes, 0);
808
809 let first = b"something";
810
811 assert!(send.write(first, false).is_ok());
812 assert_eq!(send.buffered_bytes, 9);
813
814 assert!(send.write(&[], true).is_ok());
815 assert_eq!(send.buffered_bytes, 9);
816
817 assert_eq!(send.off_front(), 0);
818
819 let (written, fin) = send.emit(&mut buf[..10]).unwrap();
820 assert_eq!(written, 9);
821 assert!(fin);
822 assert_eq!(&buf[..written], b"something");
823 assert_eq!(send.buffered_bytes, 0);
824 }
825
826 #[test]
828 fn send_buf_len_on_retransmit() {
829 let mut buf = [0; 15];
830
831 let mut send = <SendBuf>::new(u64::MAX);
832 assert_eq!(send.buffered_bytes, 0);
833 assert_eq!(send.off_front(), 0);
834
835 let first = b"something";
836
837 assert!(send.write(first, false).is_ok());
838 assert_eq!(send.off_front(), 0);
839
840 assert_eq!(send.buffered_bytes, 9);
841
842 let (written, fin) = send.emit(&mut buf[..4]).unwrap();
843 assert_eq!(written, 4);
844 assert!(!fin);
845 assert_eq!(&buf[..written], b"some");
846 assert_eq!(send.buffered_bytes, 5);
847 assert_eq!(send.off_front(), 4);
848
849 send.retransmit(3, 5);
850 assert_eq!(send.buffered_bytes, 6);
851 assert_eq!(send.off_front(), 3);
852 }
853
854 #[test]
855 fn send_buf_final_size_retransmit() {
856 let mut buf = [0; 50];
857 let mut send = <SendBuf>::new(u64::MAX);
858
859 send.write(&buf, false).unwrap();
860 assert_eq!(send.off_front(), 0);
861
862 let (written, _fin) = send.emit(&mut buf).unwrap();
864 assert_eq!(written, buf.len());
865 assert_eq!(send.off_front(), buf.len() as u64);
866
867 send.retransmit(40, 10);
870
871 let (fin_off, unsent) = send.stop(0).unwrap();
875 assert_eq!(fin_off, 50);
876 assert_eq!(unsent, 0);
877 }
878}