Optimize lock usage in proxy connection handling

- Change activeTunnel.conns from slice to map for O(1) add/remove
- Improve lock scoping in UpdateLocalSNIs: use read lock for diff
  computation, minimize write lock hold time
- Move cache invalidation outside lock (go-cache is thread-safe)
This commit is contained in:
Laurence
2026-03-13 15:23:30 +00:00
parent b9261b8fea
commit 14a3e7c531

View File

@@ -72,7 +72,7 @@ type SNIProxy struct {
} }
type activeTunnel struct { type activeTunnel struct {
conns []net.Conn conns map[net.Conn]struct{}
} }
// readOnlyConn is a wrapper for io.Reader that implements net.Conn // readOnlyConn is a wrapper for io.Reader that implements net.Conn
@@ -592,26 +592,19 @@ func (p *SNIProxy) handleConnection(clientConn net.Conn) {
p.activeTunnelsLock.Lock() p.activeTunnelsLock.Lock()
tunnel, ok := p.activeTunnels[hostname] tunnel, ok := p.activeTunnels[hostname]
if !ok { if !ok {
tunnel = &activeTunnel{} tunnel = &activeTunnel{conns: make(map[net.Conn]struct{})}
p.activeTunnels[hostname] = tunnel p.activeTunnels[hostname] = tunnel
} }
tunnel.conns = append(tunnel.conns, actualClientConn) tunnel.conns[actualClientConn] = struct{}{}
p.activeTunnelsLock.Unlock() p.activeTunnelsLock.Unlock()
defer func() { defer func() {
// Remove this conn from active tunnels // Remove this conn from active tunnels - O(1) with map
p.activeTunnelsLock.Lock() p.activeTunnelsLock.Lock()
if tunnel, ok := p.activeTunnels[hostname]; ok { if tunnel, ok := p.activeTunnels[hostname]; ok {
newConns := make([]net.Conn, 0, len(tunnel.conns)) delete(tunnel.conns, actualClientConn)
for _, c := range tunnel.conns { if len(tunnel.conns) == 0 {
if c != actualClientConn {
newConns = append(newConns, c)
}
}
if len(newConns) == 0 {
delete(p.activeTunnels, hostname) delete(p.activeTunnels, hostname)
} else {
tunnel.conns = newConns
} }
} }
p.activeTunnelsLock.Unlock() p.activeTunnelsLock.Unlock()
@@ -810,32 +803,42 @@ func (p *SNIProxy) ClearCache() {
// UpdateLocalSNIs updates the local SNIs and invalidates cache for changed domains // UpdateLocalSNIs updates the local SNIs and invalidates cache for changed domains
func (p *SNIProxy) UpdateLocalSNIs(fullDomains []string) { func (p *SNIProxy) UpdateLocalSNIs(fullDomains []string) {
newSNIs := make(map[string]struct{}) newSNIs := make(map[string]struct{}, len(fullDomains))
for _, domain := range fullDomains { for _, domain := range fullDomains {
newSNIs[domain] = struct{}{} newSNIs[domain] = struct{}{}
// Invalidate any cached route for this domain
p.cache.Delete(domain)
} }
// Update localSNIs // Get old SNIs with read lock to compute diff outside write lock
p.localSNIsLock.Lock() p.localSNIsLock.RLock()
oldSNIs := p.localSNIs
p.localSNIsLock.RUnlock()
// Compute removed SNIs outside the lock
removed := make([]string, 0) removed := make([]string, 0)
for sni := range p.localSNIs { for sni := range oldSNIs {
if _, stillLocal := newSNIs[sni]; !stillLocal { if _, stillLocal := newSNIs[sni]; !stillLocal {
removed = append(removed, sni) removed = append(removed, sni)
} }
} }
// Swap with minimal write lock hold time
p.localSNIsLock.Lock()
p.localSNIs = newSNIs p.localSNIs = newSNIs
p.localSNIsLock.Unlock() p.localSNIsLock.Unlock()
// Invalidate cache for new domains (cache is thread-safe)
for domain := range newSNIs {
p.cache.Delete(domain)
}
logger.Debug("Updated local SNIs, added %d, removed %d", len(newSNIs), len(removed)) logger.Debug("Updated local SNIs, added %d, removed %d", len(newSNIs), len(removed))
// Terminate tunnels for removed SNIs // Terminate tunnels for removed SNIs
if len(removed) > 0 { if len(removed) > 0 {
p.activeTunnelsLock.Lock() p.activeTunnelsLock.Lock()
for _, sni := range removed { for _, sni := range removed {
if tunnels, ok := p.activeTunnels[sni]; ok { if tunnel, ok := p.activeTunnels[sni]; ok {
for _, conn := range tunnels.conns { for conn := range tunnel.conns {
conn.Close() conn.Close()
} }
delete(p.activeTunnels, sni) delete(p.activeTunnels, sni)