ctxreader.go 722 B

1234567891011121314151617181920212223242526272829303132333435363738394041424344
  1. package ctxreader
  2. import (
  3. "context"
  4. "io"
  5. "sync"
  6. "sync/atomic"
  7. )
  8. type ctxReader struct {
  9. r io.ReadCloser
  10. err atomic.Value
  11. closeOnce sync.Once
  12. }
  13. func (r *ctxReader) Read(p []byte) (int, error) {
  14. if err := r.err.Load(); err != nil {
  15. return 0, err.(error)
  16. }
  17. return r.r.Read(p)
  18. }
  19. func (r *ctxReader) Close() (err error) {
  20. r.closeOnce.Do(func() { err = r.r.Close() })
  21. return
  22. }
  23. func New(ctx context.Context, r io.ReadCloser, closeOnDone bool) io.ReadCloser {
  24. if ctx.Done() == nil {
  25. return r
  26. }
  27. ctxr := ctxReader{r: r}
  28. go func(ctx context.Context) {
  29. <-ctx.Done()
  30. ctxr.err.Store(ctx.Err())
  31. if closeOnDone {
  32. ctxr.closeOnce.Do(func() { ctxr.r.Close() })
  33. }
  34. }(ctx)
  35. return &ctxr
  36. }