diff --git a/internal/consensus/state.go b/internal/consensus/state.go index c8f445e1a..b42190eeb 100644 --- a/internal/consensus/state.go +++ b/internal/consensus/state.go @@ -567,23 +567,27 @@ func (cs *State) AddVote(ctx context.Context, vote *types.Vote, peerID types.Nod // TODO: wait for event?! } +func (cs *State) getMessageChannel(peerID types.NodeID) chan msgInfo { + cs.mtx.RLock() + defer cs.mtx.RUnlock() + + if peerID == "" { + return cs.internalMsgQueue + } + + return cs.peerMsgQueue +} + // SetProposal inputs a proposal. func (cs *State) SetProposal(ctx context.Context, proposal *types.Proposal, peerID types.NodeID) error { + ch := cs.getMessageChannel(peerID) + msg := msgInfo{&ProposalMessage{proposal}, peerID, tmtime.Now()} - if peerID == "" { - select { - case <-ctx.Done(): - return ctx.Err() - case cs.internalMsgQueue <- msgInfo{&ProposalMessage{proposal}, "", tmtime.Now()}: - return nil - } - } else { - select { - case <-ctx.Done(): - return ctx.Err() - case cs.peerMsgQueue <- msgInfo{&ProposalMessage{proposal}, peerID, tmtime.Now()}: - return nil - } + select { + case <-ctx.Done(): + return ctx.Err() + case ch <- msg: + return nil } // TODO: wait for event?! @@ -591,20 +595,14 @@ func (cs *State) SetProposal(ctx context.Context, proposal *types.Proposal, peer // AddProposalBlockPart inputs a part of the proposal block. func (cs *State) AddProposalBlockPart(ctx context.Context, height int64, round int32, part *types.Part, peerID types.NodeID) error { - if peerID == "" { - select { - case <-ctx.Done(): - return ctx.Err() - case cs.internalMsgQueue <- msgInfo{&BlockPartMessage{height, round, part}, "", tmtime.Now()}: - return nil - } - } else { - select { - case <-ctx.Done(): - return ctx.Err() - case cs.peerMsgQueue <- msgInfo{&BlockPartMessage{height, round, part}, peerID, tmtime.Now()}: - return nil - } + ch := cs.getMessageChannel(peerID) + msg := msgInfo{&BlockPartMessage{height, round, part}, peerID, tmtime.Now()} + + select { + case <-ctx.Done(): + return ctx.Err() + case ch <- msg: + return nil } // TODO: wait for event?!