github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/util/quotapool/notify_queue.go (about)

     1  // Copyright 2019 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package quotapool
    12  
    13  import "sync"
    14  
    15  // bufferSize is the size of the ringBuf buf served from a notifyQueueNodePool.
    16  //
    17  // Each node is 32+8*bufferSize bytes so at 28 a node is 256 bytes which
    18  // feels like a nice number.
    19  const bufferSize = 28
    20  
    21  // notifyQueue provides an allocation efficient FIFO queue for chan struct{}.
    22  //
    23  // notifyQueue is not safe for concurrent use.
    24  type notifyQueue struct {
    25  	len  int64
    26  	head *node
    27  
    28  	pool *notifyQueueNodePool
    29  }
    30  
    31  // initializeNotifyQueue initializes a notifyQueue.
    32  // notifyQueue values should not be used until they have been initialized.
    33  // It is illegal to initialize a notifyQueue more than once and this function
    34  // will panic if called with an already initialized notifyQueue.
    35  func initializeNotifyQueue(q *notifyQueue) {
    36  	if q.pool != nil {
    37  		panic("cannot re-initialize a notifyQueue")
    38  	}
    39  	defaultNotifyQueueNodePool.initialize(q)
    40  }
    41  
    42  var defaultNotifyQueueNodePool = newNotifyQueueNodePool()
    43  
    44  // enqueue adds c to the end of the queue and returns the address of the added
    45  // notifyee.
    46  func (q *notifyQueue) enqueue(c chan struct{}) (n *notifyee) {
    47  	if q.head == nil {
    48  		q.head = q.pool.pool.Get().(*node)
    49  		q.head.prev = q.head
    50  		q.head.next = q.head
    51  	}
    52  	tail := q.head.prev
    53  	if n = tail.enqueue(c); n == nil {
    54  		newTail := q.pool.pool.Get().(*node)
    55  		tail.next = newTail
    56  		q.head.prev = newTail
    57  		newTail.prev = tail
    58  		newTail.next = q.head
    59  		if n = newTail.enqueue(c); n == nil {
    60  			panic("failed to enqueue into a fresh buffer")
    61  		}
    62  	}
    63  	q.len++
    64  	return n
    65  }
    66  
    67  // dequeue removes the current head of the queue which can be accessed with
    68  // peek().
    69  func (q *notifyQueue) dequeue() {
    70  	if q.head == nil {
    71  		return
    72  	}
    73  	q.head.dequeue()
    74  	if q.head.len == 0 {
    75  		oldHead := q.head
    76  		if oldHead.next == oldHead {
    77  			q.head = nil
    78  		} else {
    79  			q.head = oldHead.next
    80  			q.head.prev = oldHead.prev
    81  			q.head.prev.next = q.head
    82  		}
    83  		*oldHead = node{}
    84  		q.pool.pool.Put(oldHead)
    85  	}
    86  	q.len--
    87  }
    88  
    89  // peek returns the current head of the queue or nil if the queue is empty.
    90  // It does not modify the queue. It is illegal to use the returned pointer after
    91  // the next call to dequeue.
    92  func (q *notifyQueue) peek() *notifyee {
    93  	if q.head == nil {
    94  		return nil
    95  	}
    96  	return &q.head.buf[q.head.head]
    97  }
    98  
    99  // notifyQueueNodePool constructs notifyQueue objects which internally pool
   100  // their buffers.
   101  type notifyQueueNodePool struct {
   102  	pool sync.Pool
   103  }
   104  
   105  // newNotifyQueueNodePool returns a new notifyQueueNodePool which can be used
   106  // to construct notifyQueues which internally pool their buffers.
   107  func newNotifyQueueNodePool() *notifyQueueNodePool {
   108  	return &notifyQueueNodePool{
   109  		pool: sync.Pool{
   110  			New: func() interface{} { return &node{} },
   111  		},
   112  	}
   113  }
   114  
   115  // initialize initializes a which will share a sync.Pool of nodes with the
   116  // other notifyQueue instances initialized with this pool.
   117  func (p *notifyQueueNodePool) initialize(q *notifyQueue) {
   118  	*q = notifyQueue{pool: p}
   119  }
   120  
   121  type node struct {
   122  	ringBuf
   123  	prev, next *node
   124  }
   125  
   126  type notifyee struct {
   127  	c chan struct{}
   128  }
   129  
   130  type ringBuf struct {
   131  	buf  [bufferSize]notifyee
   132  	head int64
   133  	len  int64
   134  }
   135  
   136  func (rb *ringBuf) enqueue(c chan struct{}) *notifyee {
   137  	if rb.len == bufferSize {
   138  		return nil
   139  	}
   140  	i := (rb.head + rb.len) % bufferSize
   141  	rb.buf[i] = notifyee{c: c}
   142  	rb.len++
   143  	return &rb.buf[i]
   144  }
   145  
   146  func (rb *ringBuf) dequeue() {
   147  	// NB: the notifyQueue never contains an empty ringBuf.
   148  	if rb.len == 0 {
   149  		panic("cannot dequeue from an empty buffer")
   150  	}
   151  	rb.buf[rb.head] = notifyee{}
   152  	rb.head++
   153  	rb.head %= bufferSize
   154  	rb.len--
   155  }