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

     1  package mux
     2  
     3  import (
     4  	"io"
     5  )
     6  
     7  func Pair() (a, b Session) {
     8  	ar, bw := io.Pipe()
     9  	br, aw := io.Pipe()
    10  	abuf := newBufferedPipeWriter(aw, 4)
    11  	bbuf := newBufferedPipeWriter(bw, 4)
    12  	a, _ = DialIO(abuf, ar)
    13  	b, _ = DialIO(bbuf, br)
    14  	return
    15  }
    16  
    17  type bufferedPipeWriter struct {
    18  	dataCh  chan []byte
    19  	closeCh chan struct{}
    20  	closed  bool
    21  }
    22  
    23  func newBufferedPipeWriter(pw *io.PipeWriter, bufferSize int) *bufferedPipeWriter {
    24  	dataCh := make(chan []byte, bufferSize)
    25  	closeCh := make(chan struct{})
    26  
    27  	go func() {
    28  		defer pw.Close()
    29  		for {
    30  			select {
    31  			case data := <-dataCh:
    32  				pw.Write(data)
    33  			case <-closeCh:
    34  				return
    35  			}
    36  		}
    37  	}()
    38  
    39  	return &bufferedPipeWriter{
    40  		dataCh:  dataCh,
    41  		closeCh: closeCh,
    42  		closed:  false,
    43  	}
    44  }
    45  
    46  func (w *bufferedPipeWriter) Write(p []byte) (n int, err error) {
    47  	if w.closed {
    48  		return 0, io.ErrClosedPipe
    49  	}
    50  
    51  	data := make([]byte, len(p))
    52  	copy(data, p)
    53  	w.dataCh <- data
    54  	return len(p), nil
    55  }
    56  
    57  func (w *bufferedPipeWriter) Close() error {
    58  	if w.closed {
    59  		return io.ErrClosedPipe
    60  	}
    61  	w.closed = true
    62  	close(w.closeCh)
    63  	return nil
    64  }