Call wrapIfContextError from wrapIfContextDone

connectrpc · Feb 15, 2024 · c332407 · c332407
1 parent f925b2e
commit c332407
Show file tree

Hide file tree

Showing 3 changed files with 1 addition and 4 deletions.
diff --git a/envelope.go b/envelope.go
@@ -210,7 +210,6 @@ func (w *envelopeWriter) marshal(message any) *Error {
 
 func (w *envelopeWriter) write(env *envelope) *Error {
 	if _, err := w.sender.Send(env); err != nil {
-		err = wrapIfContextError(err)
 		err = wrapIfContextDone(w.ctx, err)
 		if connectErr, ok := asError(err); ok {
 			return connectErr
@@ -317,7 +316,6 @@ func (r *envelopeReader) Read(env *envelope) *Error {
 			return NewError(CodeUnknown, err)
 		}
 		err = wrapIfMaxBytesError(err, "read 5 byte message prefix")
-		err = wrapIfContextError(err)
 		err = wrapIfContextDone(r.ctx, err)
 		if connectErr, ok := asError(err); ok {
 			return connectErr
@@ -351,7 +349,6 @@ func (r *envelopeReader) Read(env *envelope) *Error {
 			)
 		}
 		err = wrapIfMaxBytesError(err, "read %d byte message", size)
-		err = wrapIfContextError(err)
 		err = wrapIfContextDone(r.ctx, err)
 		if connectErr, ok := asError(err); ok {
 			return connectErr

diff --git a/error.go b/error.go
@@ -314,6 +314,7 @@ func wrapIfContextDone(ctx context.Context, err error) error {
 	if err == nil {
 		return nil
 	}
+	err = wrapIfContextError(err)
 	if _, ok := asError(err); ok {
 		return err
 	}

diff --git a/protocol_connect.go b/protocol_connect.go
@@ -1115,7 +1115,6 @@ func (u *connectUnaryUnmarshaler) UnmarshalFunc(message any, unmarshal func([]by
 	bytesRead, err := data.ReadFrom(reader)
 	if err != nil {
 		err = wrapIfMaxBytesError(err, "read first %d bytes of message", bytesRead)
-		err = wrapIfContextError(err)
 		err = wrapIfContextDone(u.ctx, err)
 		if connectErr, ok := asError(err); ok {
 			return connectErr