pipe.go

Documentation: golang.org/x/net/internal/quic

     1  // Copyright 2023 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  //go:build go1.21
     6  
     7  package quic
     8  
     9  import (
    10  	"sync"
    11  )
    12  
    13  // A pipe is a byte buffer used in implementing streams.
    14  //
    15  // A pipe contains a window of stream data.
    16  // Random access reads and writes are supported within the window.
    17  // Writing past the end of the window extends it.
    18  // Data may be discarded from the start of the pipe, advancing the window.
    19  type pipe struct {
    20  	start int64
    21  	end   int64
    22  	head  *pipebuf
    23  	tail  *pipebuf
    24  }
    25  
    26  type pipebuf struct {
    27  	off  int64
    28  	b    []byte
    29  	next *pipebuf
    30  }
    31  
    32  func (pb *pipebuf) end() int64 {
    33  	return pb.off + int64(len(pb.b))
    34  }
    35  
    36  var pipebufPool = sync.Pool{
    37  	New: func() any {
    38  		return &pipebuf{
    39  			b: make([]byte, 4096),
    40  		}
    41  	},
    42  }
    43  
    44  func newPipebuf() *pipebuf {
    45  	return pipebufPool.Get().(*pipebuf)
    46  }
    47  
    48  func (b *pipebuf) recycle() {
    49  	b.off = 0
    50  	b.next = nil
    51  	pipebufPool.Put(b)
    52  }
    53  
    54  // writeAt writes len(b) bytes to the pipe at offset off.
    55  //
    56  // Writes to offsets before p.start are discarded.
    57  // Writes to offsets after p.end extend the pipe window.
    58  func (p *pipe) writeAt(b []byte, off int64) {
    59  	end := off + int64(len(b))
    60  	if end > p.end {
    61  		p.end = end
    62  	} else if end <= p.start {
    63  		return
    64  	}
    65  
    66  	if off < p.start {
    67  		// Discard the portion of b which falls before p.start.
    68  		trim := p.start - off
    69  		b = b[trim:]
    70  		off = p.start
    71  	}
    72  
    73  	if p.head == nil {
    74  		p.head = newPipebuf()
    75  		p.head.off = p.start
    76  		p.tail = p.head
    77  	}
    78  	pb := p.head
    79  	if off >= p.tail.off {
    80  		// Common case: Writing past the end of the pipe.
    81  		pb = p.tail
    82  	}
    83  	for {
    84  		pboff := off - pb.off
    85  		if pboff < int64(len(pb.b)) {
    86  			n := copy(pb.b[pboff:], b)
    87  			if n == len(b) {
    88  				return
    89  			}
    90  			off += int64(n)
    91  			b = b[n:]
    92  		}
    93  		if pb.next == nil {
    94  			pb.next = newPipebuf()
    95  			pb.next.off = pb.off + int64(len(pb.b))
    96  			p.tail = pb.next
    97  		}
    98  		pb = pb.next
    99  	}
   100  }
   101  
   102  // copy copies len(b) bytes into b starting from off.
   103  // The pipe must contain [off, off+len(b)).
   104  func (p *pipe) copy(off int64, b []byte) {
   105  	dst := b[:0]
   106  	p.read(off, len(b), func(c []byte) error {
   107  		dst = append(dst, c...)
   108  		return nil
   109  	})
   110  }
   111  
   112  // read calls f with the data in [off, off+n)
   113  // The data may be provided sequentially across multiple calls to f.
   114  func (p *pipe) read(off int64, n int, f func([]byte) error) error {
   115  	if off < p.start {
   116  		panic("invalid read range")
   117  	}
   118  	for pb := p.head; pb != nil && n > 0; pb = pb.next {
   119  		if off >= pb.end() {
   120  			continue
   121  		}
   122  		b := pb.b[off-pb.off:]
   123  		if len(b) > n {
   124  			b = b[:n]
   125  		}
   126  		off += int64(len(b))
   127  		n -= len(b)
   128  		if err := f(b); err != nil {
   129  			return err
   130  		}
   131  	}
   132  	if n > 0 {
   133  		panic("invalid read range")
   134  	}
   135  	return nil
   136  }
   137  
   138  // discardBefore discards all data prior to off.
   139  func (p *pipe) discardBefore(off int64) {
   140  	for p.head != nil && p.head.end() < off {
   141  		head := p.head
   142  		p.head = p.head.next
   143  		head.recycle()
   144  	}
   145  	if p.head == nil {
   146  		p.tail = nil
   147  	}
   148  	p.start = off
   149  	p.end = max(p.end, off)
   150  }
   151
View as plain text