forked from ava-labs/avalanchego
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Factor out message queue from peer implementation (ava-labs#1484)
- Loading branch information
1 parent
2fd7d90
commit e5f9ae4
Showing
7 changed files
with
457 additions
and
168 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,284 @@ | ||
| // Copyright (C) 2019-2021, Ava Labs, Inc. All rights reserved. | ||
| // See the file LICENSE for licensing terms. | ||
|
|
||
| package peer | ||
|
|
||
| import ( | ||
| "context" | ||
| "sync" | ||
|
|
||
| "github.com/ava-labs/avalanchego/ids" | ||
| "github.com/ava-labs/avalanchego/message" | ||
| "github.com/ava-labs/avalanchego/network/throttling" | ||
| "github.com/ava-labs/avalanchego/utils/logging" | ||
| ) | ||
|
|
||
| var ( | ||
| _ MessageQueue = &throttledMessageQueue{} | ||
| _ MessageQueue = &blockingMessageQueue{} | ||
| ) | ||
|
|
||
| type SendFailedCallback interface { | ||
| SendFailed(message.OutboundMessage) | ||
| } | ||
|
|
||
| type SendFailedFunc func(message.OutboundMessage) | ||
|
|
||
| func (f SendFailedFunc) SendFailed(msg message.OutboundMessage) { f(msg) } | ||
|
|
||
| type MessageQueue interface { | ||
| // Push attempts to add the message to the queue. If the context is | ||
| // canceled, then pushing the message will return `false` and the message | ||
| // will not be added to the queue. | ||
| Push(ctx context.Context, msg message.OutboundMessage) bool | ||
|
|
||
| // Pop blocks until a message is available and then returns the message. If | ||
| // the queue is closed, then `false` is returned. | ||
| Pop() (message.OutboundMessage, bool) | ||
|
|
||
| // PopNow attempts to return a message without blocking. If a message is not | ||
| // available or the queue is closed, then `false` is returned. | ||
| PopNow() (message.OutboundMessage, bool) | ||
|
|
||
| // Close empties the queue and prevents further messages from being pushed | ||
| // onto it. After calling close once, future calls to close will do nothing. | ||
| Close() | ||
| } | ||
|
|
||
| type throttledMessageQueue struct { | ||
| onFailed SendFailedCallback | ||
| // [id] of the peer we're sending messages to | ||
| id ids.NodeID | ||
| log logging.Logger | ||
| outboundMsgThrottler throttling.OutboundMsgThrottler | ||
|
|
||
| // Signalled when a message is added to the queue and when Close() is | ||
| // called. | ||
| cond *sync.Cond | ||
|
|
||
| // closed flags whether the send queue has been closed. | ||
| // [cond.L] must be held while accessing [closed]. | ||
| closed bool | ||
|
|
||
| // queue of the messages | ||
| // [cond.L] must be held while accessing [queue]. | ||
| queue []message.OutboundMessage | ||
| } | ||
|
|
||
| func NewThrottledMessageQueue( | ||
| onFailed SendFailedCallback, | ||
| id ids.NodeID, | ||
| log logging.Logger, | ||
| outboundMsgThrottler throttling.OutboundMsgThrottler, | ||
| ) MessageQueue { | ||
| return &throttledMessageQueue{ | ||
| onFailed: onFailed, | ||
| id: id, | ||
| log: log, | ||
| outboundMsgThrottler: outboundMsgThrottler, | ||
|
|
||
| cond: sync.NewCond(&sync.Mutex{}), | ||
| } | ||
| } | ||
|
|
||
| func (q *throttledMessageQueue) Push(ctx context.Context, msg message.OutboundMessage) bool { | ||
| if err := ctx.Err(); err != nil { | ||
| q.log.Debug( | ||
| "dropping %s message to %s due to a context error: %s", | ||
| msg.Op(), q.id, err, | ||
| ) | ||
| q.onFailed.SendFailed(msg) | ||
| return false | ||
| } | ||
|
|
||
| // Acquire space on the outbound message queue, or drop [msg] if we can't. | ||
| if !q.outboundMsgThrottler.Acquire(msg, q.id) { | ||
| q.log.Debug( | ||
| "dropping %s message to %s due to rate-limiting", | ||
| msg.Op(), q.id, | ||
| ) | ||
| q.onFailed.SendFailed(msg) | ||
| return false | ||
| } | ||
|
|
||
| // Invariant: must call q.outboundMsgThrottler.Release(msg, q.id) when [msg] | ||
| // is popped or, if this queue closes before [msg] is popped, when this | ||
| // queue closes. | ||
|
|
||
| q.cond.L.Lock() | ||
| defer q.cond.L.Unlock() | ||
|
|
||
| if q.closed { | ||
| q.log.Debug( | ||
| "dropping %s message to %s due to a closed queue", | ||
| msg.Op(), q.id, | ||
| ) | ||
| q.outboundMsgThrottler.Release(msg, q.id) | ||
| q.onFailed.SendFailed(msg) | ||
| return false | ||
| } | ||
|
|
||
| q.queue = append(q.queue, msg) | ||
| q.cond.Signal() | ||
| return true | ||
| } | ||
|
|
||
| func (q *throttledMessageQueue) Pop() (message.OutboundMessage, bool) { | ||
| q.cond.L.Lock() | ||
| defer q.cond.L.Unlock() | ||
|
|
||
| for { | ||
| if q.closed { | ||
| return nil, false | ||
| } | ||
| if len(q.queue) > 0 { | ||
| // There is a message | ||
| break | ||
| } | ||
| // Wait until there is a message | ||
| q.cond.Wait() | ||
| } | ||
|
|
||
| return q.pop(), true | ||
| } | ||
|
|
||
| func (q *throttledMessageQueue) PopNow() (message.OutboundMessage, bool) { | ||
| q.cond.L.Lock() | ||
| defer q.cond.L.Unlock() | ||
|
|
||
| if len(q.queue) == 0 { | ||
| // There isn't a message | ||
| return nil, false | ||
| } | ||
|
|
||
| return q.pop(), true | ||
| } | ||
|
|
||
| func (q *throttledMessageQueue) pop() message.OutboundMessage { | ||
| msg := q.queue[0] | ||
| q.queue[0] = nil | ||
| q.queue = q.queue[1:] | ||
|
|
||
| q.outboundMsgThrottler.Release(msg, q.id) | ||
| return msg | ||
| } | ||
|
|
||
| func (q *throttledMessageQueue) Close() { | ||
| q.cond.L.Lock() | ||
| defer q.cond.L.Unlock() | ||
|
|
||
| q.closed = true | ||
|
|
||
| for _, msg := range q.queue { | ||
| q.outboundMsgThrottler.Release(msg, q.id) | ||
| q.onFailed.SendFailed(msg) | ||
| } | ||
| q.queue = nil | ||
|
|
||
| q.cond.Broadcast() | ||
| } | ||
|
|
||
| type blockingMessageQueue struct { | ||
| onFailed SendFailedCallback | ||
| log logging.Logger | ||
|
|
||
| closeOnce sync.Once | ||
| closingLock sync.RWMutex | ||
| closing chan struct{} | ||
|
|
||
| // queue of the messages | ||
| queue chan message.OutboundMessage | ||
| } | ||
|
|
||
| func NewBlockingMessageQueue( | ||
| onFailed SendFailedCallback, | ||
| log logging.Logger, | ||
| bufferSize int, | ||
| ) MessageQueue { | ||
| return &blockingMessageQueue{ | ||
| onFailed: onFailed, | ||
| log: log, | ||
|
|
||
| closing: make(chan struct{}), | ||
| queue: make(chan message.OutboundMessage, bufferSize), | ||
| } | ||
| } | ||
|
|
||
| func (q *blockingMessageQueue) Push(ctx context.Context, msg message.OutboundMessage) bool { | ||
| q.closingLock.RLock() | ||
| defer q.closingLock.RUnlock() | ||
|
|
||
| ctxDone := ctx.Done() | ||
| select { | ||
| case <-q.closing: | ||
| q.log.Debug( | ||
| "dropping %s message due to a closed queue", | ||
| msg.Op(), | ||
| ) | ||
| q.onFailed.SendFailed(msg) | ||
| return false | ||
| case <-ctxDone: | ||
| q.log.Debug( | ||
| "dropping %s message due to a cancelled context", | ||
| msg.Op(), | ||
| ) | ||
| q.onFailed.SendFailed(msg) | ||
| return false | ||
| default: | ||
| } | ||
|
|
||
| select { | ||
| case q.queue <- msg: | ||
| return true | ||
| case <-ctxDone: | ||
| q.log.Debug( | ||
| "dropping %s message due to a cancelled context", | ||
| msg.Op(), | ||
| ) | ||
| q.onFailed.SendFailed(msg) | ||
| return false | ||
| case <-q.closing: | ||
| q.log.Debug( | ||
| "dropping %s message due to a closed queue", | ||
| msg.Op(), | ||
| ) | ||
| q.onFailed.SendFailed(msg) | ||
| return false | ||
| } | ||
| } | ||
|
|
||
| func (q *blockingMessageQueue) Pop() (message.OutboundMessage, bool) { | ||
| select { | ||
| case msg := <-q.queue: | ||
| return msg, true | ||
| case <-q.closing: | ||
| return nil, false | ||
| } | ||
| } | ||
|
|
||
| func (q *blockingMessageQueue) PopNow() (message.OutboundMessage, bool) { | ||
| select { | ||
| case msg := <-q.queue: | ||
| return msg, true | ||
| default: | ||
| return nil, false | ||
| } | ||
| } | ||
|
|
||
| func (q *blockingMessageQueue) Close() { | ||
| q.closeOnce.Do(func() { | ||
| close(q.closing) | ||
|
|
||
| q.closingLock.Lock() | ||
| defer q.closingLock.Unlock() | ||
|
|
||
| for { | ||
| select { | ||
| case msg := <-q.queue: | ||
| q.onFailed.SendFailed(msg) | ||
| default: | ||
| return | ||
| } | ||
| } | ||
| }) | ||
| } |
Oops, something went wrong.