// Add registers a socket for sending and/or receiving. The caller can't
// access the socket directly after this. The send channel (if any) should be
// closed by the caller.
func (io *IO) Add(s *zmq.Socket, send <-chan Data, recv chan<- Data) (err error) {
fd, err := s.GetFd()
if err != nil {
return
}
w := newWorker()
io.lock.Lock()
io.workers[int32(fd)] = w
io.lock.Unlock()
defer func() {
if err != nil {
io.lock.Lock()
delete(io.workers, int32(fd))
io.lock.Unlock()
}
}()
e := &syscall.EpollEvent{
Events: syscall.EPOLLIN | syscall.EPOLLET&0xffffffff,
Fd: int32(fd),
}
if err = syscall.EpollCtl(io.epollFd, syscall.EPOLL_CTL_ADD, fd, e); err != nil {
return
}
state, err := s.GetEvents()
if err != nil {
syscall.EpollCtl(io.epollFd, syscall.EPOLL_CTL_DEL, fd, nil)
return
}
go func() {
defer s.Close()
defer syscall.EpollCtl(io.epollFd, syscall.EPOLL_CTL_DEL, fd, nil)
w.socketLoop(s, send, recv, state)
}()
return
}
func (w *worker) socketLoop(s *zmq.Socket, send <-chan Data, recv chan<- Data, state zmq.State) {
if recv != nil {
defer close(recv)
}
var (
sendBuf Data
sendPending bool
recvBuf Data
recvPending bool
)
for {
var (
err error
sendActive <-chan Data
recvActive chan<- Data
)
if !sendPending {
sendActive = send
}
if recvPending {
recvActive = recv
}
select {
case <-w.notifier:
const fullState = zmq.POLLIN | zmq.POLLOUT
if state&fullState != fullState {
if state, err = s.GetEvents(); err != nil {
handleGeneralError(err)
return
}
}
case <-w.closer:
return
case sendBuf, sendPending = <-sendActive:
if !sendPending {
send = nil
}
case recvActive <- recvBuf:
recvPending = false
recvBuf.Bytes = nil
}
for {
loop := false
if sendPending && state&zmq.POLLOUT != 0 {
flags := zmq.DONTWAIT
if sendBuf.More {
flags |= zmq.SNDMORE
}
if _, err = s.SendBytes(sendBuf.Bytes, flags); err == nil {
sendPending = false
sendBuf.Bytes = nil
loop = true
} else if !handleIOError(err) {
return
}
if state, err = s.GetEvents(); err != nil {
handleGeneralError(err)
return
}
}
if !recvPending && state&zmq.POLLIN != 0 {
if data, err := s.RecvBytes(zmq.DONTWAIT); err == nil {
if more, err := s.GetRcvmore(); err == nil {
recvBuf.Bytes = data
recvBuf.More = more
recvPending = true
loop = true
} else {
handleGeneralError(err)
return
}
} else if !handleIOError(err) {
return
}
if state, err = s.GetEvents(); err != nil {
handleGeneralError(err)
return
}
}
if !loop {
break
}
}
}
}