encoding: call Close on decompressor reader if it implements io.Closer (#9135)

Currently, gzip decompressor readers are only returned to the sync.Pool
when an io.EOF is encountered during an active Read operation.

This PR added reader.Close in the gzip package and conditional
reader.Close() invocation via defer in decompress to ensure that reader
resources are properly reclaimed when decompression completes or fails.

RELEASE NOTES:
* encoding: gRPC now automatically calls Close() on any reader returned
by Compressor.Decompress that implements io.Closer when decompression
completes or fails

encoding/encoding.go

diff --git a/encoding/encoding.go b/encoding/encoding.go
index 296f38c..bfa8b26 100644
--- a/encoding/encoding.go
+++ b/encoding/encoding.go
@@ -66,6 +66,9 @@
 	// Decompress reads data from r, decompresses it, and provides the
 	// uncompressed data via the returned io.Reader.  If an error occurs while
 	// initializing the decompressor, that error is returned instead.
+	//
+	// The returned io.Reader may optionally implement io.ReadCloser, and if it
+	// does, gRPC will call Close() exactly once.
 	Decompress(r io.Reader) (io.Reader, error)
 	// Name is the name of the compression codec and is used to set the content
 	// coding header.  The result must be static; the result cannot change

encoding/gzip/gzip.go

diff --git a/encoding/gzip/gzip.go b/encoding/gzip/gzip.go
index 153e4db..65908d9 100644
--- a/encoding/gzip/gzip.go
+++ b/encoding/gzip/gzip.go
@@ -81,6 +81,8 @@
 	return z.Writer.Close()
 }
 
+var _ io.Closer = &reader{}
+
 type reader struct {
 	*gzip.Reader
 	pool *sync.Pool
@@ -102,14 +104,16 @@
 	return z, nil
 }
 
-func (z *reader) Read(p []byte) (n int, err error) {
-	n, err = z.Reader.Read(p)
-	if err == io.EOF {
-		z.pool.Put(z)
-	}
+func (r *reader) Read(p []byte) (n int, err error) {
+	n, err = r.Reader.Read(p)
 	return n, err
 }
 
+func (r *reader) Close() error {
+	defer r.pool.Put(r)
+	return r.Reader.Close()
+}
+
 func (c *compressor) Name() string {
 	return Name
 }

rpc_util.go

diff --git a/rpc_util.go b/rpc_util.go
index ec6af71..c3651a4 100644
--- a/rpc_util.go
+++ b/rpc_util.go
@@ -989,6 +989,9 @@
 			r.Close() // ensure buffers are reused
 			return nil, status.Errorf(codes.Internal, "grpc: failed to decompress the message: %v", err)
 		}
+		if closer, ok := dcReader.(io.Closer); ok {
+			defer closer.Close()
+		}
 
 		// Read at most one byte more than the limit from the decompressor.
 		// Unless the limit is MaxInt64, in which case, that's impossible, so

rpc_util_test.go

diff --git a/rpc_util_test.go b/rpc_util_test.go
index 79628d1..2abe751 100644
--- a/rpc_util_test.go
+++ b/rpc_util_test.go
@@ -589,3 +589,67 @@
 	}
 	wg.Wait()
 }
+
+type fakeCloseCompressor struct {
+	encoding.Compressor
+	io.Reader
+	ch chan<- struct{}
+}
+
+func (m *fakeCloseCompressor) Decompress(r io.Reader) (io.Reader, error) {
+	zr, err := gzip.NewReader(r)
+	if err != nil {
+		return nil, err
+	}
+	m.Reader = zr
+	return m, nil
+}
+
+func (m *fakeCloseCompressor) Close() error {
+	close(m.ch)
+	return nil
+}
+
+// TestDecompress_ClosesReader tests that the reader returned by Decompress is
+// closed when decompress returns.
+func (s) TestDecompress_ClosesReader(t *testing.T) {
+	testCases := []struct {
+		name                  string
+		maxReceiveMessageSize int
+		wantCode              codes.Code
+	}{
+		{
+			name:                  "close_on_failure",
+			maxReceiveMessageSize: 1,
+			wantCode:              codes.ResourceExhausted,
+		},
+		{
+			name:                  "close_on_success",
+			maxReceiveMessageSize: 100,
+			wantCode:              codes.OK,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+			defer cancel()
+
+			ch := make(chan struct{})
+			compressor := &fakeCloseCompressor{ch: ch}
+			in := compressWithDeterministicError(t, []byte("some data"))
+			out, err := decompress(compressor, in, nil, tc.maxReceiveMessageSize, mem.DefaultBufferPool())
+			if status.Code(err) != tc.wantCode {
+				t.Fatalf("decompress() failed with error code %v, want %v", status.Code(err), tc.wantCode)
+			}
+
+			out.Free()
+
+			select {
+			case <-ch:
+			case <-ctx.Done():
+				t.Fatalf("Timed out waiting for Close to be called")
+			}
+		})
+	}
+}