micro/transport/memory/memory.go

276 lines
4.9 KiB
Go
Raw Normal View History

2019-01-14 18:27:25 +03:00
// Package memory is an in-memory transport
package memory
2016-03-18 21:41:26 +03:00
import (
"context"
2016-03-18 21:41:26 +03:00
"errors"
"fmt"
"math/rand"
"net"
2016-03-18 21:41:26 +03:00
"sync"
"time"
2016-03-18 21:41:26 +03:00
"github.com/micro/go-micro/v3/transport"
maddr "github.com/micro/go-micro/v3/util/addr"
mnet "github.com/micro/go-micro/v3/util/net"
2016-03-18 21:41:26 +03:00
)
2019-01-14 18:27:25 +03:00
type memorySocket struct {
2016-03-18 21:41:26 +03:00
recv chan *transport.Message
send chan *transport.Message
// sock exit
exit chan bool
// listener exit
lexit chan bool
2018-11-14 22:41:13 +03:00
local string
remote string
// for send/recv transport.Timeout
timeout time.Duration
ctx context.Context
sync.RWMutex
2016-03-18 21:41:26 +03:00
}
2019-01-14 18:27:25 +03:00
type memoryClient struct {
*memorySocket
2016-03-18 21:41:26 +03:00
opts transport.DialOptions
}
2019-01-14 18:27:25 +03:00
type memoryListener struct {
addr string
exit chan bool
conn chan *memorySocket
lopts transport.ListenOptions
topts transport.Options
sync.RWMutex
ctx context.Context
2016-03-18 21:41:26 +03:00
}
2019-01-14 18:27:25 +03:00
type memoryTransport struct {
2016-03-18 21:41:26 +03:00
opts transport.Options
sync.RWMutex
2019-01-14 18:27:25 +03:00
listeners map[string]*memoryListener
2016-03-18 21:41:26 +03:00
}
2019-01-14 18:27:25 +03:00
func (ms *memorySocket) Recv(m *transport.Message) error {
ms.RLock()
defer ms.RUnlock()
ctx := ms.ctx
if ms.timeout > 0 {
var cancel context.CancelFunc
ctx, cancel = context.WithTimeout(ms.ctx, ms.timeout)
defer cancel()
}
2016-03-18 21:41:26 +03:00
select {
case <-ctx.Done():
return ctx.Err()
2016-03-18 21:41:26 +03:00
case <-ms.exit:
2016-03-18 22:23:44 +03:00
return errors.New("connection closed")
2016-03-18 21:41:26 +03:00
case <-ms.lexit:
2016-03-18 22:23:44 +03:00
return errors.New("server connection closed")
2016-03-18 21:41:26 +03:00
case cm := <-ms.recv:
*m = *cm
}
return nil
}
2019-01-14 18:27:25 +03:00
func (ms *memorySocket) Local() string {
2018-11-14 22:41:13 +03:00
return ms.local
}
2019-01-14 18:27:25 +03:00
func (ms *memorySocket) Remote() string {
2018-11-14 22:41:13 +03:00
return ms.remote
}
2019-01-14 18:27:25 +03:00
func (ms *memorySocket) Send(m *transport.Message) error {
ms.RLock()
defer ms.RUnlock()
ctx := ms.ctx
if ms.timeout > 0 {
var cancel context.CancelFunc
ctx, cancel = context.WithTimeout(ms.ctx, ms.timeout)
defer cancel()
}
2016-03-18 21:41:26 +03:00
select {
case <-ctx.Done():
return ctx.Err()
2016-03-18 21:41:26 +03:00
case <-ms.exit:
2016-03-18 22:23:44 +03:00
return errors.New("connection closed")
2016-03-18 21:41:26 +03:00
case <-ms.lexit:
2016-03-18 22:23:44 +03:00
return errors.New("server connection closed")
2016-03-18 21:41:26 +03:00
case ms.send <- m:
}
return nil
}
2019-01-14 18:27:25 +03:00
func (ms *memorySocket) Close() error {
ms.Lock()
defer ms.Unlock()
2016-06-03 13:08:58 +03:00
select {
case <-ms.exit:
return nil
default:
close(ms.exit)
}
2016-03-18 21:41:26 +03:00
return nil
}
2019-01-14 18:27:25 +03:00
func (m *memoryListener) Addr() string {
2016-03-18 21:41:26 +03:00
return m.addr
}
2019-01-14 18:27:25 +03:00
func (m *memoryListener) Close() error {
m.Lock()
defer m.Unlock()
2016-06-03 13:08:58 +03:00
select {
case <-m.exit:
return nil
default:
close(m.exit)
}
2016-03-18 21:41:26 +03:00
return nil
}
2019-01-14 18:27:25 +03:00
func (m *memoryListener) Accept(fn func(transport.Socket)) error {
2016-03-18 21:41:26 +03:00
for {
select {
case <-m.exit:
return nil
case c := <-m.conn:
2019-01-14 18:27:25 +03:00
go fn(&memorySocket{
lexit: c.lexit,
exit: c.exit,
send: c.recv,
recv: c.send,
local: c.Remote(),
remote: c.Local(),
timeout: m.topts.Timeout,
ctx: m.topts.Context,
2016-03-18 21:41:26 +03:00
})
}
}
}
2019-01-14 18:27:25 +03:00
func (m *memoryTransport) Dial(addr string, opts ...transport.DialOption) (transport.Client, error) {
m.RLock()
defer m.RUnlock()
2016-03-18 21:41:26 +03:00
listener, ok := m.listeners[addr]
if !ok {
return nil, errors.New("could not dial " + addr)
}
var options transport.DialOptions
for _, o := range opts {
o(&options)
}
2019-01-14 18:27:25 +03:00
client := &memoryClient{
&memorySocket{
send: make(chan *transport.Message),
recv: make(chan *transport.Message),
exit: make(chan bool),
lexit: listener.exit,
local: addr,
remote: addr,
timeout: m.opts.Timeout,
ctx: m.opts.Context,
2016-03-18 21:41:26 +03:00
},
options,
}
// pseudo connect
select {
case <-listener.exit:
return nil, errors.New("connection error")
2019-01-14 18:27:25 +03:00
case listener.conn <- client.memorySocket:
2016-03-18 21:41:26 +03:00
}
return client, nil
}
2019-01-14 18:27:25 +03:00
func (m *memoryTransport) Listen(addr string, opts ...transport.ListenOption) (transport.Listener, error) {
2016-03-18 21:41:26 +03:00
m.Lock()
defer m.Unlock()
var options transport.ListenOptions
for _, o := range opts {
o(&options)
}
host, port, err := net.SplitHostPort(addr)
if err != nil {
return nil, err
}
addr, err = maddr.Extract(host)
if err != nil {
return nil, err
}
// if zero port then randomly assign one
if len(port) > 0 && port == "0" {
i := rand.Intn(20000)
port = fmt.Sprintf("%d", 10000+i)
}
// set addr with port
addr = mnet.HostPort(addr, port)
if _, ok := m.listeners[addr]; ok {
return nil, errors.New("already listening on " + addr)
}
2019-01-14 18:27:25 +03:00
listener := &memoryListener{
lopts: options,
topts: m.opts,
addr: addr,
conn: make(chan *memorySocket),
exit: make(chan bool),
ctx: m.opts.Context,
2016-03-18 21:41:26 +03:00
}
m.listeners[addr] = listener
return listener, nil
}
2019-01-14 18:27:25 +03:00
func (m *memoryTransport) Init(opts ...transport.Option) error {
for _, o := range opts {
o(&m.opts)
}
return nil
}
2019-01-14 18:27:25 +03:00
func (m *memoryTransport) Options() transport.Options {
return m.opts
}
2019-01-14 18:27:25 +03:00
func (m *memoryTransport) String() string {
return "memory"
2016-03-18 21:41:26 +03:00
}
func NewTransport(opts ...transport.Option) transport.Transport {
var options transport.Options
rand.Seed(time.Now().UnixNano())
2016-03-18 21:41:26 +03:00
for _, o := range opts {
o(&options)
}
if options.Context == nil {
options.Context = context.Background()
}
2019-01-14 18:27:25 +03:00
return &memoryTransport{
2016-03-18 21:41:26 +03:00
opts: options,
2019-01-14 18:27:25 +03:00
listeners: make(map[string]*memoryListener),
2016-03-18 21:41:26 +03:00
}
}