From fe7335e634c0b8257fe830b1eea5a09b7c093b7d Mon Sep 17 00:00:00 2001
From: Ville Vesilehto <ville@vesilehto.fi>
Date: Mon, 24 Nov 2025 18:20:30 +0200
Subject: [PATCH] perf(proxy): avoid unnecessary alloc in Yield (#7708)

---
 plugin/pkg/proxy/persistent.go      | 12 +++++++++---
 plugin/pkg/proxy/persistent_test.go | 29 +++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/plugin/pkg/proxy/persistent.go b/plugin/pkg/proxy/persistent.go
index 280941980..0bacc851a 100644
--- a/plugin/pkg/proxy/persistent.go
+++ b/plugin/pkg/proxy/persistent.go
@@ -128,9 +128,15 @@ const yieldTimeout = 25 * time.Millisecond
 func (t *Transport) Yield(pc *persistConn) {
 	pc.used = time.Now() // update used time
 
-	// Make this non-blocking, because in the case of a very busy forwarder we will *block* on this yield. This
-	// blocks the outer go-routine and stuff will just pile up.  We timeout when the send fails to as returning
-	// these connection is an optimization anyway.
+	// Optimization: Try to return the connection immediately without creating a timer.
+	// If the receiver is not ready, we fall back to a timeout-based send to avoid blocking forever.
+	// Returning the connection is just an optimization, so dropping it on timeout is fine.
+	select {
+	case t.yield <- pc:
+		return
+	default:
+	}
+
 	select {
 	case t.yield <- pc:
 		return
diff --git a/plugin/pkg/proxy/persistent_test.go b/plugin/pkg/proxy/persistent_test.go
index 56d837113..89715ab1d 100644
--- a/plugin/pkg/proxy/persistent_test.go
+++ b/plugin/pkg/proxy/persistent_test.go
@@ -1,6 +1,7 @@
 package proxy
 
 import (
+	"runtime"
 	"testing"
 	"time"
 
@@ -107,3 +108,31 @@ func TestCleanupAll(t *testing.T) {
 		t.Error("Expected no cached connections")
 	}
 }
+
+func BenchmarkYield(b *testing.B) {
+	s := dnstest.NewServer(func(w dns.ResponseWriter, r *dns.Msg) {
+		ret := new(dns.Msg)
+		ret.SetReply(r)
+		w.WriteMsg(ret)
+	})
+	defer s.Close()
+
+	tr := newTransport("BenchmarkYield", s.Addr)
+	tr.Start()
+	defer tr.Stop()
+
+	c, _, _ := tr.Dial("udp")
+
+	b.ReportAllocs()
+
+	for b.Loop() {
+		tr.Yield(c)
+		// Drain the yield channel so we can yield again without blocking/timing out
+		// We need to simulate the consumer side slightly to keep Yield flowing
+		select {
+		case <-tr.yield:
+		default:
+		}
+		runtime.Gosched()
+	}
+}