mirror of
https://github.com/Control-D-Inc/ctrld.git
synced 2026-07-16 13:17:19 +02:00
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:
committed by
Cuong Manh Le
parent
60dd366cc4
commit
f1b93c81bc
@@ -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)
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user