Exemplo n.º 1
0
func (re *RouterElement) run(doRouting func(*RouterElement, *Peer, [][]byte) (*Peer, [][]byte, error), controlIn, controlOut *zmq.Socket) {
	var wg sync.WaitGroup
	controlChanInternal := make(chan *controlMessage, 1)
	wg.Add(1)

	// Start a goroutine to copy from the control channel to the internal
	// control channel and wake up the control socket pair on each message.
	// Quit the goroutine when the control channel itself is closed
	go func() {
		defer func() {
			controlIn.SetLinger(0)
			controlIn.Close()
			wg.Done()
		}()
		msg := [][]byte{[]byte{0}}
		for {
			select {
			case m, ok := <-re.controlChan:
				if !ok {
					close(controlChanInternal)
					// send a message to the control panel to wake it up
					controlIn.SendMessage(msg)
					// ensure we get a message back before closing it
					controlIn.RecvMessageBytes(0)
					return
				}
				controlChanInternal <- m
				controlIn.SendMessageDontwait(msg)
			}
		}
	}()

	defer func() {
		wg.Wait()
		controlOut.SetLinger(0)
		controlOut.Close()
		// Close() should already remove all the peers, so this is just belt & braces
		for _, p := range re.peers {
			p.sock.SetLinger(0)
			p.sock.Close()
			close(p.buffer)
		}
		re.wg.Done()
	}()

	for {
		poller := zmq.NewPoller()
		pollSource := make([]*Peer, len(re.peers), len(re.peers))
		blockInput := false
		i := 0
		for _, p := range re.peers {
			pollSource[i] = p
			i++
			if p.xmit == nil {
				select {
				case p.xmit = <-p.buffer:
				default:
				}
			}
			buffered := len(p.buffer)
			if p.xmit != nil {
				buffered++
			}
			if buffered >= cap(p.buffer) {
				log.Printf("Blocking input as %s cannot take more data", p.Name)
				blockInput = true
			}
		}
		for _, p := range pollSource {
			flags := zmq.State(0)
			if !blockInput {
				flags |= zmq.POLLIN
			}
			if p.xmit != nil {
				flags |= zmq.POLLOUT
			}
			poller.Add(p.sock, flags)
		}
		idxControl := poller.Add(controlOut, zmq.POLLIN)

		if polled, err := poller.PollAll(-1); err != nil {
			log.Printf("Cannot poll: %v", err)
		} else {
			for i, p := range pollSource {
				if (polled[i].Events&zmq.POLLOUT != 0) && (p.xmit != nil) {
					if _, err := p.sock.SendMessageDontwait(p.xmit); err != nil {
						log.Printf("Peer %s cannot send message: %v", p.Name, err)
					}
					p.xmit = nil
				}
			}
			for i, p := range pollSource {
				if polled[i].Events&zmq.POLLIN != 0 {
					if recv, err := p.sock.RecvMessageBytes(0); err != nil {
						log.Printf("Peer %s cannot receive message: %v", p.Name, err)
					} else {
						if dest, fwd, err := doRouting(re, p, recv); err != nil {
							log.Printf("Peer %s cannot route message: %v", p.Name, err)
						} else {
							if dest != nil && fwd != nil {
								dest.buffer <- fwd
							}
						}
					}
				}
			}
			if polled[idxControl].Events&zmq.POLLIN != 0 {
				// read from the control socket and discard. We don't care
				// if this errors as it's really only to make us read from our
				// own control socket
				controlOut.RecvMessageBytes(0)
				select {
				case msg, ok := <-controlChanInternal:
					if !ok {
						// control channel is closed
						// send a dummy message to handshake
						msg := [][]byte{[]byte{0}}
						controlOut.SendMessage(msg)
						return
					}
					msg.reply <- re.processControlMessage(msg)
				default:
				}
			}
		}
	}
}