fix: flaky test by checking ctx cancellation before racing I/O in protobuf reader/writer (#5591)

Author: akrem-chabchoub

pkg/p2p/libp2p/stream_test.go

@@ -79,6 +79,7 @@ func TestStreamVersion(t *testing.T) {
 
 			if v == nil {
 				t.Fatal("expected version to be non-nil")
+				return
 			}
 
 			expected := semver.Version{Major: tc.wantMajor, Minor: tc.wantMinor}

pkg/p2p/protobuf/protobuf.go

@@ -56,6 +56,10 @@ func newReader(r ggio.Reader) Reader {
 }
 
 func (r Reader) ReadMsgWithContext(ctx context.Context, msg proto.Message) error {
+	if err := ctx.Err(); err != nil {
+		return err
+	}
+
 	errChan := make(chan error, 1)
 	go func() {
 		errChan <- r.ReadMsg(msg)
@@ -78,6 +82,10 @@ func newWriter(r ggio.Writer) Writer {
 }
 
 func (w Writer) WriteMsgWithContext(ctx context.Context, msg proto.Message) error {
+	if err := ctx.Err(); err != nil {
+		return err
+	}
+
 	errChan := make(chan error, 1)
 	go func() {
 		errChan <- w.WriteMsg(msg)