pipe.go 3.8 KB

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