pipe.go 4.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175
  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 || p.breakErr != nil {
  83. return 0, errClosedPipeWrite
  84. }
  85. return p.b.Write(d)
  86. }
  87. // CloseWithError causes the next Read (waking up a current blocked
  88. // Read if needed) to return the provided err after all data has been
  89. // read.
  90. //
  91. // The error must be non-nil.
  92. func (p *pipe) CloseWithError(err error) { p.closeWithError(&p.err, err, nil) }
  93. // BreakWithError causes the next Read (waking up a current blocked
  94. // Read if needed) to return the provided err immediately, without
  95. // waiting for unread data.
  96. func (p *pipe) BreakWithError(err error) { p.closeWithError(&p.breakErr, err, nil) }
  97. // closeWithErrorAndCode is like CloseWithError but also sets some code to run
  98. // in the caller's goroutine before returning the error.
  99. func (p *pipe) closeWithErrorAndCode(err error, fn func()) { p.closeWithError(&p.err, err, fn) }
  100. func (p *pipe) closeWithError(dst *error, err error, fn func()) {
  101. if err == nil {
  102. panic("err must be non-nil")
  103. }
  104. p.mu.Lock()
  105. defer p.mu.Unlock()
  106. if p.c.L == nil {
  107. p.c.L = &p.mu
  108. }
  109. defer p.c.Signal()
  110. if *dst != nil {
  111. // Already been done.
  112. return
  113. }
  114. p.readFn = fn
  115. if dst == &p.breakErr {
  116. if p.b != nil {
  117. p.unread += p.b.Len()
  118. }
  119. p.b = nil
  120. }
  121. *dst = err
  122. p.closeDoneLocked()
  123. }
  124. // requires p.mu be held.
  125. func (p *pipe) closeDoneLocked() {
  126. if p.donec == nil {
  127. return
  128. }
  129. // Close if unclosed. This isn't racy since we always
  130. // hold p.mu while closing.
  131. select {
  132. case <-p.donec:
  133. default:
  134. close(p.donec)
  135. }
  136. }
  137. // Err returns the error (if any) first set by BreakWithError or CloseWithError.
  138. func (p *pipe) Err() error {
  139. p.mu.Lock()
  140. defer p.mu.Unlock()
  141. if p.breakErr != nil {
  142. return p.breakErr
  143. }
  144. return p.err
  145. }
  146. // Done returns a channel which is closed if and when this pipe is closed
  147. // with CloseWithError.
  148. func (p *pipe) Done() <-chan struct{} {
  149. p.mu.Lock()
  150. defer p.mu.Unlock()
  151. if p.donec == nil {
  152. p.donec = make(chan struct{})
  153. if p.err != nil || p.breakErr != nil {
  154. // Already hit an error.
  155. p.closeDoneLocked()
  156. }
  157. }
  158. return p.donec
  159. }