Merge pull request #640 from unistack-org/deadline
transport memory: add Send/Recv Timeout
This commit is contained in:
commit
be229438bc
@ -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),
|
||||
|
Loading…
Reference in New Issue
Block a user