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 fin_acked: bool,
127
128 shutdown: bool,
130
131 acked: ranges::RangeSet,
133
134 error: Option<u64>,
136
137 error_reported: bool,
139}
140
141impl<F: BufFactory> SendBuf<F> {
142 pub fn new(max_data: u64) -> SendBuf<F> {
144 SendBuf {
145 max_data,
146 ..SendBuf::default()
147 }
148 }
149
150 fn reserve_for_write(
152 &mut self, mut len: usize, mut fin: bool,
153 ) -> Result<SendReserve<'_, F>> {
154 let max_off = self.off + len as u64;
155
156 if len > self.cap()? {
159 len = self.cap()?;
160 fin = false;
161 }
162
163 if let Some(fin_off) = self.fin_off {
164 if max_off > fin_off {
166 return Err(Error::FinalSize);
167 }
168
169 if max_off == fin_off && !fin {
171 return Err(Error::FinalSize);
172 }
173 }
174
175 if fin {
176 self.fin_off = Some(max_off);
177 }
178
179 if self.ack_off() >= max_off {
181 return Ok(SendReserve {
182 inner: self,
183 reserved: 0,
184 fin,
185 });
186 }
187
188 Ok(SendReserve {
189 inner: self,
190 reserved: len,
191 fin,
192 })
193 }
194
195 pub fn write(&mut self, data: &[u8], fin: bool) -> Result<usize> {
201 let mut reserve = self.reserve_for_write(data.len(), fin)?;
202
203 if reserve.reserved == 0 {
204 return Ok(0);
205 }
206
207 let ret = reserve.reserved;
208
209 for chunk in data[..reserve.reserved].chunks(SEND_BUFFER_SIZE) {
212 reserve.append_buf(F::buf_from_slice(chunk))?;
213 }
214
215 Ok(ret)
216 }
217
218 pub fn append_buf(
224 &mut self, mut data: F::Buf, cap: usize, fin: bool,
225 ) -> Result<(usize, Option<F::Buf>)>
226 where
227 F::Buf: BufSplit,
228 {
229 let len = data.as_ref().len();
230 let mut reserve = self.reserve_for_write(cap.min(len), fin)?;
231
232 if reserve.reserved == 0 {
233 return Ok((0, Some(data)));
234 }
235
236 let remainder =
237 (reserve.reserved < len).then(|| data.split_at(reserve.reserved));
238
239 let ret = reserve.reserved;
240
241 reserve.append_buf(data)?;
242
243 Ok((ret, remainder))
244 }
245
246 pub fn emit(&mut self, out: &mut [u8]) -> Result<(usize, bool)> {
248 let mut out_len = out.len();
249 let out_off = self.off_front();
250
251 let mut next_off = out_off;
252
253 while out_len > 0 {
254 let off_front = self.off_front();
255
256 if self.is_empty() ||
257 off_front >= self.off ||
258 off_front != next_off ||
259 off_front >= self.max_data
260 {
261 break;
262 }
263
264 let buf = match self.data.get_mut(self.pos) {
265 Some(v) => v,
266
267 None => break,
268 };
269
270 if buf.is_empty() {
271 self.pos += 1;
272 continue;
273 }
274
275 let buf_len = cmp::min(buf.len(), out_len);
276 let partial = buf_len < buf.len();
277
278 let out_pos = (next_off - out_off) as usize;
280 out[out_pos..out_pos + buf_len].copy_from_slice(&buf[..buf_len]);
281
282 self.buffered_bytes -= buf_len as u64;
283
284 out_len -= buf_len;
285
286 next_off = buf.off() + buf_len as u64;
287
288 buf.consume(buf_len);
289
290 if partial {
291 break;
293 }
294
295 self.pos += 1;
296 }
297
298 let fin = self.fin_off == Some(next_off);
305
306 self.emit_off = cmp::max(self.emit_off, next_off);
309
310 Ok((out.len() - out_len, fin))
311 }
312
313 pub fn update_max_data(&mut self, max_data: u64) {
315 self.max_data = cmp::max(self.max_data, max_data);
316 }
317
318 pub fn update_blocked_at(&mut self, blocked_at: Option<u64>) {
320 self.blocked_at = blocked_at;
321 }
322
323 pub fn blocked_at(&self) -> Option<u64> {
325 self.blocked_at
326 }
327
328 pub fn ack(&mut self, off: u64, len: usize) {
330 self.acked.insert(off..off + len as u64);
331 }
332
333 pub(crate) fn ack_fin(&mut self) {
335 self.fin_acked = true;
336 }
337
338 pub fn ack_and_drop(&mut self, off: u64, len: usize) -> usize {
339 self.ack(off, len);
340
341 let ack_off = self.ack_off();
342
343 if self.data.is_empty() {
344 return 0;
345 }
346
347 if off > ack_off {
348 return 0;
349 }
350
351 let mut drop_until = None;
352
353 for (i, buf) in self.data.iter_mut().enumerate() {
355 if buf.off >= ack_off {
358 break;
359 }
360
361 if buf.off < ack_off && ack_off < buf.max_off() {
364 break;
365 }
366
367 drop_until = Some(i);
369 }
370
371 if let Some(drop) = drop_until {
372 let dropped_len: u64 =
375 (0..=drop).map(|i| self.data[i].len() as u64).sum();
376 self.buffered_bytes = self.buffered_bytes.saturating_sub(dropped_len);
377
378 self.data.drain(..=drop);
379
380 self.pos = self.pos.saturating_sub(drop + 1);
384
385 dropped_len as usize
386 } else {
387 0
388 }
389 }
390
391 pub fn retransmit(&mut self, off: u64, len: usize) -> usize {
392 let max_off = off + len as u64;
393 let ack_off = self.ack_off();
394
395 if self.data.is_empty() {
396 return 0;
397 }
398
399 if max_off <= ack_off {
400 return 0;
401 }
402
403 let mut total_retransmitted = 0;
404
405 for i in 0..self.data.len() {
406 let buf = &mut self.data[i];
407
408 if buf.off >= max_off {
409 break;
410 }
411
412 if off > buf.max_off() {
413 continue;
414 }
415
416 let new_buf = if buf.off < max_off && max_off < buf.max_off() {
419 Some(buf.split_off((max_off - buf.off) as usize))
420 } else {
421 None
422 };
423
424 let prev_pos = buf.pos;
425
426 buf.pos = if off > buf.off && off <= buf.max_off() {
429 cmp::min(buf.pos, buf.start + (off - buf.off) as usize)
430 } else {
431 buf.start
432 };
433
434 self.pos = cmp::min(self.pos, i);
435
436 let retransmitted = (prev_pos - buf.pos) as u64;
437 self.buffered_bytes += retransmitted;
438 total_retransmitted += retransmitted;
439
440 if let Some(b) = new_buf {
441 self.data.insert(i + 1, b);
442 }
443 }
444
445 total_retransmitted as usize
446 }
447
448 pub fn reset(&mut self) -> (u64, u64) {
450 let unsent_off = cmp::max(self.off_front(), self.emit_off);
451 let unsent_len = self.off_back().saturating_sub(unsent_off);
452
453 self.fin_off = Some(unsent_off);
454
455 self.data.clear();
457
458 self.off = unsent_off;
460 self.ack(0, self.off as usize);
461
462 self.pos = 0;
463 self.buffered_bytes = 0;
464
465 (self.emit_off, unsent_len)
466 }
467
468 pub fn stop(&mut self, error_code: u64) -> Result<(u64, u64)> {
472 if self.error.is_some() {
473 return Err(Error::Done);
474 }
475
476 let (max_off, unsent) = self.reset();
477
478 self.error = Some(error_code);
479
480 Ok((max_off, unsent))
481 }
482
483 pub fn shutdown(&mut self) -> Result<(u64, u64)> {
485 if self.shutdown {
486 return Err(Error::Done);
487 }
488
489 self.shutdown = true;
490
491 Ok(self.reset())
492 }
493
494 #[cfg(test)]
496 pub fn seed_offsets_for_test(&mut self, off: u64) {
497 assert!(self.data.is_empty());
498 self.off = off;
499 self.emit_off = off;
500 }
501
502 pub fn off_back(&self) -> u64 {
504 self.off
505 }
506
507 pub fn off_front(&self) -> u64 {
509 let mut pos = self.pos;
510
511 while let Some(b) = self.data.get(pos) {
513 if !b.is_empty() {
514 return b.off();
515 }
516
517 pos += 1;
518 }
519
520 self.off
521 }
522
523 pub fn max_off(&self) -> u64 {
525 self.max_data
526 }
527
528 pub fn is_fin(&self) -> bool {
533 if self.fin_off == Some(self.off) {
534 return true;
535 }
536
537 false
538 }
539
540 pub fn is_complete(&self) -> bool {
545 if let Some(fin_off) = self.fin_off {
546 if (self.fin_acked || self.shutdown || self.error.is_some()) &&
547 self.acked == (0..fin_off)
548 {
549 return true;
550 }
551 }
552
553 false
554 }
555
556 pub fn is_stopped(&self) -> bool {
558 self.error.is_some()
559 }
560
561 pub fn has_unreported_stop(&self) -> bool {
563 self.error.is_some() && !self.error_reported && !self.shutdown
564 }
565
566 pub fn mark_stop_reported(&mut self) {
568 self.error_reported = true;
569 }
570
571 pub fn is_shutdown(&self) -> bool {
573 self.shutdown
574 }
575
576 pub fn is_empty(&self) -> bool {
578 self.data.is_empty()
579 }
580
581 pub fn ack_off(&self) -> u64 {
583 match self.acked.iter().next() {
584 Some(std::ops::Range { start: 0, end }) => end,
587
588 Some(_) | None => 0,
589 }
590 }
591
592 pub fn cap(&self) -> Result<usize> {
594 if let Some(e) = self.error {
596 return Err(Error::StreamStopped(e));
597 }
598
599 Ok((self.max_data - self.off) as usize)
600 }
601
602 #[allow(dead_code)]
604 pub fn bufs_count(&self) -> usize {
605 self.data.len()
606 }
607
608 pub fn buffered_bytes(&self) -> u64 {
614 self.buffered_bytes
615 }
616}
617
618#[cfg(test)]
619mod tests {
620 use super::*;
621
622 #[test]
623 fn empty_write() {
624 let mut buf = [0; 5];
625
626 let mut send = <SendBuf>::new(u64::MAX);
627 assert_eq!(send.buffered_bytes, 0);
628
629 let (written, fin) = send.emit(&mut buf).unwrap();
630 assert_eq!(written, 0);
631 assert!(!fin);
632 }
633
634 #[test]
635 fn multi_write() {
636 let mut buf = [0; 128];
637
638 let mut send = <SendBuf>::new(u64::MAX);
639 assert_eq!(send.buffered_bytes, 0);
640
641 let first = b"something";
642 let second = b"helloworld";
643
644 assert!(send.write(first, false).is_ok());
645 assert_eq!(send.buffered_bytes, 9);
646
647 assert!(send.write(second, true).is_ok());
648 assert_eq!(send.buffered_bytes, 19);
649
650 let (written, fin) = send.emit(&mut buf[..128]).unwrap();
651 assert_eq!(written, 19);
652 assert!(fin);
653 assert_eq!(&buf[..written], b"somethinghelloworld");
654 assert_eq!(send.buffered_bytes, 0);
655 }
656
657 #[test]
658 fn split_write() {
659 let mut buf = [0; 10];
660
661 let mut send = <SendBuf>::new(u64::MAX);
662 assert_eq!(send.buffered_bytes, 0);
663
664 let first = b"something";
665 let second = b"helloworld";
666
667 assert!(send.write(first, false).is_ok());
668 assert_eq!(send.buffered_bytes, 9);
669
670 assert!(send.write(second, true).is_ok());
671 assert_eq!(send.buffered_bytes, 19);
672
673 assert_eq!(send.off_front(), 0);
674
675 let (written, fin) = send.emit(&mut buf[..10]).unwrap();
676 assert_eq!(written, 10);
677 assert!(!fin);
678 assert_eq!(&buf[..written], b"somethingh");
679 assert_eq!(send.buffered_bytes, 9);
680
681 assert_eq!(send.off_front(), 10);
682
683 let (written, fin) = send.emit(&mut buf[..5]).unwrap();
684 assert_eq!(written, 5);
685 assert!(!fin);
686 assert_eq!(&buf[..written], b"ellow");
687 assert_eq!(send.buffered_bytes, 4);
688
689 assert_eq!(send.off_front(), 15);
690
691 let (written, fin) = send.emit(&mut buf[..10]).unwrap();
692 assert_eq!(written, 4);
693 assert!(fin);
694 assert_eq!(&buf[..written], b"orld");
695 assert_eq!(send.buffered_bytes, 0);
696
697 assert_eq!(send.off_front(), 19);
698 }
699
700 #[test]
701 fn resend() {
702 let mut buf = [0; 15];
703
704 let mut send = <SendBuf>::new(u64::MAX);
705 assert_eq!(send.buffered_bytes, 0);
706 assert_eq!(send.off_front(), 0);
707
708 let first = b"something";
709 let second = b"helloworld";
710
711 assert!(send.write(first, false).is_ok());
712 assert_eq!(send.off_front(), 0);
713
714 assert!(send.write(second, true).is_ok());
715 assert_eq!(send.off_front(), 0);
716
717 assert_eq!(send.buffered_bytes, 19);
718
719 let (written, fin) = send.emit(&mut buf[..4]).unwrap();
720 assert_eq!(written, 4);
721 assert!(!fin);
722 assert_eq!(&buf[..written], b"some");
723 assert_eq!(send.buffered_bytes, 15);
724 assert_eq!(send.off_front(), 4);
725
726 let (written, fin) = send.emit(&mut buf[..5]).unwrap();
727 assert_eq!(written, 5);
728 assert!(!fin);
729 assert_eq!(&buf[..written], b"thing");
730 assert_eq!(send.buffered_bytes, 10);
731 assert_eq!(send.off_front(), 9);
732
733 let (written, fin) = send.emit(&mut buf[..5]).unwrap();
734 assert_eq!(written, 5);
735 assert!(!fin);
736 assert_eq!(&buf[..written], b"hello");
737 assert_eq!(send.buffered_bytes, 5);
738 assert_eq!(send.off_front(), 14);
739
740 send.retransmit(4, 5);
741 assert_eq!(send.buffered_bytes, 10);
742 assert_eq!(send.off_front(), 4);
743
744 send.retransmit(0, 4);
745 assert_eq!(send.buffered_bytes, 14);
746 assert_eq!(send.off_front(), 0);
747
748 let (written, fin) = send.emit(&mut buf[..11]).unwrap();
749 assert_eq!(written, 9);
750 assert!(!fin);
751 assert_eq!(&buf[..written], b"something");
752 assert_eq!(send.buffered_bytes, 5);
753 assert_eq!(send.off_front(), 14);
754
755 let (written, fin) = send.emit(&mut buf[..11]).unwrap();
756 assert_eq!(written, 5);
757 assert!(fin);
758 assert_eq!(&buf[..written], b"world");
759 assert_eq!(send.buffered_bytes, 0);
760 assert_eq!(send.off_front(), 19);
761 }
762
763 #[test]
764 fn write_blocked_by_off() {
765 let mut buf = [0; 10];
766
767 let mut send = <SendBuf>::default();
768 assert_eq!(send.buffered_bytes, 0);
769
770 let first = b"something";
771 let second = b"helloworld";
772
773 assert_eq!(send.write(first, false), Ok(0));
774 assert_eq!(send.buffered_bytes, 0);
775
776 assert_eq!(send.write(second, true), Ok(0));
777 assert_eq!(send.buffered_bytes, 0);
778
779 send.update_max_data(5);
780
781 assert_eq!(send.write(first, false), Ok(5));
782 assert_eq!(send.buffered_bytes, 5);
783
784 assert_eq!(send.write(second, true), Ok(0));
785 assert_eq!(send.buffered_bytes, 5);
786
787 assert_eq!(send.off_front(), 0);
788
789 let (written, fin) = send.emit(&mut buf[..10]).unwrap();
790 assert_eq!(written, 5);
791 assert!(!fin);
792 assert_eq!(&buf[..written], b"somet");
793 assert_eq!(send.buffered_bytes, 0);
794
795 assert_eq!(send.off_front(), 5);
796
797 let (written, fin) = send.emit(&mut buf[..10]).unwrap();
798 assert_eq!(written, 0);
799 assert!(!fin);
800 assert_eq!(&buf[..written], b"");
801 assert_eq!(send.buffered_bytes, 0);
802
803 send.update_max_data(15);
804
805 assert_eq!(send.write(&first[5..], false), Ok(4));
806 assert_eq!(send.buffered_bytes, 4);
807
808 assert_eq!(send.write(second, true), Ok(6));
809 assert_eq!(send.buffered_bytes, 10);
810
811 assert_eq!(send.off_front(), 5);
812
813 let (written, fin) = send.emit(&mut buf[..10]).unwrap();
814 assert_eq!(written, 10);
815 assert!(!fin);
816 assert_eq!(&buf[..10], b"hinghellow");
817 assert_eq!(send.buffered_bytes, 0);
818
819 send.update_max_data(25);
820
821 assert_eq!(send.write(&second[6..], true), Ok(4));
822 assert_eq!(send.buffered_bytes, 4);
823
824 assert_eq!(send.off_front(), 15);
825
826 let (written, fin) = send.emit(&mut buf[..10]).unwrap();
827 assert_eq!(written, 4);
828 assert!(fin);
829 assert_eq!(&buf[..written], b"orld");
830 assert_eq!(send.buffered_bytes, 0);
831 }
832
833 #[test]
834 fn zero_len_write() {
835 let mut buf = [0; 10];
836
837 let mut send = <SendBuf>::new(u64::MAX);
838 assert_eq!(send.buffered_bytes, 0);
839
840 let first = b"something";
841
842 assert!(send.write(first, false).is_ok());
843 assert_eq!(send.buffered_bytes, 9);
844
845 assert!(send.write(&[], true).is_ok());
846 assert_eq!(send.buffered_bytes, 9);
847
848 assert_eq!(send.off_front(), 0);
849
850 let (written, fin) = send.emit(&mut buf[..10]).unwrap();
851 assert_eq!(written, 9);
852 assert!(fin);
853 assert_eq!(&buf[..written], b"something");
854 assert_eq!(send.buffered_bytes, 0);
855 }
856
857 #[test]
859 fn send_buf_len_on_retransmit() {
860 let mut buf = [0; 15];
861
862 let mut send = <SendBuf>::new(u64::MAX);
863 assert_eq!(send.buffered_bytes, 0);
864 assert_eq!(send.off_front(), 0);
865
866 let first = b"something";
867
868 assert!(send.write(first, false).is_ok());
869 assert_eq!(send.off_front(), 0);
870
871 assert_eq!(send.buffered_bytes, 9);
872
873 let (written, fin) = send.emit(&mut buf[..4]).unwrap();
874 assert_eq!(written, 4);
875 assert!(!fin);
876 assert_eq!(&buf[..written], b"some");
877 assert_eq!(send.buffered_bytes, 5);
878 assert_eq!(send.off_front(), 4);
879
880 send.retransmit(3, 5);
881 assert_eq!(send.buffered_bytes, 6);
882 assert_eq!(send.off_front(), 3);
883 }
884
885 #[test]
886 fn send_buf_final_size_retransmit() {
887 let mut buf = [0; 50];
888 let mut send = <SendBuf>::new(u64::MAX);
889
890 send.write(&buf, false).unwrap();
891 assert_eq!(send.off_front(), 0);
892
893 let (written, _fin) = send.emit(&mut buf).unwrap();
895 assert_eq!(written, buf.len());
896 assert_eq!(send.off_front(), buf.len() as u64);
897
898 send.retransmit(40, 10);
901
902 let (fin_off, unsent) = send.stop(0).unwrap();
906 assert_eq!(fin_off, 50);
907 assert_eq!(unsent, 0);
908 }
909}