go-gost/core
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add keepAlive option for udp Listener

Browse Source
This commit is contained in:
ginuerzh
2022-04-03 22:23:27 +08:00
parent fc1e6e8ff2
commit 6340d5198f
9 changed files with 109 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
common/util/udp/conn.go → common/net/udp/conn.go
View File

@ -9,8 +9,8 @@ import (
"github.com/go-gost/core/common/bufpool"
)
// Conn is a server side connection for UDP client peer, it implements net.Conn and net.PacketConn.
type Conn struct {
// conn is a server side connection for UDP client peer, it implements net.Conn and net.PacketConn.
type conn struct {
net.PacketConn
localAddr net.Addr
remoteAddr net.Addr
@ -18,19 +18,21 @@ type Conn struct {
idle int32 // indicate the connection is idle
closed chan struct{}
closeMutex sync.Mutex
keepAlive bool
}
func NewConn(c net.PacketConn, localAddr, remoteAddr net.Addr, queueSize int) *Conn {
return &Conn{
func newConn(c net.PacketConn, laddr, remoteAddr net.Addr, queueSize int, keepAlive bool) *conn {
return &conn{
PacketConn: c,
localAddr: localAddr,
localAddr: laddr,
remoteAddr: remoteAddr,
rc: make(chan []byte, queueSize),
closed: make(chan struct{}),
keepAlive: keepAlive,
}
}
func (c *Conn) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
func (c *conn) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
select {
case bb := <-c.rc:
n = copy(b, bb)
@ -47,16 +49,20 @@ func (c *Conn) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
return
}
func (c *Conn) Read(b []byte) (n int, err error) {
func (c *conn) Read(b []byte) (n int, err error) {
n, _, err = c.ReadFrom(b)
return
}
func (c *Conn) Write(b []byte) (n int, err error) {
return c.WriteTo(b, c.remoteAddr)
func (c *conn) Write(b []byte) (n int, err error) {
n, err = c.WriteTo(b, c.remoteAddr)
if !c.keepAlive {
c.Close()
}
return
}
func (c *Conn) Close() error {
func (c *conn) Close() error {
c.closeMutex.Lock()
defer c.closeMutex.Unlock()
@ -68,19 +74,19 @@ func (c *Conn) Close() error {
return nil
}
func (c *Conn) LocalAddr() net.Addr {
func (c *conn) LocalAddr() net.Addr {
return c.localAddr
}
func (c *Conn) RemoteAddr() net.Addr {
func (c *conn) RemoteAddr() net.Addr {
return c.remoteAddr
}
func (c *Conn) IsIdle() bool {
func (c *conn) IsIdle() bool {
return atomic.LoadInt32(&c.idle) > 0
}
func (c *Conn) SetIdle(idle bool) {
func (c *conn) SetIdle(idle bool) {
v := int32(0)
if idle {
v = 1
@ -88,7 +94,7 @@ func (c *Conn) SetIdle(idle bool) {
atomic.StoreInt32(&c.idle, v)
}
func (c *Conn) WriteQueue(b []byte) error {
func (c *conn) WriteQueue(b []byte) error {
select {
case c.rc <- b:
return nil

62
common/util/udp/listener.go → common/net/udp/listener.go
View File

@ -9,30 +9,37 @@ import (
"github.com/go-gost/core/logger"
)
type ListenConfig struct {
Addr net.Addr
Backlog int
ReadQueueSize int
ReadBufferSize int
TTL time.Duration
KeepAlive bool
Logger logger.Logger
}
type listener struct {
addr net.Addr
conn net.PacketConn
cqueue chan net.Conn
readQueueSize int
readBufferSize int
connPool *ConnPool
mux sync.Mutex
closed chan struct{}
errChan chan error
logger logger.Logger
conn net.PacketConn
cqueue chan net.Conn
connPool *connPool
mux sync.Mutex
closed chan struct{}
errChan chan error
config *ListenConfig
}
func NewListener(conn net.PacketConn, addr net.Addr, backlog, dataQueueSize, dataBufferSize int, ttl time.Duration, logger logger.Logger) net.Listener {
func NewListener(conn net.PacketConn, cfg *ListenConfig) net.Listener {
if cfg == nil {
cfg = &ListenConfig{}
}
ln := &listener{
conn: conn,
addr: addr,
cqueue: make(chan net.Conn, backlog),
connPool: NewConnPool(ttl).WithLogger(logger),
readQueueSize: dataQueueSize,
readBufferSize: dataBufferSize,
closed: make(chan struct{}),
errChan: make(chan error, 1),
logger: logger,
conn: conn,
cqueue: make(chan net.Conn, cfg.Backlog),
connPool: newConnPool(cfg.TTL).WithLogger(cfg.Logger),
closed: make(chan struct{}),
errChan: make(chan error, 1),
config: cfg,
}
go ln.listenLoop()
@ -61,7 +68,7 @@ func (ln *listener) listenLoop() {
default:
}
b := bufpool.Get(ln.readBufferSize)
b := bufpool.Get(ln.config.ReadBufferSize)
n, raddr, err := ln.conn.ReadFrom(*b)
if err != nil {
@ -77,13 +84,16 @@ func (ln *listener) listenLoop() {
}
if err := c.WriteQueue((*b)[:n]); err != nil {
ln.logger.Warn("data discarded: ", err)
ln.config.Logger.Warn("data discarded: ", err)
}
}
}
func (ln *listener) Addr() net.Addr {
return ln.addr
if ln.config.Addr != nil {
return ln.config.Addr
}
return ln.conn.LocalAddr()
}
func (ln *listener) Close() error {
@ -98,7 +108,7 @@ func (ln *listener) Close() error {
return nil
}
func (ln *listener) getConn(raddr net.Addr) *Conn {
func (ln *listener) getConn(raddr net.Addr) *conn {
ln.mux.Lock()
defer ln.mux.Unlock()
@ -107,14 +117,14 @@ func (ln *listener) getConn(raddr net.Addr) *Conn {
return c
}
c = NewConn(ln.conn, ln.addr, raddr, ln.readQueueSize)
c = newConn(ln.conn, ln.Addr(), raddr, ln.config.ReadQueueSize, ln.config.KeepAlive)
select {
case ln.cqueue <- c:
ln.connPool.Set(raddr.String(), c)
return c
default:
c.Close()
ln.logger.Warnf("connection queue is full, client %s discarded", raddr)
ln.config.Logger.Warnf("connection queue is full, client %s discarded", raddr)
return nil
}
}

24
common/util/udp/pool.go → common/net/udp/pool.go
View File

@ -7,15 +7,15 @@ import (
"github.com/go-gost/core/logger"
)
type ConnPool struct {
type connPool struct {
m sync.Map
ttl time.Duration
closed chan struct{}
logger logger.Logger
}
func NewConnPool(ttl time.Duration) *ConnPool {
p := &ConnPool{
func newConnPool(ttl time.Duration) *connPool {
p := &connPool{
ttl: ttl,
closed: make(chan struct{}),
}
@ -23,28 +23,28 @@ func NewConnPool(ttl time.Duration) *ConnPool {
return p
}
func (p *ConnPool) WithLogger(logger logger.Logger) *ConnPool {
func (p *connPool) WithLogger(logger logger.Logger) *connPool {
p.logger = logger
return p
}
func (p *ConnPool) Get(key any) (c *Conn, ok bool) {
func (p *connPool) Get(key any) (c *conn, ok bool) {
v, ok := p.m.Load(key)
if ok {
c, ok = v.(*Conn)
c, ok = v.(*conn)
}
return
}
func (p *ConnPool) Set(key any, c *Conn) {
func (p *connPool) Set(key any, c *conn) {
p.m.Store(key, c)
}
func (p *ConnPool) Delete(key any) {
func (p *connPool) Delete(key any) {
p.m.Delete(key)
}
func (p *ConnPool) Close() {
func (p *connPool) Close() {
select {
case <-p.closed:
return
@ -54,14 +54,14 @@ func (p *ConnPool) Close() {
close(p.closed)
p.m.Range(func(k, v any) bool {
if c, ok := v.(*Conn); ok && c != nil {
if c, ok := v.(*conn); ok && c != nil {
c.Close()
}
return true
})
}
func (p *ConnPool) idleCheck() {
func (p *connPool) idleCheck() {
ticker := time.NewTicker(p.ttl)
defer ticker.Stop()
@ -71,7 +71,7 @@ func (p *ConnPool) idleCheck() {
size := 0
idles := 0
p.m.Range(func(key, value any) bool {
c, ok := value.(*Conn)
c, ok := value.(*conn)
if !ok || c == nil {
p.Delete(key)
return true

4
common/net/udp.go → common/net/udp/udp.go
View File

@ -1,4 +1,4 @@
package net
package udp
import (
"io"
@ -6,7 +6,7 @@ import (
"syscall"
)
type UDPConn interface {
type Conn interface {
net.PacketConn
io.Reader
io.Writer