tractor.dev/toolkit-go@v0.0.0-20241010005851-214d91207d07/duplex/mux/util_buffer.go (about)

     1  package mux
     2  
     3  import (
     4  	"io"
     5  	"sync"
     6  )
     7  
     8  // buffer provides a linked list buffer for data exchange
     9  // between producer and consumer. Theoretically the buffer is
    10  // of unlimited capacity as it does no allocation of its own.
    11  type buffer struct {
    12  	// protects concurrent access to head, tail and closed
    13  	*sync.Cond
    14  
    15  	head *element // the buffer that will be read first
    16  	tail *element // the buffer that will be read last
    17  
    18  	closed bool
    19  }
    20  
    21  // An element represents a single link in a linked list.
    22  type element struct {
    23  	buf  []byte
    24  	next *element
    25  }
    26  
    27  // newBuffer returns an empty buffer that is not closed.
    28  func newBuffer() *buffer {
    29  	e := new(element)
    30  	b := &buffer{
    31  		Cond: sync.NewCond(new(sync.Mutex)),
    32  		head: e,
    33  		tail: e,
    34  	}
    35  	return b
    36  }
    37  
    38  // write makes buf available for Read to receive.
    39  // buf must not be modified after the call to write.
    40  func (b *buffer) write(buf []byte) {
    41  	b.Cond.L.Lock()
    42  	e := &element{buf: buf}
    43  	b.tail.next = e
    44  	b.tail = e
    45  	b.Cond.Signal()
    46  	b.Cond.L.Unlock()
    47  }
    48  
    49  // eof closes the buffer. Reads from the buffer once all
    50  // the data has been consumed will receive io.EOF.
    51  func (b *buffer) eof() {
    52  	b.Cond.L.Lock()
    53  	b.closed = true
    54  	b.Cond.Signal()
    55  	b.Cond.L.Unlock()
    56  }
    57  
    58  // Read reads data from the internal buffer in buf.  Reads will block
    59  // if no data is available, or until the buffer is closed.
    60  func (b *buffer) Read(buf []byte) (n int, err error) {
    61  	b.Cond.L.Lock()
    62  	defer b.Cond.L.Unlock()
    63  
    64  	for len(buf) > 0 {
    65  		// if there is data in b.head, copy it
    66  		if len(b.head.buf) > 0 {
    67  			r := copy(buf, b.head.buf)
    68  			buf, b.head.buf = buf[r:], b.head.buf[r:]
    69  			n += r
    70  			continue
    71  		}
    72  		// if there is a next buffer, make it the head
    73  		if len(b.head.buf) == 0 && b.head != b.tail {
    74  			b.head = b.head.next
    75  			continue
    76  		}
    77  
    78  		// if at least one byte has been copied, return
    79  		if n > 0 {
    80  			break
    81  		}
    82  
    83  		// if nothing was read, and there is nothing outstanding
    84  		// check to see if the buffer is closed.
    85  		if b.closed {
    86  			err = io.EOF
    87  			break
    88  		}
    89  		// out of buffers, wait for producer
    90  		b.Cond.Wait()
    91  	}
    92  	return
    93  }