github.com/koko1123/flow-go-1@v0.29.6/network/queue/messageQueue.go (about)

     1  package queue
     2  
     3  import (
     4  	"container/heap"
     5  	"context"
     6  	"fmt"
     7  	"sync"
     8  	"time"
     9  
    10  	"github.com/koko1123/flow-go-1/module"
    11  )
    12  
    13  type Priority int
    14  
    15  const LowPriority = Priority(1)
    16  const MediumPriority = Priority(5)
    17  const HighPriority = Priority(10)
    18  
    19  // MessagePriorityFunc - the callback function to derive priority of a message
    20  type MessagePriorityFunc func(message interface{}) (Priority, error)
    21  
    22  // MessageQueue is the heap based priority queue implementation of the MessageQueue implementation
    23  type MessageQueue struct {
    24  	pq           *priorityQueue
    25  	cond         *sync.Cond
    26  	priorityFunc MessagePriorityFunc
    27  	ctx          context.Context
    28  	metrics      module.NetworkInboundQueueMetrics
    29  }
    30  
    31  func (mq *MessageQueue) Insert(message interface{}) error {
    32  
    33  	if err := mq.ctx.Err(); err != nil {
    34  		return err
    35  	}
    36  
    37  	// determine the message priority
    38  	priority, err := mq.priorityFunc(message)
    39  	if err != nil {
    40  		return fmt.Errorf("failed to derive message priority: %w", err)
    41  	}
    42  
    43  	// create the queue item
    44  	item := &item{
    45  		message:   message,
    46  		priority:  int(priority),
    47  		timestamp: time.Now(),
    48  	}
    49  
    50  	// lock the underlying mutex
    51  	mq.cond.L.Lock()
    52  
    53  	// push message to the underlying priority queue
    54  	heap.Push(mq.pq, item)
    55  
    56  	// record metrics
    57  	mq.metrics.MessageAdded(item.priority)
    58  
    59  	// signal a waiting routine that a message is now available
    60  	mq.cond.Signal()
    61  
    62  	// unlock the underlying mutex
    63  	mq.cond.L.Unlock()
    64  
    65  	return nil
    66  }
    67  
    68  func (mq *MessageQueue) Remove() interface{} {
    69  	mq.cond.L.Lock()
    70  	defer mq.cond.L.Unlock()
    71  	for mq.pq.Len() == 0 {
    72  
    73  		// if the context has been canceled, don't wait
    74  		if err := mq.ctx.Err(); err != nil {
    75  			return nil
    76  		}
    77  
    78  		mq.cond.Wait()
    79  	}
    80  	item := heap.Pop(mq.pq).(*item)
    81  
    82  	// record metrics
    83  	mq.metrics.QueueDuration(time.Since(item.timestamp), item.priority)
    84  	mq.metrics.MessageRemoved(item.priority)
    85  
    86  	return item.message
    87  }
    88  
    89  func (mq *MessageQueue) Len() int {
    90  	mq.cond.L.Lock()
    91  	defer mq.cond.L.Unlock()
    92  	return mq.pq.Len()
    93  }
    94  
    95  func NewMessageQueue(ctx context.Context, priorityFunc MessagePriorityFunc, metrics module.NetworkInboundQueueMetrics) *MessageQueue {
    96  	var items = make([]*item, 0)
    97  	pq := priorityQueue(items)
    98  	mq := &MessageQueue{
    99  		pq:           &pq,
   100  		priorityFunc: priorityFunc,
   101  		ctx:          ctx,
   102  		metrics:      metrics,
   103  	}
   104  	m := sync.Mutex{}
   105  	mq.cond = sync.NewCond(&m)
   106  
   107  	// kick off a go routine to unblock queue readers on shutdown
   108  	go func() {
   109  		<-ctx.Done()
   110  		// unblock receive
   111  		mq.cond.Broadcast()
   112  	}()
   113  
   114  	return mq
   115  }