micro/util/socket/socket.go

124 lines
2.3 KiB
Go
Raw Normal View History

2019-08-15 17:09:56 +03:00
// Package socket provides a pseudo socket
package socket
import (
"io"
"github.com/unistack-org/micro/v3/network/transport"
2019-08-15 17:09:56 +03:00
)
2019-08-15 22:08:49 +03:00
// Socket is our pseudo socket for transport.Socket
type Socket struct {
2019-11-27 20:12:07 +03:00
id string
2019-08-15 17:09:56 +03:00
// closed
closed chan bool
// remote addr
remote string
// local addr
local string
// send chan
send chan *transport.Message
// recv chan
recv chan *transport.Message
}
// SetLocal sets the local addr
2019-08-15 22:08:49 +03:00
func (s *Socket) SetLocal(l string) {
2019-08-15 17:09:56 +03:00
s.local = l
}
// SetRemote sets the remote addr
2019-08-15 22:08:49 +03:00
func (s *Socket) SetRemote(r string) {
2019-08-15 17:09:56 +03:00
s.remote = r
}
// Accept passes a message to the socket which will be processed by the call to Recv
2019-08-15 22:08:49 +03:00
func (s *Socket) Accept(m *transport.Message) error {
2019-08-15 17:09:56 +03:00
select {
case s.recv <- m:
return nil
case <-s.closed:
return io.EOF
2019-08-15 17:09:56 +03:00
}
}
// Process takes the next message off the send queue created by a call to Send
2019-08-15 22:08:49 +03:00
func (s *Socket) Process(m *transport.Message) error {
2019-08-15 17:09:56 +03:00
select {
case msg := <-s.send:
*m = *msg
case <-s.closed:
// see if we need to drain
select {
case msg := <-s.send:
*m = *msg
return nil
default:
return io.EOF
}
2019-08-15 17:09:56 +03:00
}
return nil
}
// Remote returns remote addr
2019-08-15 22:08:49 +03:00
func (s *Socket) Remote() string {
2019-08-15 17:09:56 +03:00
return s.remote
}
// Local returns local addr
2019-08-15 22:08:49 +03:00
func (s *Socket) Local() string {
2019-08-15 17:09:56 +03:00
return s.local
}
// Send message by via transport
2019-08-15 22:08:49 +03:00
func (s *Socket) Send(m *transport.Message) error {
2019-08-15 17:09:56 +03:00
// send a message
select {
case s.send <- m:
2019-08-15 17:09:56 +03:00
case <-s.closed:
return io.EOF
}
return nil
}
// Recv message from transport
2019-08-15 22:08:49 +03:00
func (s *Socket) Recv(m *transport.Message) error {
2019-08-15 17:09:56 +03:00
// receive a message
select {
case msg := <-s.recv:
// set message
*m = *msg
case <-s.closed:
return io.EOF
}
// return nil
return nil
}
// Close closes the socket
2019-08-15 22:08:49 +03:00
func (s *Socket) Close() error {
2019-08-15 17:09:56 +03:00
select {
case <-s.closed:
// no op
default:
close(s.closed)
}
return nil
}
// New returns a new pseudo socket which can be used in the place of a transport socket.
// Messages are sent to the socket via Accept and receives from the socket via Process.
// SetLocal/SetRemote should be called before using the socket.
2019-11-27 20:12:07 +03:00
func New(id string) *Socket {
2019-08-15 22:08:49 +03:00
return &Socket{
2019-11-27 20:12:07 +03:00
id: id,
2019-08-15 17:09:56 +03:00
closed: make(chan bool),
local: "local",
remote: "remote",
send: make(chan *transport.Message, 128),
recv: make(chan *transport.Message, 128),
}
}