pipe.go 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179
  1. // Copyright 2014 The Go Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style
  3. // license that can be found in the LICENSE file.
  4. package http2
  5. import (
  6. "errors"
  7. "io"
  8. "sync"
  9. )
  10. // pipe is a goroutine-safe io.Reader/io.Writer pair. It's like
  11. // io.Pipe except there are no PipeReader/PipeWriter halves, and the
  12. // underlying buffer is an interface. (io.Pipe is always unbuffered)
  13. type pipe struct {
  14. mu sync.Mutex
  15. c sync.Cond // c.L lazily initialized to &p.mu
  16. b pipeBuffer // nil when done reading
  17. unread int // bytes unread when done
  18. err error // read error once empty. non-nil means closed.
  19. breakErr error // immediate read error (caller doesn't see rest of b)
  20. donec chan struct{} // closed on error
  21. readFn func() // optional code to run in Read before error
  22. }
  23. type pipeBuffer interface {
  24. Len() int
  25. io.Writer
  26. io.Reader
  27. }
  28. // setBuffer initializes the pipe buffer.
  29. // It has no effect if the pipe is already closed.
  30. func (p *pipe) setBuffer(b pipeBuffer) {
  31. p.mu.Lock()
  32. defer p.mu.Unlock()
  33. if p.err != nil || p.breakErr != nil {
  34. return
  35. }
  36. p.b = b
  37. }
  38. func (p *pipe) Len() int {
  39. p.mu.Lock()
  40. defer p.mu.Unlock()
  41. if p.b == nil {
  42. return p.unread
  43. }
  44. return p.b.Len()
  45. }
  46. // Read waits until data is available and copies bytes
  47. // from the buffer into p.
  48. func (p *pipe) Read(d []byte) (n int, err error) {
  49. p.mu.Lock()
  50. defer p.mu.Unlock()
  51. if p.c.L == nil {
  52. p.c.L = &p.mu
  53. }
  54. for {
  55. if p.breakErr != nil {
  56. return 0, p.breakErr
  57. }
  58. if p.b != nil && p.b.Len() > 0 {
  59. return p.b.Read(d)
  60. }
  61. if p.err != nil {
  62. if p.readFn != nil {
  63. p.readFn() // e.g. copy trailers
  64. p.readFn = nil // not sticky like p.err
  65. }
  66. p.b = nil
  67. return 0, p.err
  68. }
  69. p.c.Wait()
  70. }
  71. }
  72. var errClosedPipeWrite = errors.New("write on closed buffer")
  73. // Write copies bytes from p into the buffer and wakes a reader.
  74. // It is an error to write more data than the buffer can hold.
  75. func (p *pipe) Write(d []byte) (n int, err error) {
  76. p.mu.Lock()
  77. defer p.mu.Unlock()
  78. if p.c.L == nil {
  79. p.c.L = &p.mu
  80. }
  81. defer p.c.Signal()
  82. if p.err != nil {
  83. return 0, errClosedPipeWrite
  84. }
  85. if p.breakErr != nil {
  86. p.unread += len(d)
  87. return len(d), nil // discard when there is no reader
  88. }
  89. return p.b.Write(d)
  90. }
  91. // CloseWithError causes the next Read (waking up a current blocked
  92. // Read if needed) to return the provided err after all data has been
  93. // read.
  94. //
  95. // The error must be non-nil.
  96. func (p *pipe) CloseWithError(err error) { p.closeWithError(&p.err, err, nil) }
  97. // BreakWithError causes the next Read (waking up a current blocked
  98. // Read if needed) to return the provided err immediately, without
  99. // waiting for unread data.
  100. func (p *pipe) BreakWithError(err error) { p.closeWithError(&p.breakErr, err, nil) }
  101. // closeWithErrorAndCode is like CloseWithError but also sets some code to run
  102. // in the caller's goroutine before returning the error.
  103. func (p *pipe) closeWithErrorAndCode(err error, fn func()) { p.closeWithError(&p.err, err, fn) }
  104. func (p *pipe) closeWithError(dst *error, err error, fn func()) {
  105. if err == nil {
  106. panic("err must be non-nil")
  107. }
  108. p.mu.Lock()
  109. defer p.mu.Unlock()
  110. if p.c.L == nil {
  111. p.c.L = &p.mu
  112. }
  113. defer p.c.Signal()
  114. if *dst != nil {
  115. // Already been done.
  116. return
  117. }
  118. p.readFn = fn
  119. if dst == &p.breakErr {
  120. if p.b != nil {
  121. p.unread += p.b.Len()
  122. }
  123. p.b = nil
  124. }
  125. *dst = err
  126. p.closeDoneLocked()
  127. }
  128. // requires p.mu be held.
  129. func (p *pipe) closeDoneLocked() {
  130. if p.donec == nil {
  131. return
  132. }
  133. // Close if unclosed. This isn't racy since we always
  134. // hold p.mu while closing.
  135. select {
  136. case <-p.donec:
  137. default:
  138. close(p.donec)
  139. }
  140. }
  141. // Err returns the error (if any) first set by BreakWithError or CloseWithError.
  142. func (p *pipe) Err() error {
  143. p.mu.Lock()
  144. defer p.mu.Unlock()
  145. if p.breakErr != nil {
  146. return p.breakErr
  147. }
  148. return p.err
  149. }
  150. // Done returns a channel which is closed if and when this pipe is closed
  151. // with CloseWithError.
  152. func (p *pipe) Done() <-chan struct{} {
  153. p.mu.Lock()
  154. defer p.mu.Unlock()
  155. if p.donec == nil {
  156. p.donec = make(chan struct{})
  157. if p.err != nil || p.breakErr != nil {
  158. // Already hit an error.
  159. p.closeDoneLocked()
  160. }
  161. }
  162. return p.donec
  163. }