// NewPullEngine creates an instance of a PullEngine with a certain sleep time
// between pull initiations
func NewPullEngine(participant PullAdapter, sleepTime time.Duration) *PullEngine {
engine := &PullEngine{
PullAdapter: participant,
stopFlag: int32(0),
state: util.NewSet(),
item2owners: make(map[string][]string),
peers2nonces: make(map[string]uint64),
nonces2peers: make(map[uint64]string),
acceptingDigests: int32(0),
acceptingResponses: int32(0),
incomingNONCES: util.NewSet(),
outgoingNONCES: util.NewSet(),
}
go func() {
for !engine.toDie() {
time.Sleep(sleepTime)
if engine.toDie() {
return
}
engine.initiatePull()
}
}()
return engine
}
func createPullInstance(endpoint string, peer2PullInst map[string]*pullInstance) *pullInstance {
inst := &pullInstance{
items: util.NewSet(),
stopChan: make(chan struct{}),
peer2PullInst: peer2PullInst,
self: discovery.NetworkMember{Endpoint: endpoint, Metadata: []byte{}, PKIid: []byte(endpoint)},
msgChan: make(chan *pullMsg, 10),
}
peer2PullInst[endpoint] = inst
conf := PullConfig{
MsgType: proto.PullMsgType_BlockMessage,
Channel: []byte(""),
Id: endpoint,
PeerCountToSelect: 3,
PullInterval: pullInterval,
Tag: proto.GossipMessage_EMPTY,
}
seqNumFromMsg := func(msg *proto.GossipMessage) string {
dataMsg := msg.GetDataMsg()
if dataMsg == nil {
return ""
}
if dataMsg.Payload == nil {
return ""
}
return fmt.Sprintf("%d", dataMsg.Payload.SeqNum)
}
blockConsumer := func(msg *proto.GossipMessage) {
inst.items.Add(msg.GetDataMsg().Payload.SeqNum)
}
inst.mediator = NewPullMediator(conf, inst, inst, seqNumFromMsg, blockConsumer)
go func() {
for {
select {
case <-inst.stopChan:
return
case msg := <-inst.msgChan:
inst.mediator.HandleMessage(msg)
}
}
}()
return inst
}
func TestRegisterMsgHook(t *testing.T) {
t.Parallel()
peer2pullInst := make(map[string]*pullInstance)
inst1 := createPullInstance("localhost:5611", peer2pullInst)
inst2 := createPullInstance("localhost:5612", peer2pullInst)
defer inst1.stop()
defer inst2.stop()
receivedMsgTypes := util.NewSet()
for _, msgType := range []PullMsgType{HelloMsgType, DigestMsgType, RequestMsgType, ResponseMsgType} {
mType := msgType
inst1.mediator.RegisterMsgHook(mType, func(_ []string, items []*proto.GossipMessage, _ comm.ReceivedMessage) {
receivedMsgTypes.Add(mType)
})
}
inst1.mediator.Add(dataMsg(1))
inst2.mediator.Add(dataMsg(2))
// Ensure all message types are received
waitUntilOrFail(t, func() bool { return len(receivedMsgTypes.ToArray()) == 4 })
}