github.com/koko1123/flow-go-1@v0.29.6/engine/common/synchronization/request_heap.go (about)

     1  package synchronization
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/koko1123/flow-go-1/engine"
     7  	"github.com/koko1123/flow-go-1/model/flow"
     8  )
     9  
    10  // RequestHeap is a special structure that implements engine.MessageStore interface and
    11  // indexes requests by originator. If request will be sent by same originator then it will replace the old one.
    12  // Comparing to default FIFO queue this one can contain MAX one request for origin ID.
    13  // Getting value from queue as well as ejecting is pseudo-random.
    14  type RequestHeap struct {
    15  	lock     sync.Mutex
    16  	limit    uint
    17  	requests map[flow.Identifier]*engine.Message
    18  }
    19  
    20  func NewRequestHeap(limit uint) *RequestHeap {
    21  	return &RequestHeap{
    22  		limit:    limit,
    23  		requests: make(map[flow.Identifier]*engine.Message),
    24  	}
    25  }
    26  
    27  // Put stores message into requests map using OriginID as key.
    28  // Returns always true
    29  func (q *RequestHeap) Put(message *engine.Message) bool {
    30  	q.lock.Lock()
    31  	defer q.lock.Unlock()
    32  	// first try to eject if we are at max capacity, we need to do this way
    33  	// to prevent a situation where just inserted item gets ejected
    34  	if _, found := q.requests[message.OriginID]; !found {
    35  		// if no message from the origin is stored, make sure we have room to store the new message:
    36  		q.reduce()
    37  	}
    38  	// at this point we can be sure that there is at least one slot
    39  	q.requests[message.OriginID] = message
    40  	return true
    41  }
    42  
    43  // Get returns pseudo-random element from request storage using go map properties.
    44  func (q *RequestHeap) Get() (*engine.Message, bool) {
    45  	q.lock.Lock()
    46  	defer q.lock.Unlock()
    47  
    48  	if len(q.requests) == 0 {
    49  		return nil, false
    50  	}
    51  
    52  	// pick arbitrary element using go map randomness property
    53  	var originID flow.Identifier
    54  	var msg *engine.Message
    55  	for originID, msg = range q.requests {
    56  		break
    57  	}
    58  	delete(q.requests, originID)
    59  
    60  	return msg, true
    61  }
    62  
    63  // reduce will reduce the size of the kept entities until we are within the
    64  // configured memory pool size limit. If called on max capacity will eject at least one element.
    65  func (q *RequestHeap) reduce() {
    66  	for overCapacity := len(q.requests) - int(q.limit); overCapacity >= 0; overCapacity-- {
    67  		for originID := range q.requests { // pick first element using go map randomness property
    68  			delete(q.requests, originID)
    69  			break
    70  		}
    71  	}
    72  }