github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/network/stub/buffer.go (about)

     1  package stub
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/onflow/flow-go/model/flow"
     7  	"github.com/onflow/flow-go/network/channels"
     8  )
     9  
    10  // PendingMessage is a pending message to be sent
    11  type PendingMessage struct {
    12  	// The sender node id
    13  	From    flow.Identifier
    14  	Channel channels.Channel
    15  	Event   interface{}
    16  	// The id of the receiver nodes
    17  	TargetIDs []flow.Identifier
    18  }
    19  
    20  // Buffer buffers all the pending messages to be sent over the mock network from one node to a list of nodes
    21  type Buffer struct {
    22  	sync.Mutex
    23  	pending []*PendingMessage
    24  }
    25  
    26  // NewBuffer initialize the Buffer
    27  func NewBuffer() *Buffer {
    28  	return &Buffer{
    29  		pending: make([]*PendingMessage, 0),
    30  	}
    31  }
    32  
    33  // Save stores a pending message to the buffer
    34  func (b *Buffer) Save(m *PendingMessage) {
    35  	b.Lock()
    36  	defer b.Unlock()
    37  
    38  	b.pending = append(b.pending, m)
    39  }
    40  
    41  // DeliverRecursive recursively delivers all pending messages using the provided
    42  // sendOne method until the buffer is empty. If sendOne does not deliver the
    43  // message, it is permanently dropped.
    44  func (b *Buffer) DeliverRecursive(sendOne func(*PendingMessage)) {
    45  	for {
    46  		// get all pending messages, and clear the buffer
    47  		messages := b.takeAll()
    48  
    49  		// This check is necessary to exit the endless for loop
    50  		if len(messages) == 0 {
    51  			return
    52  		}
    53  
    54  		for _, msg := range messages {
    55  			sendOne(msg)
    56  		}
    57  	}
    58  }
    59  
    60  // Deliver delivers all pending messages currently in the buffer using the
    61  // provided sendOne method. If sendOne returns false, the message was not sent
    62  // and will remain in the buffer.
    63  func (b *Buffer) Deliver(sendOne func(*PendingMessage) bool) {
    64  
    65  	messages := b.takeAll()
    66  	var unsent []*PendingMessage
    67  
    68  	for _, msg := range messages {
    69  		ok := sendOne(msg)
    70  		if !ok {
    71  			unsent = append(unsent, msg)
    72  		}
    73  	}
    74  
    75  	b.Lock()
    76  	defer b.Unlock()
    77  
    78  	// add the unsent messages back to the buffer
    79  	b.pending = append(unsent, b.pending...)
    80  }
    81  
    82  // takeAll takes all pending messages from the buffer and empties the buffer.
    83  func (b *Buffer) takeAll() []*PendingMessage {
    84  	b.Lock()
    85  	defer b.Unlock()
    86  
    87  	toSend := b.pending[:]
    88  	b.pending = nil
    89  
    90  	return toSend
    91  }