github.com/criyle/go-sandbox@v0.10.3/pkg/pipe/buffer.go (about)

     1  // Package pipe provides a wrapper to create a pipe and
     2  // collect at most max bytes from the reader side
     3  package pipe
     4  
     5  import (
     6  	"bytes"
     7  	"fmt"
     8  	"io"
     9  	"os"
    10  )
    11  
    12  // Buffer is used to create a writable pipe and read
    13  // at most max bytes to a buffer
    14  type Buffer struct {
    15  	W      *os.File
    16  	Buffer *bytes.Buffer
    17  	Done   <-chan struct{}
    18  	Max    int64
    19  }
    20  
    21  // NewPipe create a pipe with a goroutine to copy its read-end to writer
    22  // returns the write end and signal for finish
    23  // caller need to close w
    24  func NewPipe(writer io.Writer, n int64) (<-chan struct{}, *os.File, error) {
    25  	r, w, err := os.Pipe()
    26  	if err != nil {
    27  		return nil, nil, err
    28  	}
    29  	done := make(chan struct{})
    30  	go func() {
    31  		io.CopyN(writer, r, int64(n))
    32  		close(done)
    33  		// ensure no blocking / SIGPIPE on the other end
    34  		io.Copy(io.Discard, r)
    35  		r.Close()
    36  	}()
    37  	return done, w, nil
    38  }
    39  
    40  // NewBuffer creates a os pipe, caller need to
    41  // caller need to close w
    42  // Notice: if rely on done for finish, w need be closed in parent process
    43  func NewBuffer(max int64) (*Buffer, error) {
    44  	buffer := new(bytes.Buffer)
    45  	done, w, err := NewPipe(buffer, max+1)
    46  	if err != nil {
    47  		return nil, err
    48  	}
    49  	return &Buffer{
    50  		W:      w,
    51  		Max:    max,
    52  		Buffer: buffer,
    53  		Done:   done,
    54  	}, nil
    55  }
    56  
    57  func (b Buffer) String() string {
    58  	return fmt.Sprintf("Buffer[%d/%d]", b.Buffer.Len(), b.Max)
    59  }