From ce0cb913bc7589687f2019071bac5f51f0a9c415 Mon Sep 17 00:00:00 2001
From: Klaus Post <klauspost@gmail.com>
Date: Fri, 2 Feb 2024 08:48:12 -0800
Subject: [PATCH] Fix ineffective recycling (#18952)

Recycle would always be called on the dummy value `any(newRT())` instead of the actual value given to the recycle function.

Caught by race tests, but mostly harmless, except for reduced perf.

Other minor cleanups. Introduced in #18940 (unreleased)
---
 internal/grid/handlers.go | 17 ++++++++++-------
 internal/grid/types.go    |  1 -
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/internal/grid/handlers.go b/internal/grid/handlers.go
index 2773d3229..e0ea53bb9 100644
--- a/internal/grid/handlers.go
+++ b/internal/grid/handlers.go
@@ -348,9 +348,14 @@ type SingleHandler[Req, Resp RoundTripper] struct {
 
 func recycleFunc[RT RoundTripper](newRT func() RT) (newFn func() RT, recycle func(r RT)) {
 	rAny := any(newRT())
-	if rc, ok := rAny.(Recycler); ok {
+	var rZero RT
+	if _, ok := rAny.(Recycler); ok {
 		return newRT, func(r RT) {
-			rc.Recycle()
+			if r != rZero {
+				if rc, ok := any(r).(Recycler); ok {
+					rc.Recycle()
+				}
+			}
 		}
 	}
 	pool := sync.Pool{
@@ -358,7 +363,6 @@ func recycleFunc[RT RoundTripper](newRT func() RT) (newFn func() RT, recycle fun
 			return newRT()
 		},
 	}
-	var rZero RT
 	return func() RT { return pool.Get().(RT) },
 		func(r RT) {
 			if r != rZero {
@@ -436,6 +440,7 @@ func (h *SingleHandler[Req, Resp]) Register(m *Manager, handle func(req Req) (re
 		}
 		resp, rerr := handle(req)
 		h.recycleReq(req)
+
 		if rerr != nil {
 			PutByteBuffer(payload)
 			return nil, rerr
@@ -474,22 +479,20 @@ func (h *SingleHandler[Req, Resp]) Call(ctx context.Context, c Requester, req Re
 		ctx = context.WithValue(ctx, TraceParamsKey{}, fmt.Sprintf("type=%T", req))
 	}
 	if h.callReuseReq {
-		defer func() {
-			h.recycleReq(req)
-		}()
+		defer h.recycleReq(req)
 	}
 	res, err := c.Request(ctx, h.id, payload)
 	PutByteBuffer(payload)
 	if err != nil {
 		return resp, err
 	}
+	defer PutByteBuffer(res)
 	r := h.NewResponse()
 	_, err = r.UnmarshalMsg(res)
 	if err != nil {
 		h.PutResponse(r)
 		return resp, err
 	}
-	PutByteBuffer(res)
 	return r, err
 }
 
diff --git a/internal/grid/types.go b/internal/grid/types.go
index e29387b2c..70a281c00 100644
--- a/internal/grid/types.go
+++ b/internal/grid/types.go
@@ -235,7 +235,6 @@ func (b *Bytes) UnmarshalMsg(bytes []byte) ([]byte, error) {
 		copy(*b, val)
 	} else {
 		if cap(*b) == 0 && len(val) <= maxBufferSize {
-			PutByteBuffer(*b)
 			*b = GetByteBuffer()[:0]
 		} else {
 			PutByteBuffer(*b)