{trpc, attachment}: pre-check empty attachment while getting or setting to avoid memory allocation

client/attachment_test.go

@@ -20,10 +20,10 @@ import (
 func TestAttachment(t *testing.T) {
 	attm := NewAttachment(bytes.NewReader([]byte("attachment")))
 	require.Equal(t, attachment.NoopAttachment{}, attm.Response())
-
 	msg := codec.Message(context.Background())
 	setAttachment(msg, &attm.attachment)
-	attcher := attachment.GetClientRequestAttachment(msg)
+	attcher, ok := attachment.ClientRequestAttachment(msg)
+	require.True(t, ok)
 	bts, err := io.ReadAll(attcher)
 	require.Nil(t, err)
 	require.Equal(t, []byte("attachment"), bts)

codec.go

@@ -326,7 +326,9 @@ func msgWithRequestProtocol(msg codec.Msg, req *trpcpb.RequestProtocol, attm []b
 	}
 	// set call type
 	msg.WithCallType(codec.RequestType(req.GetCallType()))
-	attachment.SetServerRequestAttachment(msg, attm)
+	if len(attm) != 0 {
+		attachment.SetServerRequestAttachment(msg, attm)
+	}
 }
 
 // Encode implements codec.Codec.
@@ -342,9 +344,12 @@ func (s *ServerCodec) Encode(msg codec.Msg, rspBody []byte) ([]byte, error) {
 
 	rspProtocol := getAndInitResponseProtocol(msg)
 
-	attm, err := io.ReadAll(attachment.GetServerResponseAttachment(msg))
-	if err != nil {
-		return nil, fmt.Errorf("encoding attachment : %w", err)
+	var attm []byte
+	if a, ok := attachment.ServerResponseAttachment(msg); ok {
+		var err error
+		if attm, err = io.ReadAll(a); err != nil {
+			return nil, fmt.Errorf("encoding attachment: %w", err)
+		}
 	}
 	rspProtocol.AttachmentSize = uint32(len(attm))
 
@@ -467,9 +472,11 @@ func (c *ClientCodec) Encode(msg codec.Msg, reqBody []byte) (reqBuf []byte, err
 	frameHead.upgradeProtocol(curProtocolVersion, requestID)
 	msg.WithRequestID(requestID)
 
-	attm, err := io.ReadAll(attachment.GetClientRequestAttachment(msg))
-	if err != nil {
-		return nil, fmt.Errorf("encoding attachment : %w", err)
+	var attm []byte
+	if a, ok := attachment.ClientRequestAttachment(msg); ok {
+		if attm, err = io.ReadAll(a); err != nil {
+			return nil, fmt.Errorf("encoding attachment: %w", err)
+		}
 	}
 	req.AttachmentSize = uint32(len(attm))
 
@@ -700,6 +707,8 @@ func updateMsg(msg codec.Msg, frameHead *FrameHead, rsp *trpcpb.ResponseProtocol
 	frameHead.upgradeProtocol(curProtocolVersion, rsp.RequestId)
 	msg.WithRequestID(rsp.RequestId)
 
-	attachment.SetClientResponseAttachment(msg, attm)
+	if len(attm) != 0 {
+		attachment.SetClientResponseAttachment(msg, attm)
+	}
 	return nil
 }

codec_test.go

@@ -9,6 +9,7 @@ import (
 	"bytes"
 	"context"
 	"encoding/binary"
+	"errors"
 	"log"
 	"net"
 	"regexp"
@@ -18,6 +19,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/protobuf/proto"
+	"trpc.group/trpc-go/trpc-go/internal/attachment"
 	trpcpb "trpc.group/trpc/trpc-protocol/pb/go/trpc"
 
 	"trpc.group/trpc-go/trpc-go"
@@ -256,34 +258,66 @@ func TestServerCodec_CallTypeDecode(t *testing.T) {
 }
 
 func TestClientCodec_EncodeErr(t *testing.T) {
-	cc := trpc.ClientCodec{}
-	msg := codec.Message(trpc.BackgroundContext())
-	msg.WithClientMetaData(codec.MetaData{"overHeadLengthU16": make([]byte, 64*1024)})
-	_, err := cc.Encode(msg, nil)
-	assert.EqualError(t, err, "head len overflows uint16")
+	t.Run("head len overflows uint16", func(t *testing.T) {
+		cc := trpc.ClientCodec{}
+		msg := codec.Message(trpc.BackgroundContext())
+		msg.WithClientMetaData(codec.MetaData{"overHeadLengthU16": make([]byte, 64*1024)})
+		_, err := cc.Encode(msg, nil)
+		assert.EqualError(t, err, "head len overflows uint16")
+	})
+	t.Run("frame len is too large", func(t *testing.T) {
+		cc := trpc.ClientCodec{}
+		msg := codec.Message(trpc.BackgroundContext())
+		_, err := cc.Encode(msg, make([]byte, trpc.DefaultMaxFrameSize))
+		assert.EqualError(t, err, "frame len is larger than MaxFrameSize(10485760)")
+	})
+	t.Run("encoding attachment failed", func(t *testing.T) {
+		cc := trpc.ClientCodec{}
+		msg := codec.Message(trpc.BackgroundContext())
+		msg.WithCommonMeta(codec.CommonMeta{attachment.ClientAttachmentKey{}: &attachment.Attachment{Request: &errorReader{}, Response: attachment.NoopAttachment{}}})
+		_, err := cc.Encode(msg, nil)
+		assert.EqualError(t, err, "encoding attachment: reading errorReader always returns error")
+	})
 
-	msg = codec.Message(trpc.BackgroundContext())
-	_, err = cc.Encode(msg, make([]byte, trpc.DefaultMaxFrameSize))
-	assert.EqualError(t, err, "frame len is larger than MaxFrameSize(10485760)")
+}
+
+type errorReader struct{}
+
+func (*errorReader) Read(p []byte) (n int, err error) {
+	return 0, errors.New("reading errorReader always returns error")
 }
 
 func TestServerCodec_EncodeErr(t *testing.T) {
-	sc := trpc.ServerCodec{}
-	msg := codec.Message(trpc.BackgroundContext())
-	msg.WithServerMetaData(codec.MetaData{"overHeadLengthU16": make([]byte, 64*1024)})
-	rspBuf, err := sc.Encode(msg, nil)
-	assert.Nil(t, err)
-	head := &trpcpb.ResponseProtocol{}
-	err = proto.Unmarshal(rspBuf[16:], head)
-	assert.Nil(t, err)
-	assert.Equal(t, int32(errs.RetServerEncodeFail), head.GetRet())
+	t.Run("head len overflows uint16", func(t *testing.T) {
+		msg := codec.Message(trpc.BackgroundContext())
+		sc := trpc.ServerCodec{}
+		msg.WithServerMetaData(codec.MetaData{"overHeadLengthU16": make([]byte, 64*1024)})
+		rspBuf, err := sc.Encode(msg, nil)
+		assert.Nil(t, err)
 
-	msg = codec.Message(trpc.BackgroundContext())
-	rspBuf, err = sc.Encode(msg, make([]byte, trpc.DefaultMaxFrameSize))
-	assert.Nil(t, err)
-	err = proto.Unmarshal(rspBuf[16:], head)
-	assert.Nil(t, err)
-	assert.Equal(t, int32(errs.RetServerEncodeFail), head.GetRet())
+		head := &trpcpb.ResponseProtocol{}
+		err = proto.Unmarshal(rspBuf[16:], head)
+		assert.Nil(t, err)
+		assert.Equal(t, int32(errs.RetServerEncodeFail), head.GetRet())
+	})
+	t.Run("frame len is too large", func(t *testing.T) {
+		msg := codec.Message(trpc.BackgroundContext())
+		sc := trpc.ServerCodec{}
+		rspBuf, err := sc.Encode(msg, make([]byte, trpc.DefaultMaxFrameSize))
+		assert.Nil(t, err)
+
+		head := &trpcpb.ResponseProtocol{}
+		err = proto.Unmarshal(rspBuf[16:], head)
+		assert.Nil(t, err)
+		assert.Equal(t, int32(errs.RetServerEncodeFail), head.GetRet())
+	})
+	t.Run("encoding attachment failed", func(t *testing.T) {
+		msg := codec.Message(trpc.BackgroundContext())
+		msg.WithCommonMeta(codec.CommonMeta{attachment.ServerAttachmentKey{}: &attachment.Attachment{Request: attachment.NoopAttachment{}, Response: &errorReader{}}})
+		sc := trpc.ServerCodec{}
+		_, err := sc.Encode(msg, nil)
+		assert.EqualError(t, err, "encoding attachment: reading errorReader always returns error")
+	})
 }
 
 func TestMultiplexFrame(t *testing.T) {

internal/attachment/attachment.go

@@ -25,28 +25,28 @@ type Attachment struct {
 	Response io.Reader
 }
 
-// NoopAttachment is a empty attachment.
+// NoopAttachment is an empty attachment.
 type NoopAttachment struct{}
 
 // Read implements the io.Reader interface, which always returns (0, io.EOF)
 func (a NoopAttachment) Read(_ []byte) (n int, err error) {
 	return 0, io.EOF
 }
 
-// GetClientRequestAttachment returns client's Request Attachment from msg.
-func GetClientRequestAttachment(msg codec.Msg) io.Reader {
+// ClientRequestAttachment returns client's Request Attachment from msg.
+func ClientRequestAttachment(msg codec.Msg) (io.Reader, bool) {
 	if a, _ := msg.CommonMeta()[ClientAttachmentKey{}].(*Attachment); a != nil {
-		return a.Request
+		return a.Request, true
 	}
-	return NoopAttachment{}
+	return nil, false
 }
 
-// GetServerResponseAttachment returns server's Response Attachment from msg.
-func GetServerResponseAttachment(msg codec.Msg) io.Reader {
+// ServerResponseAttachment returns server's Response Attachment from msg.
+func ServerResponseAttachment(msg codec.Msg) (io.Reader, bool) {
 	if a, _ := msg.CommonMeta()[ServerAttachmentKey{}].(*Attachment); a != nil {
-		return a.Response
+		return a.Response, true
 	}
-	return NoopAttachment{}
+	return nil, false
 }
 
 // SetClientResponseAttachment sets client's Response attachment to msg.

internal/attachment/attachment_test.go

@@ -23,61 +23,57 @@ import (
 func TestGetClientRequestAttachment(t *testing.T) {
 	t.Run("nil message", func(t *testing.T) {
 		require.Panics(t, func() {
-			attachment.GetClientRequestAttachment(nil)
+			attachment.ClientRequestAttachment(nil)
 		})
 	})
 	t.Run("empty message", func(t *testing.T) {
 		msg := trpc.Message(context.Background())
-		want := attachment.NoopAttachment{}
-		if got := attachment.GetClientRequestAttachment(msg); !reflect.DeepEqual(got, want) {
-			t.Errorf("GetClientRequestAttachment() = %v, want %v", got, want)
-		}
+		_, ok := attachment.ClientRequestAttachment(msg)
+		require.False(t, ok)
 	})
 	t.Run("message contains nil attachment", func(t *testing.T) {
 		msg := trpc.Message(context.Background())
 		msg.WithCommonMeta(codec.CommonMeta{attachment.ClientAttachmentKey{}: nil})
-		want := attachment.NoopAttachment{}
-		if got := attachment.GetClientRequestAttachment(msg); !reflect.DeepEqual(got, want) {
-			t.Errorf("GetClientRequestAttachment() = %v, want %v", got, want)
-		}
+		_, ok := attachment.ClientRequestAttachment(msg)
+		require.False(t, ok)
 	})
 	t.Run("message contains non-empty Request attachment", func(t *testing.T) {
 		msg := trpc.Message(context.Background())
 		want := bytes.NewReader([]byte("attachment"))
 		msg.WithCommonMeta(codec.CommonMeta{attachment.ClientAttachmentKey{}: &attachment.Attachment{Request: want}})
-		if got := attachment.GetClientRequestAttachment(msg); !reflect.DeepEqual(got, want) {
-			t.Errorf("GetClientRequestAttachment() = %v, want %v", got, want)
+		got, ok := attachment.ClientRequestAttachment(msg)
+		require.True(t, ok)
+		if !reflect.DeepEqual(got, want) {
+			t.Errorf("ServerResponseAttachment() = %v, want %v", got, want)
 		}
 	})
 }
 
 func TestGetServerResponseAttachment(t *testing.T) {
 	t.Run("nil message", func(t *testing.T) {
 		require.Panics(t, func() {
-			attachment.GetServerResponseAttachment(nil)
+			attachment.ServerResponseAttachment(nil)
 		})
 	})
 	t.Run("empty message", func(t *testing.T) {
 		msg := trpc.Message(context.Background())
-		want := attachment.NoopAttachment{}
-		if got := attachment.GetServerResponseAttachment(msg); !reflect.DeepEqual(got, want) {
-			t.Errorf("GetServerResponseAttachment() = %v, want %v", got, want)
-		}
+		_, ok := attachment.ServerResponseAttachment(msg)
+		require.False(t, ok)
 	})
 	t.Run("message contains nil attachment", func(t *testing.T) {
 		msg := trpc.Message(context.Background())
 		msg.WithCommonMeta(codec.CommonMeta{attachment.ClientAttachmentKey{}: nil})
-		want := attachment.NoopAttachment{}
-		if got := attachment.GetServerResponseAttachment(msg); !reflect.DeepEqual(got, want) {
-			t.Errorf("GetServerResponseAttachment() = %v, want %v", got, want)
-		}
+		_, ok := attachment.ClientRequestAttachment(msg)
+		require.False(t, ok)
 	})
 	t.Run("message contains non-empty response attachment", func(t *testing.T) {
 		msg := trpc.Message(context.Background())
 		want := bytes.NewReader([]byte("attachment"))
 		msg.WithCommonMeta(codec.CommonMeta{attachment.ServerAttachmentKey{}: &attachment.Attachment{Response: want}})
-		if got := attachment.GetServerResponseAttachment(msg); !reflect.DeepEqual(got, want) {
-			t.Errorf("GetServerResponseAttachment() = %v, want %v", got, want)
+		got, ok := attachment.ServerResponseAttachment(msg)
+		require.True(t, ok)
+		if !reflect.DeepEqual(got, want) {
+			t.Errorf("ServerResponseAttachment() = %v, want %v", got, want)
 		}
 	})
 }

server/attachment_test.go

@@ -24,7 +24,8 @@ func TestAttachment(t *testing.T) {
 	require.Equal(t, attachment.NoopAttachment{}, attm.Request())
 
 	attm.SetResponse(bytes.NewReader([]byte("attachment")))
-	responseAttm := attachment.GetServerResponseAttachment(msg)
+	responseAttm, ok := attachment.ServerResponseAttachment(msg)
+	require.True(t, ok)
 	bts, err := io.ReadAll(responseAttm)
 	require.Nil(t, err)
 	require.Equal(t, []byte("attachment"), bts)