databuffer.go 3.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146
  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. "fmt"
  8. "sync"
  9. )
  10. // Buffer chunks are allocated from a pool to reduce pressure on GC.
  11. // The maximum wasted space per dataBuffer is 2x the largest size class,
  12. // which happens when the dataBuffer has multiple chunks and there is
  13. // one unread byte in both the first and last chunks. We use a few size
  14. // classes to minimize overheads for servers that typically receive very
  15. // small request bodies.
  16. //
  17. // TODO: Benchmark to determine if the pools are necessary. The GC may have
  18. // improved enough that we can instead allocate chunks like this:
  19. // make([]byte, max(16<<10, expectedBytesRemaining))
  20. var (
  21. dataChunkSizeClasses = []int{
  22. 1 << 10,
  23. 2 << 10,
  24. 4 << 10,
  25. 8 << 10,
  26. 16 << 10,
  27. }
  28. dataChunkPools = [...]sync.Pool{
  29. {New: func() interface{} { return make([]byte, 1<<10) }},
  30. {New: func() interface{} { return make([]byte, 2<<10) }},
  31. {New: func() interface{} { return make([]byte, 4<<10) }},
  32. {New: func() interface{} { return make([]byte, 8<<10) }},
  33. {New: func() interface{} { return make([]byte, 16<<10) }},
  34. }
  35. )
  36. func getDataBufferChunk(size int64) []byte {
  37. i := 0
  38. for ; i < len(dataChunkSizeClasses)-1; i++ {
  39. if size <= int64(dataChunkSizeClasses[i]) {
  40. break
  41. }
  42. }
  43. return dataChunkPools[i].Get().([]byte)
  44. }
  45. func putDataBufferChunk(p []byte) {
  46. for i, n := range dataChunkSizeClasses {
  47. if len(p) == n {
  48. dataChunkPools[i].Put(p)
  49. return
  50. }
  51. }
  52. panic(fmt.Sprintf("unexpected buffer len=%v", len(p)))
  53. }
  54. // dataBuffer is an io.ReadWriter backed by a list of data chunks.
  55. // Each dataBuffer is used to read DATA frames on a single stream.
  56. // The buffer is divided into chunks so the server can limit the
  57. // total memory used by a single connection without limiting the
  58. // request body size on any single stream.
  59. type dataBuffer struct {
  60. chunks [][]byte
  61. r int // next byte to read is chunks[0][r]
  62. w int // next byte to write is chunks[len(chunks)-1][w]
  63. size int // total buffered bytes
  64. expected int64 // we expect at least this many bytes in future Write calls (ignored if <= 0)
  65. }
  66. var errReadEmpty = errors.New("read from empty dataBuffer")
  67. // Read copies bytes from the buffer into p.
  68. // It is an error to read when no data is available.
  69. func (b *dataBuffer) Read(p []byte) (int, error) {
  70. if b.size == 0 {
  71. return 0, errReadEmpty
  72. }
  73. var ntotal int
  74. for len(p) > 0 && b.size > 0 {
  75. readFrom := b.bytesFromFirstChunk()
  76. n := copy(p, readFrom)
  77. p = p[n:]
  78. ntotal += n
  79. b.r += n
  80. b.size -= n
  81. // If the first chunk has been consumed, advance to the next chunk.
  82. if b.r == len(b.chunks[0]) {
  83. putDataBufferChunk(b.chunks[0])
  84. end := len(b.chunks) - 1
  85. copy(b.chunks[:end], b.chunks[1:])
  86. b.chunks[end] = nil
  87. b.chunks = b.chunks[:end]
  88. b.r = 0
  89. }
  90. }
  91. return ntotal, nil
  92. }
  93. func (b *dataBuffer) bytesFromFirstChunk() []byte {
  94. if len(b.chunks) == 1 {
  95. return b.chunks[0][b.r:b.w]
  96. }
  97. return b.chunks[0][b.r:]
  98. }
  99. // Len returns the number of bytes of the unread portion of the buffer.
  100. func (b *dataBuffer) Len() int {
  101. return b.size
  102. }
  103. // Write appends p to the buffer.
  104. func (b *dataBuffer) Write(p []byte) (int, error) {
  105. ntotal := len(p)
  106. for len(p) > 0 {
  107. // If the last chunk is empty, allocate a new chunk. Try to allocate
  108. // enough to fully copy p plus any additional bytes we expect to
  109. // receive. However, this may allocate less than len(p).
  110. want := int64(len(p))
  111. if b.expected > want {
  112. want = b.expected
  113. }
  114. chunk := b.lastChunkOrAlloc(want)
  115. n := copy(chunk[b.w:], p)
  116. p = p[n:]
  117. b.w += n
  118. b.size += n
  119. b.expected -= int64(n)
  120. }
  121. return ntotal, nil
  122. }
  123. func (b *dataBuffer) lastChunkOrAlloc(want int64) []byte {
  124. if len(b.chunks) != 0 {
  125. last := b.chunks[len(b.chunks)-1]
  126. if b.w < len(last) {
  127. return last
  128. }
  129. }
  130. chunk := getDataBufferChunk(want)
  131. b.chunks = append(b.chunks, chunk)
  132. b.w = 0
  133. return chunk
  134. }