safe.go (2206B)
1 package io 2 3 import ( 4 "io" 5 "sync" 6 ) 7 8 // NewSafeReadCloser returns a new safeReadCloser that wraps readCloser. 9 func NewSafeReadCloser(readCloser io.ReadCloser) io.ReadCloser { 10 sr := &safeReadCloser{ 11 readCloser: readCloser, 12 } 13 14 if _, ok := readCloser.(io.WriterTo); ok { 15 return &safeWriteToReadCloser{safeReadCloser: sr} 16 } 17 18 return sr 19 } 20 21 // safeWriteToReadCloser wraps a safeReadCloser but exposes a WriteTo interface implementation. This will panic 22 // if the underlying io.ReadClose does not support WriteTo. Use NewSafeReadCloser to ensure the proper handling of this 23 // type. 24 type safeWriteToReadCloser struct { 25 *safeReadCloser 26 } 27 28 // WriteTo implements the io.WriteTo interface. 29 func (r *safeWriteToReadCloser) WriteTo(w io.Writer) (int64, error) { 30 r.safeReadCloser.mtx.Lock() 31 defer r.safeReadCloser.mtx.Unlock() 32 33 if r.safeReadCloser.closed { 34 return 0, io.EOF 35 } 36 37 return r.safeReadCloser.readCloser.(io.WriterTo).WriteTo(w) 38 } 39 40 // safeReadCloser wraps a io.ReadCloser and presents an io.ReadCloser interface. When Close is called on safeReadCloser 41 // the underlying Close method will be executed, and then the reference to the reader will be dropped. This type 42 // is meant to be used with the net/http library which will retain a reference to the request body for the lifetime 43 // of a goroutine connection. Wrapping in this manner will ensure that no data race conditions are falsely reported. 44 // This type is thread-safe. 45 type safeReadCloser struct { 46 readCloser io.ReadCloser 47 closed bool 48 mtx sync.Mutex 49 } 50 51 // Read reads up to len(p) bytes into p from the underlying read. If the reader is closed io.EOF will be returned. 52 func (r *safeReadCloser) Read(p []byte) (n int, err error) { 53 r.mtx.Lock() 54 defer r.mtx.Unlock() 55 if r.closed { 56 return 0, io.EOF 57 } 58 59 return r.readCloser.Read(p) 60 } 61 62 // Close calls the underlying io.ReadCloser's Close method, removes the reference to the reader, and returns any error 63 // reported from Close. Subsequent calls to Close will always return a nil error. 64 func (r *safeReadCloser) Close() error { 65 r.mtx.Lock() 66 defer r.mtx.Unlock() 67 if r.closed { 68 return nil 69 } 70 71 r.closed = true 72 rc := r.readCloser 73 r.readCloser = nil 74 return rc.Close() 75 }