/internal/helper/lines/send.go

https://gitlab.com/gitlab-org/gitaly · Go · 158 lines · 93 code · 27 blank · 38 comment · 26 complexity · d8c780369f12a5decb2bc2aa9fe89d9e MD5 · raw file

  1. package lines
  2. import (
  3. "bufio"
  4. "bytes"
  5. "fmt"
  6. "io"
  7. "regexp"
  8. )
  9. // SenderOpts contains fields that Send() uses to determine what is considered
  10. // a line, and how to handle pagination. That is, how many lines to skip, before
  11. // a line gets fed into the Sender.
  12. type SenderOpts struct {
  13. // Delimiter is the separator used to split the sender's output into
  14. // lines. Defaults to an empty byte (0).
  15. Delimiter byte
  16. // Limit is the upper limit of how many lines will be sent. The zero
  17. // value will cause no lines to be sent.
  18. Limit int
  19. // IsPageToken allows control over which results are sent as part of the
  20. // response. When IsPageToken evaluates to true for the first time,
  21. // results will start to be sent as part of the response. This function
  22. // will be called with an empty slice previous to sending the first line
  23. // in order to allow sending everything right from the beginning.
  24. IsPageToken func([]byte) bool
  25. // When PageTokenError is true than Sender will return an error when provided
  26. // PageToken is not found.
  27. PageTokenError bool
  28. // Filter limits sent results to those that pass the filter. The zero
  29. // value (nil) disables filtering.
  30. Filter *regexp.Regexp
  31. }
  32. // ItemsPerMessage establishes the threshold to flush the buffer when using the
  33. // `Send` function. It's a variable instead of a constant to make it possible to
  34. // override in tests.
  35. var ItemsPerMessage = 20
  36. // Sender handles a buffer of lines from a Git command
  37. type Sender func([][]byte) error
  38. type writer struct {
  39. sender Sender
  40. lines [][]byte
  41. options SenderOpts
  42. }
  43. // CopyAndAppend adds a newly allocated copy of `e` to the `s` slice. Useful to
  44. // avoid io buffer shennanigans
  45. func CopyAndAppend(s [][]byte, e []byte) [][]byte {
  46. line := make([]byte, len(e))
  47. copy(line, e)
  48. return append(s, line)
  49. }
  50. // flush calls the `sender` handler function with the accumulated lines and
  51. // clears the lines buffer.
  52. func (w *writer) flush() error {
  53. if len(w.lines) == 0 { // No message to send, just return
  54. return nil
  55. }
  56. if err := w.sender(w.lines); err != nil {
  57. return err
  58. }
  59. // Reset the message
  60. w.lines = nil
  61. return nil
  62. }
  63. // addLine adds a new line to the writer buffer, and flushes if the maximum
  64. // size has been achieved
  65. func (w *writer) addLine(p []byte) error {
  66. w.lines = CopyAndAppend(w.lines, p)
  67. if len(w.lines) >= ItemsPerMessage {
  68. return w.flush()
  69. }
  70. return nil
  71. }
  72. // consume reads from an `io.Reader` and writes each line to the buffer. It
  73. // flushes after being done reading.
  74. func (w *writer) consume(r io.Reader) error {
  75. buf := bufio.NewReader(r)
  76. // As `IsPageToken` will instruct us to send the _next_ line only, we
  77. // need to call it before the first iteration to allow for the case
  78. // where we want to send right from the beginning.
  79. pastPageToken := w.options.IsPageToken([]byte{})
  80. for i := 0; i < w.options.Limit; {
  81. var line []byte
  82. for {
  83. // delim can be multiple bytes, so we read till the end byte of it ...
  84. chunk, err := buf.ReadBytes(w.delimiter())
  85. if err != nil && err != io.EOF {
  86. return err
  87. }
  88. line = append(line, chunk...)
  89. // ... then we check if the last bytes of line are the same as delim
  90. if bytes.HasSuffix(line, []byte{w.delimiter()}) {
  91. break
  92. }
  93. if err == io.EOF {
  94. i = w.options.Limit // Implicit exit clause for the loop
  95. break
  96. }
  97. }
  98. line = bytes.TrimSuffix(line, []byte{w.delimiter()})
  99. if len(line) == 0 {
  100. break
  101. }
  102. // If a page token is given, we need to skip all lines until we've found it.
  103. // All remaining lines will then be sent until we reach the pagination limit.
  104. if !pastPageToken {
  105. pastPageToken = w.options.IsPageToken(line)
  106. continue
  107. }
  108. if w.filter() != nil && !w.filter().Match(line) {
  109. continue
  110. }
  111. i++ // Only increment the counter if the result wasn't skipped
  112. if err := w.addLine(line); err != nil {
  113. return err
  114. }
  115. }
  116. if !pastPageToken && w.options.PageTokenError {
  117. return fmt.Errorf("could not find page token")
  118. }
  119. return w.flush()
  120. }
  121. func (w *writer) delimiter() byte { return w.options.Delimiter }
  122. func (w *writer) filter() *regexp.Regexp { return w.options.Filter }
  123. // Send reads output from `r`, splits it at `opts.Delimiter`, then handles the
  124. // buffered lines using `sender`.
  125. func Send(r io.Reader, sender Sender, opts SenderOpts) error {
  126. if opts.IsPageToken == nil {
  127. opts.IsPageToken = func(_ []byte) bool { return true }
  128. }
  129. writer := &writer{sender: sender, options: opts}
  130. return writer.consume(r)
  131. }