transport memory: add Send/Recv Timeout

Signed-off-by: Vasiliy Tolstov <v.tolstov@unistack.org>
This commit is contained in:
Василий Толстов 2019-08-03 15:39:44 +03:00
parent d250ac736f
commit e1709026e4

View File

@ -2,6 +2,7 @@
package memory
import (
"context"
"errors"
"fmt"
"math/rand"
@ -24,6 +25,10 @@ type memorySocket struct {
local string
remote string
// for send/recv transport.Timeout
timeout time.Duration
ctx context.Context
sync.RWMutex
}
@ -36,8 +41,10 @@ type memoryListener struct {
addr string
exit chan bool
conn chan *memorySocket
opts transport.ListenOptions
lopts transport.ListenOptions
topts transport.Options
sync.RWMutex
ctx context.Context
}
type memoryTransport struct {
@ -49,7 +56,17 @@ type memoryTransport struct {
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()
}
select {
case <-ctx.Done():
return ctx.Err()
case <-ms.exit:
return errors.New("connection closed")
case <-ms.lexit:
@ -71,7 +88,17 @@ func (ms *memorySocket) Remote() string {
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()
}
select {
case <-ctx.Done():
return ctx.Err()
case <-ms.exit:
return errors.New("connection closed")
case <-ms.lexit:
@ -122,6 +149,8 @@ func (m *memoryListener) Accept(fn func(transport.Socket)) error {
recv: c.send,
local: c.Remote(),
remote: c.Local(),
timeout: m.topts.Timeout,
ctx: m.topts.Context,
})
}
}
@ -149,6 +178,8 @@ func (m *memoryTransport) Dial(addr string, opts ...transport.DialOption) (trans
lexit: listener.exit,
local: addr,
remote: addr,
timeout: m.opts.Timeout,
ctx: m.opts.Context,
},
options,
}
@ -196,10 +227,12 @@ func (m *memoryTransport) Listen(addr string, opts ...transport.ListenOption) (t
}
listener := &memoryListener{
opts: options,
lopts: options,
topts: m.opts,
addr: addr,
conn: make(chan *memorySocket),
exit: make(chan bool),
ctx: m.opts.Context,
}
m.listeners[addr] = listener
@ -223,12 +256,18 @@ func (m *memoryTransport) String() string {
}
func NewTransport(opts ...transport.Option) transport.Transport {
rand.Seed(time.Now().UnixNano())
var options transport.Options
rand.Seed(time.Now().UnixNano())
for _, o := range opts {
o(&options)
}
if options.Context == nil {
options.Context = context.Background()
}
return &memoryTransport{
opts: options,
listeners: make(map[string]*memoryListener),