mirror of
https://github.com/criyle/go-sandbox.git
synced 2025-11-04 14:49:53 +08:00
60 lines
1.3 KiB
Go
60 lines
1.3 KiB
Go
// Package pipe provides a wrapper to create a pipe and
|
|
// collect at most max bytes from the reader side
|
|
package pipe
|
|
|
|
import (
|
|
"bytes"
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
)
|
|
|
|
// Buffer is used to create a writable pipe and read
|
|
// at most max bytes to a buffer
|
|
type Buffer struct {
|
|
W *os.File
|
|
Buffer *bytes.Buffer
|
|
Done <-chan struct{}
|
|
Max int64
|
|
}
|
|
|
|
// NewPipe create a pipe with a goroutine to copy its read-end to writer
|
|
// returns the write end and signal for finish
|
|
// caller need to close w
|
|
func NewPipe(writer io.Writer, n int64) (<-chan struct{}, *os.File, error) {
|
|
r, w, err := os.Pipe()
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
done := make(chan struct{})
|
|
go func() {
|
|
io.CopyN(writer, r, int64(n))
|
|
close(done)
|
|
// ensure no blocking / SIGPIPE on the other end
|
|
io.Copy(io.Discard, r)
|
|
r.Close()
|
|
}()
|
|
return done, w, nil
|
|
}
|
|
|
|
// NewBuffer creates a os pipe, caller need to
|
|
// caller need to close w
|
|
// Notice: if rely on done for finish, w need be closed in parent process
|
|
func NewBuffer(max int64) (*Buffer, error) {
|
|
buffer := new(bytes.Buffer)
|
|
done, w, err := NewPipe(buffer, max+1)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &Buffer{
|
|
W: w,
|
|
Max: max,
|
|
Buffer: buffer,
|
|
Done: done,
|
|
}, nil
|
|
}
|
|
|
|
func (b Buffer) String() string {
|
|
return fmt.Sprintf("Buffer[%d/%d]", b.Buffer.Len(), b.Max)
|
|
}
|