iterator.go 8.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227
  1. // Copyright 2016 Google LLC.
  2. // Use of this source code is governed by a BSD-style
  3. // license that can be found in the LICENSE file.
  4. // Package iterator provides support for standard Google API iterators.
  5. // See https://github.com/GoogleCloudPlatform/gcloud-golang/wiki/Iterator-Guidelines.
  6. package iterator
  7. import (
  8. "errors"
  9. "fmt"
  10. "reflect"
  11. )
  12. // Done is returned by an iterator's Next method when the iteration is
  13. // complete; when there are no more items to return.
  14. var Done = errors.New("no more items in iterator")
  15. // We don't support mixed calls to Next and NextPage because they play
  16. // with the paging state in incompatible ways.
  17. var errMixed = errors.New("iterator: Next and NextPage called on same iterator")
  18. // PageInfo contains information about an iterator's paging state.
  19. type PageInfo struct {
  20. // Token is the token used to retrieve the next page of items from the
  21. // API. You may set Token immediately after creating an iterator to
  22. // begin iteration at a particular point. If Token is the empty string,
  23. // the iterator will begin with the first eligible item.
  24. //
  25. // The result of setting Token after the first call to Next is undefined.
  26. //
  27. // After the underlying API method is called to retrieve a page of items,
  28. // Token is set to the next-page token in the response.
  29. Token string
  30. // MaxSize is the maximum number of items returned by a call to the API.
  31. // Set MaxSize as a hint to optimize the buffering behavior of the iterator.
  32. // If zero, the page size is determined by the underlying service.
  33. //
  34. // Use Pager to retrieve a page of a specific, exact size.
  35. MaxSize int
  36. // The error state of the iterator. Manipulated by PageInfo.next and Pager.
  37. // This is a latch: it starts as nil, and once set should never change.
  38. err error
  39. // If true, no more calls to fetch should be made. Set to true when fetch
  40. // returns an empty page token. The iterator is Done when this is true AND
  41. // the buffer is empty.
  42. atEnd bool
  43. // Function that fetches a page from the underlying service. It should pass
  44. // the pageSize and pageToken arguments to the service, fill the buffer
  45. // with the results from the call, and return the next-page token returned
  46. // by the service. The function must not remove any existing items from the
  47. // buffer. If the underlying RPC takes an int32 page size, pageSize should
  48. // be silently truncated.
  49. fetch func(pageSize int, pageToken string) (nextPageToken string, err error)
  50. // Function that returns the number of currently buffered items.
  51. bufLen func() int
  52. // Function that returns the buffer, after setting the buffer variable to nil.
  53. takeBuf func() interface{}
  54. // Set to true on first call to PageInfo.next or Pager.NextPage. Used to check
  55. // for calls to both Next and NextPage with the same iterator.
  56. nextCalled, nextPageCalled bool
  57. }
  58. // NewPageInfo exposes internals for iterator implementations.
  59. // It is not a stable interface.
  60. var NewPageInfo = newPageInfo
  61. // newPageInfo creates and returns a PageInfo and a next func. If an iterator can
  62. // support paging, its iterator-creating method should call this. Each time the
  63. // iterator's Next is called, it should call the returned next fn to determine
  64. // whether a next item exists, and if so it should pop an item from the buffer.
  65. //
  66. // The fetch, bufLen and takeBuf arguments provide access to the iterator's
  67. // internal slice of buffered items. They behave as described in PageInfo, above.
  68. //
  69. // The return value is the PageInfo.next method bound to the returned PageInfo value.
  70. // (Returning it avoids exporting PageInfo.next.)
  71. //
  72. // Note: the returned PageInfo and next fn do not remove items from the buffer.
  73. // It is up to the iterator using these to remove items from the buffer:
  74. // typically by performing a pop in its Next. If items are not removed from the
  75. // buffer, memory may grow unbounded.
  76. func newPageInfo(fetch func(int, string) (string, error), bufLen func() int, takeBuf func() interface{}) (pi *PageInfo, next func() error) {
  77. pi = &PageInfo{
  78. fetch: fetch,
  79. bufLen: bufLen,
  80. takeBuf: takeBuf,
  81. }
  82. return pi, pi.next
  83. }
  84. // Remaining returns the number of items available before the iterator makes another API call.
  85. func (pi *PageInfo) Remaining() int { return pi.bufLen() }
  86. // next provides support for an iterator's Next function. An iterator's Next
  87. // should return the error returned by next if non-nil; else it can assume
  88. // there is at least one item in its buffer, and it should return that item and
  89. // remove it from the buffer.
  90. func (pi *PageInfo) next() error {
  91. pi.nextCalled = true
  92. if pi.err != nil { // Once we get an error, always return it.
  93. // TODO(jba): fix so users can retry on transient errors? Probably not worth it.
  94. return pi.err
  95. }
  96. if pi.nextPageCalled {
  97. pi.err = errMixed
  98. return pi.err
  99. }
  100. // Loop until we get some items or reach the end.
  101. for pi.bufLen() == 0 && !pi.atEnd {
  102. if err := pi.fill(pi.MaxSize); err != nil {
  103. pi.err = err
  104. return pi.err
  105. }
  106. if pi.Token == "" {
  107. pi.atEnd = true
  108. }
  109. }
  110. // Either the buffer is non-empty or pi.atEnd is true (or both).
  111. if pi.bufLen() == 0 {
  112. // The buffer is empty and pi.atEnd is true, i.e. the service has no
  113. // more items.
  114. pi.err = Done
  115. }
  116. return pi.err
  117. }
  118. // Call the service to fill the buffer, using size and pi.Token. Set pi.Token to the
  119. // next-page token returned by the call.
  120. // If fill returns a non-nil error, the buffer will be empty.
  121. func (pi *PageInfo) fill(size int) error {
  122. tok, err := pi.fetch(size, pi.Token)
  123. if err != nil {
  124. pi.takeBuf() // clear the buffer
  125. return err
  126. }
  127. pi.Token = tok
  128. return nil
  129. }
  130. // Pageable is implemented by iterators that support paging.
  131. type Pageable interface {
  132. // PageInfo returns paging information associated with the iterator.
  133. PageInfo() *PageInfo
  134. }
  135. // Pager supports retrieving iterator items a page at a time.
  136. type Pager struct {
  137. pageInfo *PageInfo
  138. pageSize int
  139. }
  140. // NewPager returns a pager that uses iter. Calls to its NextPage method will
  141. // obtain exactly pageSize items, unless fewer remain. The pageToken argument
  142. // indicates where to start the iteration. Pass the empty string to start at
  143. // the beginning, or pass a token retrieved from a call to Pager.NextPage.
  144. //
  145. // If you use an iterator with a Pager, you must not call Next on the iterator.
  146. func NewPager(iter Pageable, pageSize int, pageToken string) *Pager {
  147. p := &Pager{
  148. pageInfo: iter.PageInfo(),
  149. pageSize: pageSize,
  150. }
  151. p.pageInfo.Token = pageToken
  152. if pageSize <= 0 {
  153. p.pageInfo.err = errors.New("iterator: page size must be positive")
  154. }
  155. return p
  156. }
  157. // NextPage retrieves a sequence of items from the iterator and appends them
  158. // to slicep, which must be a pointer to a slice of the iterator's item type.
  159. // Exactly p.pageSize items will be appended, unless fewer remain.
  160. //
  161. // The first return value is the page token to use for the next page of items.
  162. // If empty, there are no more pages. Aside from checking for the end of the
  163. // iteration, the returned page token is only needed if the iteration is to be
  164. // resumed a later time, in another context (possibly another process).
  165. //
  166. // The second return value is non-nil if an error occurred. It will never be
  167. // the special iterator sentinel value Done. To recognize the end of the
  168. // iteration, compare nextPageToken to the empty string.
  169. //
  170. // It is possible for NextPage to return a single zero-length page along with
  171. // an empty page token when there are no more items in the iteration.
  172. func (p *Pager) NextPage(slicep interface{}) (nextPageToken string, err error) {
  173. p.pageInfo.nextPageCalled = true
  174. if p.pageInfo.err != nil {
  175. return "", p.pageInfo.err
  176. }
  177. if p.pageInfo.nextCalled {
  178. p.pageInfo.err = errMixed
  179. return "", p.pageInfo.err
  180. }
  181. if p.pageInfo.bufLen() > 0 {
  182. return "", errors.New("must call NextPage with an empty buffer")
  183. }
  184. // The buffer must be empty here, so takeBuf is a no-op. We call it just to get
  185. // the buffer's type.
  186. wantSliceType := reflect.PtrTo(reflect.ValueOf(p.pageInfo.takeBuf()).Type())
  187. if slicep == nil {
  188. return "", errors.New("nil passed to Pager.NextPage")
  189. }
  190. vslicep := reflect.ValueOf(slicep)
  191. if vslicep.Type() != wantSliceType {
  192. return "", fmt.Errorf("slicep should be of type %s, got %T", wantSliceType, slicep)
  193. }
  194. for p.pageInfo.bufLen() < p.pageSize {
  195. if err := p.pageInfo.fill(p.pageSize - p.pageInfo.bufLen()); err != nil {
  196. p.pageInfo.err = err
  197. return "", p.pageInfo.err
  198. }
  199. if p.pageInfo.Token == "" {
  200. break
  201. }
  202. }
  203. e := vslicep.Elem()
  204. e.Set(reflect.AppendSlice(e, reflect.ValueOf(p.pageInfo.takeBuf())))
  205. return p.pageInfo.Token, nil
  206. }