github.com/grailbio/base@v0.0.11/syncqueue/ordered_queue.go (about)

     1  package syncqueue
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  )
     7  
     8  // OrderedQueue is a queue that orders entries on their way out.  An
     9  // inserter enqueues each entry with an index, and when the receiver
    10  // dequeues an entry, the entries arrive in the sequential order of
    11  // their indices.  An OrderedQueue has a maximum size; if the queue
    12  // contains the next entry, the queue will accept up to maxSize
    13  // entries.  If the queue does not yet contain the next entry, the
    14  // queue will block an insert that would make the size equal to
    15  // maxSize unless the new entry is the next entry to be dequeued.
    16  type OrderedQueue struct {
    17  	next    int
    18  	maxSize int
    19  	pending map[int]interface{}
    20  	cond    *sync.Cond
    21  	closed  bool
    22  	err     error
    23  }
    24  
    25  // Create a new OrderedQueue with size maxSize.
    26  func NewOrderedQueue(maxSize int) *OrderedQueue {
    27  	if maxSize < 1 {
    28  		panic("OrderedQueue must have length at least 1")
    29  	}
    30  	return &OrderedQueue{
    31  		next:    0,
    32  		maxSize: maxSize,
    33  		pending: make(map[int]interface{}),
    34  		cond:    sync.NewCond(&sync.Mutex{}),
    35  		closed:  false,
    36  	}
    37  }
    38  
    39  // Insert an entry into the queue with the specified sequence index.
    40  // Insert blocks if the insert would make the queue full and the new
    41  // entry is not the one next one in the output sequence.  The sequence
    42  // index should start at zero.
    43  func (q *OrderedQueue) Insert(index int, value interface{}) error {
    44  	q.cond.L.Lock()
    45  	defer q.cond.L.Unlock()
    46  
    47  	_, haveNext := q.pending[q.next]
    48  	for q.err == nil && ((haveNext && len(q.pending) == q.maxSize) || (!haveNext && index != q.next && len(q.pending) == q.maxSize-1)) {
    49  		q.cond.Wait()
    50  		_, haveNext = q.pending[q.next]
    51  	}
    52  	if q.err != nil {
    53  		return q.err
    54  	}
    55  	if q.closed {
    56  		panic("closed OrderedQueue before Insert finished")
    57  	}
    58  
    59  	q.pending[index] = value
    60  	if index == q.next {
    61  		q.cond.Broadcast()
    62  	}
    63  	return nil
    64  }
    65  
    66  // Close the ordered queue.  In the normal case, err should be nil,
    67  // and Close tells the OrderedQueue that all calls to Insert have
    68  // completed.  The user may not call Insert() after calling Close().
    69  // The user may close the OrderedQueue prematurely with a non-nil err
    70  // while some calls to Insert() and/or Next() are blocking; in this
    71  // case those calls that are blocking will return with err.
    72  func (q *OrderedQueue) Close(err error) error {
    73  	q.cond.L.Lock()
    74  	defer q.cond.L.Unlock()
    75  	if q.err == nil {
    76  		q.err = err
    77  	}
    78  	q.closed = true
    79  	q.cond.Broadcast()
    80  	return q.err
    81  }
    82  
    83  // Next returns true and the next entry in sequence if the next entry
    84  // is available.  Next blocks if the next entry is not yet present.
    85  // Next returns (nil, false) if there are no more entries, and the
    86  // queue is closed.
    87  func (q *OrderedQueue) Next() (value interface{}, ok bool, err error) {
    88  	q.cond.L.Lock()
    89  	defer q.cond.L.Unlock()
    90  
    91  	value, found := q.pending[q.next]
    92  	for q.err == nil && (!found && !q.closed) {
    93  		q.cond.Wait()
    94  		value, found = q.pending[q.next]
    95  	}
    96  	if q.err != nil {
    97  		return nil, false, q.err
    98  	}
    99  	if q.closed && len(q.pending) == 0 {
   100  		return nil, false, nil
   101  	}
   102  	if q.closed && !found {
   103  		panic(fmt.Sprintf("OrderedQueue is closed, but entry %d is not present", q.next))
   104  	}
   105  
   106  	delete(q.pending, q.next)
   107  	q.next++
   108  	q.cond.Broadcast()
   109  	return value, true, nil
   110  }