github.com/philippseith/signalr@v0.6.3/ctxpipe.go (about) 1 package signalr 2 3 import ( 4 "context" 5 "errors" 6 "io" 7 "sync" 8 ) 9 10 // onceError is an object that will only store an error once. 11 type onceError struct { 12 sync.Mutex // guards following 13 err error 14 } 15 16 func (a *onceError) Store(err error) { 17 a.Lock() 18 defer a.Unlock() 19 if a.err != nil { 20 return 21 } 22 a.err = err 23 } 24 func (a *onceError) Load() error { 25 a.Lock() 26 defer a.Unlock() 27 return a.err 28 } 29 30 // ErrClosedPipe is the error used for read or write operations on a closed pipe. 31 var ErrClosedPipe = errors.New("io: read/write on closed pipe") 32 33 // A pipe is the shared pipe structure underlying PipeReader and PipeWriter. 34 type pipe struct { 35 wrMu sync.Mutex // Serializes Write operations 36 wrCh chan []byte 37 rdCh chan int 38 39 once sync.Once // Protects closing done 40 ctx context.Context 41 cancel context.CancelFunc 42 rErr onceError 43 wErr onceError 44 } 45 46 func (p *pipe) Read(b []byte) (n int, err error) { 47 select { 48 case <-p.ctx.Done(): 49 return 0, p.readCloseError() 50 default: 51 } 52 53 select { 54 case bw := <-p.wrCh: 55 nr := copy(b, bw) 56 p.rdCh <- nr 57 return nr, nil 58 case <-p.ctx.Done(): 59 return 0, p.readCloseError() 60 } 61 } 62 63 func (p *pipe) readCloseError() error { 64 rErr := p.rErr.Load() 65 if wErr := p.wErr.Load(); rErr == nil && wErr != nil { 66 return wErr 67 } 68 return ErrClosedPipe 69 } 70 71 func (p *pipe) CloseRead(err error) error { 72 if err == nil { 73 err = ErrClosedPipe 74 } 75 p.rErr.Store(err) 76 p.once.Do(func() { p.cancel() }) 77 return nil 78 } 79 80 func (p *pipe) Write(b []byte) (n int, err error) { 81 select { 82 case <-p.ctx.Done(): 83 return 0, p.writeCloseError() 84 default: 85 p.wrMu.Lock() 86 defer p.wrMu.Unlock() 87 } 88 89 for once := true; once || len(b) > 0; once = false { 90 select { 91 case p.wrCh <- b: 92 nw := <-p.rdCh 93 b = b[nw:] 94 n += nw 95 case <-p.ctx.Done(): 96 return n, p.writeCloseError() 97 } 98 } 99 return n, nil 100 } 101 102 func (p *pipe) writeCloseError() error { 103 wErr := p.wErr.Load() 104 if rErr := p.rErr.Load(); wErr == nil && rErr != nil { 105 return rErr 106 } 107 return ErrClosedPipe 108 } 109 110 func (p *pipe) CloseWrite(err error) error { 111 if err == nil { 112 err = io.EOF 113 } 114 p.wErr.Store(err) 115 p.once.Do(func() { p.cancel() }) 116 return nil 117 } 118 119 // A PipeReader is the read half of a pipe. 120 type PipeReader struct { 121 p *pipe 122 } 123 124 // Read implements the standard Read interface: 125 // it reads data from the pipe, blocking until a writer 126 // arrives or the write end is closed. 127 // If the write end is closed with an error, that error is 128 // returned as err; otherwise err is EOF. 129 func (r *PipeReader) Read(data []byte) (n int, err error) { 130 return r.p.Read(data) 131 } 132 133 // Close closes the reader; subsequent writes to the 134 // write half of the pipe will return the error ErrClosedPipe. 135 func (r *PipeReader) Close() error { 136 return r.CloseWithError(nil) 137 } 138 139 // CloseWithError closes the reader; subsequent writes 140 // to the write half of the pipe will return the error err. 141 // 142 // CloseWithError never overwrites the previous error if it exists 143 // and always returns nil. 144 func (r *PipeReader) CloseWithError(err error) error { 145 return r.p.CloseRead(err) 146 } 147 148 // A PipeWriter is the write half of a pipe. 149 type PipeWriter struct { 150 p *pipe 151 } 152 153 // Write implements the standard Write interface: 154 // it writes data to the pipe, blocking until one or more readers 155 // have consumed all the data or the read end is closed. 156 // If the read end is closed with an error, that err is 157 // returned as err; otherwise err is ErrClosedPipe. 158 func (w *PipeWriter) Write(data []byte) (n int, err error) { 159 return w.p.Write(data) 160 } 161 162 // Close closes the writer; subsequent reads from the 163 // read half of the pipe will return no bytes and EOF. 164 func (w *PipeWriter) Close() error { 165 return w.CloseWithError(nil) 166 } 167 168 // CloseWithError closes the writer; subsequent reads from the 169 // read half of the pipe will return no bytes and the error err, 170 // or EOF if err is nil. 171 // 172 // CloseWithError never overwrites the previous error if it exists 173 // and always returns nil. 174 func (w *PipeWriter) CloseWithError(err error) error { 175 return w.p.CloseWrite(err) 176 } 177 178 // CtxPipe creates a synchronous in-memory pipe. 179 // It can be used to connect code expecting an io.Reader 180 // with code expecting an io.Writer. 181 // 182 // By canceling the context, Read and Write can be canceled 183 // 184 // Reads and Writes on the pipe are matched one to one 185 // except when multiple Reads are needed to consume a single Write. 186 // That is, each Write to the PipeWriter blocks until it has satisfied 187 // one or more Reads from the PipeReader that fully consume 188 // the written data. 189 // The data is copied directly from the Write to the corresponding 190 // Read (or Reads); there is no internal buffering. 191 // 192 // It is safe to call Read and Write in parallel with each other or with Close. 193 // Parallel calls to Read and parallel calls to Write are also safe: 194 // the individual calls will be gated sequentially. 195 func CtxPipe(ctx context.Context) (*PipeReader, *PipeWriter) { 196 p := &pipe{ 197 wrCh: make(chan []byte), 198 rdCh: make(chan int), 199 } 200 p.ctx, p.cancel = context.WithCancel(ctx) 201 return &PipeReader{p}, &PipeWriter{p} 202 }