micro/util/socket/socket.go

132 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/micro/go-micro/transport"
)
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
}
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
}
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
}
2019-08-15 22:08:49 +03:00
func (s *Socket) Remote() string {
2019-08-15 17:09:56 +03:00
return s.remote
}
2019-08-15 22:08:49 +03:00
func (s *Socket) Local() string {
2019-08-15 17:09:56 +03:00
return s.local
}
2019-08-15 22:08:49 +03:00
func (s *Socket) Send(m *transport.Message) error {
2019-08-15 17:09:56 +03:00
// make copy
msg := &transport.Message{
Header: make(map[string]string),
2019-08-15 22:54:09 +03:00
Body: make([]byte, len(m.Body)),
2019-08-15 17:09:56 +03:00
}
2019-08-15 22:08:49 +03:00
// copy headers
2019-08-15 17:09:56 +03:00
for k, v := range m.Header {
msg.Header[k] = v
}
2019-08-15 22:08:49 +03:00
// copy body
copy(msg.Body, m.Body)
2019-08-15 17:09:56 +03:00
// send a message
select {
case s.send <- msg:
case <-s.closed:
return io.EOF
}
return nil
}
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),
}
}