From 7f581cb668b19f4114591d1ff73f33cc50e764ab Mon Sep 17 00:00:00 2001
From: ginuerzh <ginuerzh@gmail.com>
Date: Tue, 7 Nov 2023 23:09:07 +0800
Subject: [PATCH] improve udp listener

---
 common/net/udp/listener.go | 14 ++++++++------
 common/net/udp/pool.go     | 15 +++++++++++++++
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/common/net/udp/listener.go b/common/net/udp/listener.go
index 3c31edb..00216fa 100644
--- a/common/net/udp/listener.go
+++ b/common/net/udp/listener.go
@@ -33,12 +33,14 @@ func NewListener(conn net.PacketConn, cfg *ListenConfig) net.Listener {
 	}
 
 	ln := &listener{
-		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,
+		conn:    conn,
+		cqueue:  make(chan net.Conn, cfg.Backlog),
+		closed:  make(chan struct{}),
+		errChan: make(chan error, 1),
+		config:  cfg,
+	}
+	if cfg.KeepAlive {
+		ln.connPool = newConnPool(cfg.TTL).WithLogger(cfg.Logger)
 	}
 	go ln.listenLoop()
 
diff --git a/common/net/udp/pool.go b/common/net/udp/pool.go
index 1d69578..232d03d 100644
--- a/common/net/udp/pool.go
+++ b/common/net/udp/pool.go
@@ -29,6 +29,10 @@ func (p *connPool) WithLogger(logger logger.Logger) *connPool {
 }
 
 func (p *connPool) Get(key any) (c *conn, ok bool) {
+	if p == nil {
+		return
+	}
+
 	v, ok := p.m.Load(key)
 	if ok {
 		c, ok = v.(*conn)
@@ -37,14 +41,25 @@ func (p *connPool) Get(key any) (c *conn, ok bool) {
 }
 
 func (p *connPool) Set(key any, c *conn) {
+	if p == nil {
+		return
+	}
+
 	p.m.Store(key, c)
 }
 
 func (p *connPool) Delete(key any) {
+	if p == nil {
+		return
+	}
 	p.m.Delete(key)
 }
 
 func (p *connPool) Close() {
+	if p == nil {
+		return
+	}
+
 	select {
 	case <-p.closed:
 		return