1234567891011121314151617181920212223242526272829303132333435363738394041424344 |
- package ctxreader
- import (
- "context"
- "io"
- "sync"
- "sync/atomic"
- )
- type ctxReader struct {
- r io.ReadCloser
- err atomic.Value
- closeOnce sync.Once
- }
- func (r *ctxReader) Read(p []byte) (int, error) {
- if err := r.err.Load(); err != nil {
- return 0, err.(error)
- }
- return r.r.Read(p)
- }
- func (r *ctxReader) Close() (err error) {
- r.closeOnce.Do(func() { err = r.r.Close() })
- return
- }
- func New(ctx context.Context, r io.ReadCloser, closeOnDone bool) io.ReadCloser {
- if ctx.Done() == nil {
- return r
- }
- ctxr := ctxReader{r: r}
- go func(ctx context.Context) {
- <-ctx.Done()
- ctxr.err.Store(ctx.Err())
- if closeOnDone {
- ctxr.closeOnce.Do(func() { ctxr.r.Close() })
- }
- }(ctx)
- return &ctxr
- }
|