refactor(doq): simplify DoQ connection pool implementation

Replace the map-based pool and refCount bookkeeping with a channel-based
pool. Drop the closed state, per-connection address tracking, and extra
mutexes so the pool relies on the channel for concurrency and lifecycle,
matching the approach used in the DoT pool.
This commit is contained in:
Cuong Manh Le
2026-01-28 23:50:53 +07:00
committed by Cuong Manh Le
parent 60dd366cc4
commit f1b93c81bc
+50 -97
View File
@@ -9,7 +9,6 @@ import (
"io" "io"
"net" "net"
"runtime" "runtime"
"sync"
"time" "time"
"github.com/miekg/dns" "github.com/miekg/dns"
@@ -48,22 +47,19 @@ func (r *doqResolver) Resolve(ctx context.Context, msg *dns.Msg) (*dns.Msg, erro
return answer, err return answer, err
} }
// doqConnPool manages a pool of QUIC connections for DoQ queries. const doqPoolSize = 16
// doqConnPool manages a pool of QUIC connections for DoQ queries using a buffered channel.
type doqConnPool struct { type doqConnPool struct {
uc *UpstreamConfig uc *UpstreamConfig
addrs []string addrs []string
port string port string
tlsConfig *tls.Config tlsConfig *tls.Config
mu sync.RWMutex conns chan *doqConn
conns map[string]*doqConn
closed bool
} }
type doqConn struct { type doqConn struct {
conn *quic.Conn conn *quic.Conn
lastUsed time.Time
refCount int
mu sync.Mutex
} }
func newDOQConnPool(_ context.Context, uc *UpstreamConfig, addrs []string) *doqConnPool { func newDOQConnPool(_ context.Context, uc *UpstreamConfig, addrs []string) *doqConnPool {
@@ -83,7 +79,7 @@ func newDOQConnPool(_ context.Context, uc *UpstreamConfig, addrs []string) *doqC
addrs: addrs, addrs: addrs,
port: port, port: port,
tlsConfig: tlsConfig, tlsConfig: tlsConfig,
conns: make(map[string]*doqConn), conns: make(chan *doqConn, doqPoolSize),
} }
// Use SetFinalizer here because we need to call a method on the pool itself. // Use SetFinalizer here because we need to call a method on the pool itself.
@@ -116,7 +112,7 @@ func (p *doqConnPool) Resolve(ctx context.Context, msg *dns.Msg) (*dns.Msg, erro
} }
func (p *doqConnPool) doResolve(ctx context.Context, msg *dns.Msg) (*dns.Msg, error) { func (p *doqConnPool) doResolve(ctx context.Context, msg *dns.Msg) (*dns.Msg, error) {
conn, addr, err := p.getConn(ctx) conn, err := p.getConn(ctx)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@@ -124,14 +120,14 @@ func (p *doqConnPool) doResolve(ctx context.Context, msg *dns.Msg) (*dns.Msg, er
// Pack the DNS message // Pack the DNS message
msgBytes, err := msg.Pack() msgBytes, err := msg.Pack()
if err != nil { if err != nil {
p.putConn(addr, conn, false) p.putConn(conn, false)
return nil, err return nil, err
} }
// Open a new stream for this query // Open a new stream for this query
stream, err := conn.OpenStream() stream, err := conn.OpenStream()
if err != nil { if err != nil {
p.putConn(addr, conn, false) p.putConn(conn, false)
return nil, err return nil, err
} }
@@ -147,13 +143,13 @@ func (p *doqConnPool) doResolve(ctx context.Context, msg *dns.Msg) (*dns.Msg, er
var msgLenBytes = []byte{byte(msgLen >> 8), byte(msgLen & 0xFF)} var msgLenBytes = []byte{byte(msgLen >> 8), byte(msgLen & 0xFF)}
if _, err := stream.Write(msgLenBytes); err != nil { if _, err := stream.Write(msgLenBytes); err != nil {
stream.Close() stream.Close()
p.putConn(addr, conn, false) p.putConn(conn, false)
return nil, err return nil, err
} }
if _, err := stream.Write(msgBytes); err != nil { if _, err := stream.Write(msgBytes); err != nil {
stream.Close() stream.Close()
p.putConn(addr, conn, false) p.putConn(conn, false)
return nil, err return nil, err
} }
@@ -163,7 +159,7 @@ func (p *doqConnPool) doResolve(ctx context.Context, msg *dns.Msg) (*dns.Msg, er
// Return connection to pool (mark as potentially bad if error occurred) // Return connection to pool (mark as potentially bad if error occurred)
isGood := err == nil && len(buf) > 0 isGood := err == nil && len(buf) > 0
p.putConn(addr, conn, isGood) p.putConn(conn, isGood)
if err != nil { if err != nil {
return nil, err return nil, err
@@ -184,79 +180,42 @@ func (p *doqConnPool) doResolve(ctx context.Context, msg *dns.Msg) (*dns.Msg, er
} }
// getConn gets a QUIC connection from the pool or creates a new one. // getConn gets a QUIC connection from the pool or creates a new one.
func (p *doqConnPool) getConn(ctx context.Context) (*quic.Conn, string, error) { // A connection is taken from the channel while in use; putConn returns it.
p.mu.Lock() func (p *doqConnPool) getConn(ctx context.Context) (*quic.Conn, error) {
defer p.mu.Unlock() for {
select {
if p.closed { case dc := <-p.conns:
return nil, "", io.EOF if dc.conn != nil && dc.conn.Context().Err() == nil {
} return dc.conn, nil
// Try to reuse an existing connection
for addr, doqConn := range p.conns {
doqConn.mu.Lock()
if doqConn.refCount == 0 && doqConn.conn != nil {
// Check if connection is still alive
select {
case <-doqConn.conn.Context().Done():
// Connection is closed, remove it
doqConn.mu.Unlock()
delete(p.conns, addr)
continue
default:
} }
if dc.conn != nil {
doqConn.refCount++ dc.conn.CloseWithError(quic.ApplicationErrorCode(quic.NoError), "")
doqConn.lastUsed = time.Now() }
conn := doqConn.conn default:
doqConn.mu.Unlock() _, conn, err := p.dialConn(ctx)
return conn, addr, nil if err != nil {
return nil, err
}
return conn, nil
} }
doqConn.mu.Unlock()
} }
// No available connection, create a new one
addr, conn, err := p.dialConn(ctx)
if err != nil {
return nil, "", err
}
doqConn := &doqConn{
conn: conn,
lastUsed: time.Now(),
refCount: 1,
}
p.conns[addr] = doqConn
return conn, addr, nil
} }
// putConn returns a connection to the pool. // putConn returns a connection to the pool for reuse by other goroutines.
func (p *doqConnPool) putConn(addr string, conn *quic.Conn, isGood bool) { func (p *doqConnPool) putConn(conn *quic.Conn, isGood bool) {
p.mu.Lock() if !isGood || conn == nil || conn.Context().Err() != nil {
defer p.mu.Unlock() if conn != nil {
conn.CloseWithError(quic.ApplicationErrorCode(quic.NoError), "")
doqConn, ok := p.conns[addr] }
if !ok {
return return
} }
dc := &doqConn{conn: conn}
doqConn.mu.Lock() select {
defer doqConn.mu.Unlock() case p.conns <- dc:
default:
doqConn.refCount-- // Channel full, close the connection
if doqConn.refCount < 0 { dc.conn.CloseWithError(quic.ApplicationErrorCode(quic.NoError), "")
doqConn.refCount = 0
} }
// If connection is bad or closed, remove it from pool
if !isGood || conn.Context().Err() != nil {
delete(p.conns, addr)
conn.CloseWithError(quic.ApplicationErrorCode(quic.NoError), "")
return
}
doqConn.lastUsed = time.Now()
} }
// dialConn creates a new QUIC connection using parallel dialing like DoH3. // dialConn creates a new QUIC connection using parallel dialing like DoH3.
@@ -301,23 +260,17 @@ func (p *doqConnPool) dialConn(ctx context.Context) (string, *quic.Conn, error)
return addr, conn, nil return addr, conn, nil
} }
// CloseIdleConnections closes all idle connections in the pool. // CloseIdleConnections closes all connections in the pool.
// When called during cleanup (e.g., from finalizer), it closes all connections // Connections currently checked out (in use) are not closed.
// regardless of refCount to prevent resource leaks.
func (p *doqConnPool) CloseIdleConnections() { func (p *doqConnPool) CloseIdleConnections() {
p.mu.Lock() for {
defer p.mu.Unlock() select {
case dc := <-p.conns:
p.closed = true if dc.conn != nil {
dc.conn.CloseWithError(quic.ApplicationErrorCode(quic.NoError), "")
for addr, dc := range p.conns { }
dc.mu.Lock() default:
if dc.conn != nil { return
// Close all connections to ensure proper cleanup, even if in use
// This prevents resource leaks when the pool is being destroyed
dc.conn.CloseWithError(quic.ApplicationErrorCode(quic.NoError), "")
} }
dc.mu.Unlock()
delete(p.conns, addr)
} }
} }