You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

232 lines
7.1 KiB

  1. // Package queue implements a dynamic FIFO queue with a fixed upper bound
  2. // and a flexible quota mechanism to handle bursty load.
  3. package queue
  4. import (
  5. "context"
  6. "errors"
  7. "sync"
  8. )
  9. var (
  10. // ErrQueueFull is returned by the Add method of a queue when the queue has
  11. // reached its hard capacity limit.
  12. ErrQueueFull = errors.New("queue is full")
  13. // ErrNoCredit is returned by the Add method of a queue when the queue has
  14. // exceeded its soft quota and there is insufficient burst credit.
  15. ErrNoCredit = errors.New("insufficient burst credit")
  16. // ErrQueueClosed is returned by the Add method of a closed queue, and by
  17. // the Wait method of a closed empty queue.
  18. ErrQueueClosed = errors.New("queue is closed")
  19. // Sentinel errors reported by the New constructor.
  20. errHardLimit = errors.New("hard limit must be > 0 and ≥ soft quota")
  21. errBurstCredit = errors.New("burst credit must be non-negative")
  22. )
  23. // A Queue is a limited-capacity FIFO queue of arbitrary data items.
  24. //
  25. // A queue has a soft quota and a hard limit on the number of items that may be
  26. // contained in the queue. Adding items in excess of the hard limit will fail
  27. // unconditionally.
  28. //
  29. // For items in excess of the soft quota, a credit system applies: Each queue
  30. // maintains a burst credit score. Adding an item in excess of the soft quota
  31. // costs 1 unit of burst credit. If there is not enough burst credit, the add
  32. // will fail.
  33. //
  34. // The initial burst credit is assigned when the queue is constructed. Removing
  35. // items from the queue adds additional credit if the resulting queue length is
  36. // less than the current soft quota. Burst credit is capped by the hard limit.
  37. //
  38. // A Queue is safe for concurrent use by multiple goroutines.
  39. type Queue struct {
  40. mu sync.Mutex // protects the fields below
  41. softQuota int // adjusted dynamically (see Add, Remove)
  42. hardLimit int // fixed for the lifespan of the queue
  43. queueLen int // number of entries in the queue list
  44. credit float64 // current burst credit
  45. closed bool
  46. nempty *sync.Cond
  47. back *entry
  48. front *entry
  49. // The queue is singly-linked. Front points to the sentinel and back points
  50. // to the newest entry. The oldest entry is front.link if it exists.
  51. }
  52. // New constructs a new empty queue with the specified options. It reports an
  53. // error if any of the option values are invalid.
  54. func New(opts Options) (*Queue, error) {
  55. if opts.HardLimit <= 0 || opts.HardLimit < opts.SoftQuota {
  56. return nil, errHardLimit
  57. }
  58. if opts.BurstCredit < 0 {
  59. return nil, errBurstCredit
  60. }
  61. if opts.SoftQuota <= 0 {
  62. opts.SoftQuota = opts.HardLimit
  63. }
  64. if opts.BurstCredit == 0 {
  65. opts.BurstCredit = float64(opts.SoftQuota)
  66. }
  67. sentinel := new(entry)
  68. q := &Queue{
  69. softQuota: opts.SoftQuota,
  70. hardLimit: opts.HardLimit,
  71. credit: opts.BurstCredit,
  72. back: sentinel,
  73. front: sentinel,
  74. }
  75. q.nempty = sync.NewCond(&q.mu)
  76. return q, nil
  77. }
  78. // Add adds item to the back of the queue. It reports an error and does not
  79. // enqueue the item if the queue is full or closed, or if it exceeds its soft
  80. // quota and there is not enough burst credit.
  81. func (q *Queue) Add(item interface{}) error {
  82. q.mu.Lock()
  83. defer q.mu.Unlock()
  84. if q.closed {
  85. return ErrQueueClosed
  86. }
  87. if q.queueLen >= q.softQuota {
  88. if q.queueLen == q.hardLimit {
  89. return ErrQueueFull
  90. } else if q.credit < 1 {
  91. return ErrNoCredit
  92. }
  93. // Successfully exceeding the soft quota deducts burst credit and raises
  94. // the soft quota. This has the effect of reducing the credit cap and the
  95. // amount of credit given for removing items to better approximate the
  96. // rate at which the consumer is servicing the queue.
  97. q.credit--
  98. q.softQuota = q.queueLen + 1
  99. }
  100. e := &entry{item: item}
  101. q.back.link = e
  102. q.back = e
  103. q.queueLen++
  104. if q.queueLen == 1 { // was empty
  105. q.nempty.Signal()
  106. }
  107. return nil
  108. }
  109. // Remove removes and returns the frontmost (oldest) item in the queue and
  110. // reports whether an item was available. If the queue is empty, Remove
  111. // returns nil, false.
  112. func (q *Queue) Remove() (interface{}, bool) {
  113. q.mu.Lock()
  114. defer q.mu.Unlock()
  115. if q.queueLen == 0 {
  116. return nil, false
  117. }
  118. return q.popFront(), true
  119. }
  120. // Wait blocks until q is non-empty or closed, and then returns the frontmost
  121. // (oldest) item from the queue. If ctx ends before an item is available, Wait
  122. // returns a nil value and a context error. If the queue is closed while it is
  123. // still empty, Wait returns nil, ErrQueueClosed.
  124. func (q *Queue) Wait(ctx context.Context) (interface{}, error) {
  125. // If the context terminates, wake the waiter.
  126. ctx, cancel := context.WithCancel(ctx)
  127. defer cancel()
  128. go func() { <-ctx.Done(); q.nempty.Broadcast() }()
  129. q.mu.Lock()
  130. defer q.mu.Unlock()
  131. for q.queueLen == 0 {
  132. if q.closed {
  133. return nil, ErrQueueClosed
  134. }
  135. select {
  136. case <-ctx.Done():
  137. return nil, ctx.Err()
  138. default:
  139. q.nempty.Wait()
  140. }
  141. }
  142. return q.popFront(), nil
  143. }
  144. // Close closes the queue. After closing, any further Add calls will report an
  145. // error, but items that were added to the queue prior to closing will still be
  146. // available for Remove and Wait. Wait will report an error without blocking if
  147. // it is called on a closed, empty queue.
  148. func (q *Queue) Close() error {
  149. q.mu.Lock()
  150. defer q.mu.Unlock()
  151. q.closed = true
  152. q.nempty.Broadcast()
  153. return nil
  154. }
  155. // popFront removes the frontmost item of q and returns its value after
  156. // updating quota and credit settings.
  157. //
  158. // Preconditions: The caller holds q.mu and q is not empty.
  159. func (q *Queue) popFront() interface{} {
  160. e := q.front.link
  161. q.front.link = e.link
  162. if e == q.back {
  163. q.back = q.front
  164. }
  165. q.queueLen--
  166. if q.queueLen < q.softQuota {
  167. // Successfully removing items from the queue below half the soft quota
  168. // lowers the soft quota. This has the effect of increasing the credit cap
  169. // and the amount of credit given for removing items to better approximate
  170. // the rate at which the consumer is servicing the queue.
  171. if q.softQuota > 1 && q.queueLen < q.softQuota/2 {
  172. q.softQuota--
  173. }
  174. // Give credit for being below the soft quota. Note we do this after
  175. // adjusting the quota so the credit reflects the item we just removed.
  176. q.credit += float64(q.softQuota-q.queueLen) / float64(q.softQuota)
  177. if cap := float64(q.hardLimit - q.softQuota); q.credit > cap {
  178. q.credit = cap
  179. }
  180. }
  181. return e.item
  182. }
  183. // Options are the initial settings for a Queue.
  184. type Options struct {
  185. // The maximum number of items the queue will ever be permitted to hold.
  186. // This value must be positive, and greater than or equal to SoftQuota. The
  187. // hard limit is fixed and does not change as the queue is used.
  188. //
  189. // The hard limit should be chosen to exceed the largest burst size expected
  190. // under normal operating conditions.
  191. HardLimit int
  192. // The initial expected maximum number of items the queue should contain on
  193. // an average workload. If this value is zero, it is initialized to the hard
  194. // limit. The soft quota is adjusted from the initial value dynamically as
  195. // the queue is used.
  196. SoftQuota int
  197. // The initial burst credit score. This value must be greater than or equal
  198. // to zero. If it is zero, the soft quota is used.
  199. BurstCredit float64
  200. }
  201. type entry struct {
  202. item interface{}
  203. link *entry
  204. }