Add bufferpool

encoding/thrift/outbound_nowire.go

@@ -32,6 +32,7 @@ import (
 	encodingapi "go.uber.org/yarpc/api/encoding"
 	"go.uber.org/yarpc/api/transport"
 	"go.uber.org/yarpc/encoding/thrift/internal"
+	"go.uber.org/yarpc/internal/bufferpool"
 	"go.uber.org/yarpc/pkg/encoding"
 	"go.uber.org/yarpc/pkg/errors"
 	"go.uber.org/yarpc/pkg/procedure"
@@ -138,10 +139,11 @@ func (c noWireThriftClient) Call(ctx context.Context, reqBody stream.Enveloper,
 
 	out := c.cc.GetUnaryOutbound()
 
-	treq, proto, err := c.buildTransportRequest(reqBody)
+	treq, proto, cleanup, err := c.buildTransportRequest(reqBody)
 	if err != nil {
 		return err
 	}
+	defer cleanup()
 
 	call := encodingapi.NewOutboundCall(encoding.FromOptions(opts)...)
 	ctx, err = call.WriteToRequest(ctx, treq)
@@ -196,10 +198,11 @@ func (c noWireThriftClient) Call(ctx context.Context, reqBody stream.Enveloper,
 func (c noWireThriftClient) CallOneway(ctx context.Context, reqBody stream.Enveloper, opts ...yarpc.CallOption) (transport.Ack, error) {
 	out := c.cc.GetOnewayOutbound()
 
-	treq, _, err := c.buildTransportRequest(reqBody)
+	treq, _, cleanup, err := c.buildTransportRequest(reqBody)
 	if err != nil {
 		return nil, err
 	}
+	defer cleanup()
 
 	call := encodingapi.NewOutboundCall(encoding.FromOptions(opts)...)
 	ctx, err = call.WriteToRequest(ctx, treq)
@@ -214,7 +217,7 @@ func (c noWireThriftClient) Enabled() bool {
 	return c.NoWire
 }
 
-func (c noWireThriftClient) buildTransportRequest(reqBody stream.Enveloper) (*transport.Request, stream.Protocol, error) {
+func (c noWireThriftClient) buildTransportRequest(reqBody stream.Enveloper) (*transport.Request, stream.Protocol, func(), error) {
 	proto := c.p
 	if !c.Enveloping {
 		proto = disableEnvelopingNoWireProtocol{
@@ -232,32 +235,38 @@ func (c noWireThriftClient) buildTransportRequest(reqBody stream.Enveloper) (*tr
 
 	envType := reqBody.EnvelopeType()
 	if envType != wire.Call && envType != wire.OneWay {
-		return nil, nil, errors.RequestBodyEncodeError(
+		return nil, nil, nil, errors.RequestBodyEncodeError(
 			&treq, errUnexpectedEnvelopeType(envType),
 		)
 	}
 
-	var buffer bytes.Buffer
-	sw := proto.Writer(&buffer)
+	// Free the buffer with the cleanup function after the transport has finished reading it.
+	buffer := bufferpool.Get()
+	cleanup := func() { bufferpool.Put(buffer) }
+
+	sw := proto.Writer(buffer)
 	defer sw.Close()
 
 	if err := sw.WriteEnvelopeBegin(stream.EnvelopeHeader{
 		Name:  reqBody.MethodName(),
 		Type:  envType,
 		SeqID: 1, // don't care
 	}); err != nil {
-		return nil, nil, errors.RequestBodyEncodeError(&treq, err)
+		cleanup()
+		return nil, nil, func() {}, errors.RequestBodyEncodeError(&treq, err)
 	}
 
 	if err := reqBody.Encode(sw); err != nil {
-		return nil, nil, errors.RequestBodyEncodeError(&treq, err)
+		cleanup()
+		return nil, nil, func() {}, errors.RequestBodyEncodeError(&treq, err)
 	}
 
 	if err := sw.WriteEnvelopeEnd(); err != nil {
-		return nil, nil, errors.RequestBodyEncodeError(&treq, err)
+		cleanup()
+		return nil, nil, func() {}, errors.RequestBodyEncodeError(&treq, err)
 	}
 
-	treq.Body = &buffer
+	treq.Body = bytes.NewReader(buffer.Bytes())
 	treq.BodySize = buffer.Len()
-	return &treq, proto, nil
+	return &treq, proto, cleanup, nil
 }

encoding/thrift/outbound_nowire_test.go

@@ -355,7 +355,8 @@ func TestBuildTransportRequestWriteError(t *testing.T) {
 		sw.EXPECT().Close().Return(nil)
 		sw.EXPECT().WriteEnvelopeBegin(wantEnvHeader).Return(errors.New("writeenvelopebegin error"))
 
-		_, _, err := nwc.buildTransportRequest(fakeEnveloper(wire.Call))
+		_, _, cleanup, err := nwc.buildTransportRequest(fakeEnveloper(wire.Call))
+		defer cleanup()
 		require.Error(t, err)
 		assert.Contains(t, err.Error(), `failed to encode "thrift" request body for procedure "::someMethod" of service "service": writeenvelopebegin error`)
 	})
@@ -364,7 +365,8 @@ func TestBuildTransportRequestWriteError(t *testing.T) {
 		sw.EXPECT().Close().Return(nil)
 		sw.EXPECT().WriteEnvelopeBegin(wantEnvHeader).Return(nil)
 
-		_, _, err := nwc.buildTransportRequest(errorEnveloper{envelopeType: wire.Call, err: errors.New("encode error")})
+		_, _, cleanup, err := nwc.buildTransportRequest(errorEnveloper{envelopeType: wire.Call, err: errors.New("encode error")})
+		defer cleanup()
 		require.Error(t, err)
 		assert.Contains(t, err.Error(), `failed to encode "thrift" request body for procedure "::someMethod" of service "service": encode error`)
 	})
@@ -375,7 +377,8 @@ func TestBuildTransportRequestWriteError(t *testing.T) {
 		sw.EXPECT().WriteString(_irrelevant).Return(nil)
 		sw.EXPECT().WriteEnvelopeEnd().Return(errors.New("writeenvelopeend error"))
 
-		_, _, err := nwc.buildTransportRequest(fakeEnveloper(wire.Call))
+		_, _, cleanup, err := nwc.buildTransportRequest(fakeEnveloper(wire.Call))
+		defer cleanup()
 		require.Error(t, err)
 		assert.Contains(t, err.Error(), `failed to encode "thrift" request body for procedure "::someMethod" of service "service": writeenvelopeend error`)
 	})