Made tcp connection context-aware

sei-tendermint/internal/mux/mux.go

@@ -34,7 +34,6 @@ import (
 	"github.com/tendermint/tendermint/libs/utils"
 	"github.com/tendermint/tendermint/libs/utils/scope"
 	"google.golang.org/protobuf/proto"
-	"io"
 )
 
 const handshakeMaxSize = 10 * 1024 // 10kB
@@ -46,8 +45,6 @@ var errTooManyMsgs = errors.New("too many messages")
 var errTooLargeMsg = errors.New("message too large")
 var errUnknownKind = errors.New("unknown kind")
 
-type errConn struct{ error }
-
 type Config struct {
 	// Maximal number of bytes in a frame (excluding header).
 	FrameSize uint64
@@ -127,15 +124,13 @@ func newRunner(mux *Mux) *runner {
 // If the stream does not exist yet, it tries to create it as an accept (inbound) stream.
 // In that case the inbound stream limit for the given kind is checked.
 func (r *runner) getOrAccept(id streamID, kind StreamKind) (*streamState, error) {
-	fmt.Printf("getOrAccept(%v)\n", id)
 	for inner := range r.inner.RLock() {
 		s, ok := inner.streams[id]
 		if ok {
 			return s, nil
 		}
 	}
 	for inner := range r.inner.Lock() {
-		fmt.Printf("accepting stream %v\n", id)
 		if id.isConnect() {
 			return nil, errUnknownStream
 		}
@@ -220,36 +215,36 @@ func (r *runner) runSend(ctx context.Context, conn conn.Conn) error {
 			if err != nil {
 				panic(err)
 			}
-			if _, err := conn.Write([]byte{byte(len(headerRaw))}); err != nil {
-				return errConn{err}
+			if err := conn.Write(ctx, []byte{byte(len(headerRaw))}); err != nil {
+				return err
 			}
-			if _, err := conn.Write(headerRaw); err != nil {
-				return errConn{err}
+			if err := conn.Write(ctx, headerRaw); err != nil {
+				return err
 			}
-			if _, err := conn.Write(f.Payload); err != nil {
-				return errConn{err}
+			if err := conn.Write(ctx, f.Payload); err != nil {
+				return err
 			}
 		}
 		if flush {
-			if err := conn.Flush(); err != nil {
-				return errConn{err}
+			if err := conn.Flush(ctx); err != nil {
+				return err
 			}
 		}
 	}
 }
 
 // runRecv receives and processes the incoming frames sequentially.
-func (r *runner) runRecv(conn conn.Conn) error {
+func (r *runner) runRecv(ctx context.Context, conn conn.Conn) error {
 	for {
 		// frame size is hard capped here at 255B.
 		// Currently we have 7 varint fields (up to 77B)
 		var headerSize [1]byte
-		if _, err := conn.Read(headerSize[:]); err != nil {
-			return errConn{err}
+		if err := conn.Read(ctx, headerSize[:]); err != nil {
+			return err
 		}
 		headerRaw := make([]byte, headerSize[0])
-		if _, err := io.ReadFull(conn, headerRaw[:]); err != nil {
-			return errConn{err}
+		if err := conn.Read(ctx, headerRaw[:]); err != nil {
+			return err
 		}
 		var h pb.Header
 		if err := proto.Unmarshal(headerRaw, &h); err != nil {
@@ -263,7 +258,6 @@ func (r *runner) runRecv(conn conn.Conn) error {
 			return err
 		}
 		for sInner := range s.inner.Lock() {
-			fmt.Printf("sInner.closed.remote[%v] = %v, %v\n", s.id, sInner.closed.remote, &h)
 			if sInner.closed.remote {
 				return errFrameAfterClose
 			}
@@ -286,7 +280,7 @@ func (r *runner) runRecv(conn conn.Conn) error {
 			}
 			// Read the payload.
 			payload := make([]byte, ps)
-			if _, err := io.ReadFull(conn, payload[:]); err != nil {
+			if err := conn.Read(ctx, payload[:]); err != nil {
 				return err
 			}
 			s.RemotePayload(payload)
@@ -307,52 +301,37 @@ func (r *runner) runRecv(conn conn.Conn) error {
 
 // Run runs the multiplexer for the given connection.
 // It closes the connection before return.
-func (m *Mux) Run(ctx context.Context, conn conn.Conn) error {
-	return scope.Run(ctx, func(ctx context.Context, s scope.Scope) error {
-		// Close on cancel.
-		s.Spawn(func() error {
-			<-ctx.Done()
-			s.Cancel(ctx.Err())
-			conn.Close()
-			return nil
-		})
-
+func (m *Mux) Run(ctx context.Context, conn conn.Conn) (err error) {
+	return utils.IgnoreCancel(scope.Run(ctx, func(ctx context.Context, s scope.Scope) error {
 		// Handshake exchange.
 		handshake, err := scope.Run1(ctx, func(ctx context.Context, s scope.Scope) (*handshake, error) {
 			s.Spawn(func() error {
-				handshakeRaw, err := proto.Marshal(handshakeConv.Encode(&handshake{Kinds: m.cfg.Kinds}))
-				if err != nil {
-					panic(err)
-				}
+				handshakeRaw := handshakeConv.Marshal(&handshake{Kinds: m.cfg.Kinds})
 				sizeRaw := binary.LittleEndian.AppendUint32(nil, uint32(len(handshakeRaw)))
-				if _, err := conn.Write(sizeRaw); err != nil {
-					return errConn{err}
+				if err := conn.Write(ctx, sizeRaw); err != nil {
+					return err
 				}
-				if _, err := conn.Write(handshakeRaw); err != nil {
-					return errConn{err}
+				if err := conn.Write(ctx, handshakeRaw); err != nil {
+					return err
 				}
-				if err := conn.Flush(); err != nil {
-					return errConn{err}
+				if err := conn.Flush(ctx); err != nil {
+					return err
 				}
 				return nil
 			})
 			var sizeRaw [4]byte
-			if _, err := io.ReadFull(conn, sizeRaw[:]); err != nil {
-				return nil, errConn{err}
+			if err := conn.Read(ctx, sizeRaw[:]); err != nil {
+				return nil, err
 			}
 			size := binary.LittleEndian.Uint32(sizeRaw[:])
 			if size > handshakeMaxSize {
 				return nil, fmt.Errorf("handshake too large")
 			}
 			handshakeRaw := make([]byte, size)
-			if _, err := io.ReadFull(conn, handshakeRaw[:]); err != nil {
-				return nil, errConn{err}
-			}
-			var handshakeProto pb.Handshake
-			if err := proto.Unmarshal(handshakeRaw, &handshakeProto); err != nil {
+			if err := conn.Read(ctx, handshakeRaw[:]); err != nil {
 				return nil, err
 			}
-			return handshakeConv.Decode(&handshakeProto)
+			return handshakeConv.Unmarshal(handshakeRaw)
 		})
 		if err != nil {
 			return err
@@ -374,9 +353,9 @@ func (m *Mux) Run(ctx context.Context, conn conn.Conn) error {
 		}
 		// Run the tasks.
 		s.Spawn(func() error { return r.runSend(ctx, conn) })
-		s.Spawn(func() error { return r.runRecv(conn) })
+		s.Spawn(func() error { return r.runRecv(ctx, conn) })
 		return nil
-	})
+	}))
 }
 
 // queue is a queue of frames to send, consumed by runSend.

sei-tendermint/internal/mux/mux_test.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"github.com/tendermint/tendermint/crypto/ed25519"
 	"github.com/tendermint/tendermint/internal/p2p/conn"
 	"github.com/tendermint/tendermint/libs/utils"
 	"github.com/tendermint/tendermint/libs/utils/require"
@@ -14,25 +13,21 @@ import (
 	"testing"
 )
 
-// Ignores cancellation and connection errors.
-// Wrap Mux.Run() calls in this - in tests we manage both ends of the connection,
-// so there is a race condition between disconnects and cancellation, when test
-// is shutting down.
-func ignoreDisconnect(err error) error {
-	if utils.ErrorAs[errConn](err).IsPresent() {
-		return nil
-	}
-	return utils.IgnoreCancel(err)
-}
-
-func testConn(t *testing.T) (*conn.SecretConnection, *conn.SecretConnection) {
+func spawnBgPipe(ctx context.Context, s scope.Scope) (*conn.SecretConnection, *conn.SecretConnection) {
 	c1, c2 := tcp.TestPipe()
+	s.SpawnBg(func() error {
+		_ = c1.Run(ctx)
+		return nil
+	})
+	s.SpawnBg(func() error {
+		_ = c2.Run(ctx)
+		return nil
+	})
 	var scs [2]*conn.SecretConnection
-	utils.OrPanic(scope.Run(t.Context(), func(ctx context.Context, s scope.Scope) error {
+	utils.OrPanic(scope.Run(ctx, func(ctx context.Context, s scope.Scope) error {
 		for i, c := range utils.Slice(c1, c2) {
-			t.Cleanup(func() { c.Close() })
 			s.Spawn(func() error {
-				scs[i] = utils.OrPanic1(conn.MakeSecretConnection(ctx, c, ed25519.GenerateSecretKey()))
+				scs[i] = utils.OrPanic1(conn.MakeSecretConnection(ctx, c))
 				return nil
 			})
 		}
@@ -53,7 +48,7 @@ func transform(msg []byte) []byte {
 }
 
 func runServer(ctx context.Context, rng utils.Rng, mux *Mux) error {
-	return scope.Run(ctx, func(ctx context.Context, s scope.Scope) error {
+	return utils.IgnoreCancel(scope.Run(ctx, func(ctx context.Context, s scope.Scope) error {
 		for kind := range mux.cfg.Kinds {
 			_ = rng.Split()
 			s.Spawn(func() error {
@@ -94,7 +89,7 @@ func runServer(ctx context.Context, rng utils.Rng, mux *Mux) error {
 			})
 		}
 		return nil
-	})
+	}))
 }
 
 type clientSet struct {
@@ -230,13 +225,13 @@ func runClients(ctx context.Context, rng utils.Rng, mux *Mux) error {
 func TestHappyPath(t *testing.T) {
 	rng := utils.TestRng()
 	kindCount := 5
-	c1, c2 := testConn(t)
 	err := scope.Run(t.Context(), func(ctx context.Context, s scope.Scope) error {
+		c1, c2 := spawnBgPipe(ctx, s)
 		for _, c := range utils.Slice(c1, c2) {
 			mux := makeMux(rng, kindCount)
 			serverRng := rng.Split()
-			s.SpawnBgNamed("mux", func() error { return ignoreDisconnect(mux.Run(ctx, c)) })
-			s.SpawnBgNamed("server", func() error { return utils.IgnoreCancel(runServer(ctx, serverRng, mux)) })
+			s.SpawnBgNamed("mux", func() error { return mux.Run(ctx, c) })
+			s.SpawnBgNamed("server", func() error { return runServer(ctx, serverRng, mux) })
 			clientRng := rng.Split()
 			s.SpawnNamed("client", func() error { return runClients(ctx, clientRng, mux) })
 		}
@@ -281,15 +276,15 @@ func makeConfig(kinds ...StreamKind) *Config {
 }
 
 func TestStreamKindsMismatch(t *testing.T) {
-	c1, c2 := testConn(t)
 	err := scope.Run(t.Context(), func(ctx context.Context, s scope.Scope) error {
+		c1, c2 := spawnBgPipe(ctx, s)
 		var k0, k1, k2 StreamKind = 0, 1, 2
 		muxs := utils.Slice(
 			NewMux(makeConfig(k0, k1)),
 			NewMux(makeConfig(k1, k2)),
 		)
 		for i, c := range utils.Slice(c1, c2) {
-			s.SpawnBg(func() error { return ignoreDisconnect(muxs[i].Run(ctx, c)) })
+			s.SpawnBg(func() error { return muxs[i].Run(ctx, c) })
 		}
 		// Connecting/accepting of unconfigured kind should error.
 		if _, err := muxs[0].Connect(ctx, k2, 10, 10); !errors.Is(err, errUnknownKind) {
@@ -346,8 +341,8 @@ func TestStreamKindsMismatch(t *testing.T) {
 // Test checking that closing a stream does not drop messages on the floor:
 // sending and receiving still works as long as messages fit into a window.
 func TestClosedStream(t *testing.T) {
-	c1, c2 := testConn(t)
 	err := scope.Run(t.Context(), func(ctx context.Context, s scope.Scope) error {
+		c1, c2 := spawnBgPipe(ctx, s)
 		kind := StreamKind(0)
 		window := uint64(4)
 		msg := []byte("hello")
@@ -356,7 +351,7 @@ func TestClosedStream(t *testing.T) {
 			NewMux(makeConfig(kind)),
 		)
 		for i, c := range utils.Slice(c1, c2) {
-			s.SpawnBg(func() error { return ignoreDisconnect(muxs[i].Run(ctx, c)) })
+			s.SpawnBg(func() error { return muxs[i].Run(ctx, c) })
 		}
 		s.Spawn(func() error {
 			// Just accept a single stream and close immediately.
@@ -407,14 +402,14 @@ func TestClosedStream(t *testing.T) {
 }
 
 func TestProtocol_TooLargeMsg(t *testing.T) {
-	c1, c2 := testConn(t)
 	err := scope.Run(t.Context(), func(ctx context.Context, s scope.Scope) error {
+		c1, c2 := spawnBgPipe(ctx, s)
 		kind := StreamKind(0)
 		maxMsgSize := uint64(10)
 
 		// Bad mux.
 		badMux := NewMux(makeConfig(kind))
-		s.SpawnBg(func() error { return ignoreDisconnect(badMux.Run(ctx, c1)) })
+		s.SpawnBg(func() error { return badMux.Run(ctx, c1) })
 		s.SpawnBg(func() error {
 			t.Log("Connect stream.")
 			stream, err := badMux.Connect(ctx, kind, 0, 0)
@@ -455,15 +450,15 @@ func TestProtocol_TooLargeMsg(t *testing.T) {
 }
 
 func TestProtocol_TooManyMsgs(t *testing.T) {
-	c1, c2 := testConn(t)
 	err := scope.Run(t.Context(), func(ctx context.Context, s scope.Scope) error {
+		c1, c2 := spawnBgPipe(ctx, s)
 		kind := StreamKind(0)
 		maxMsgSize := uint64(10)
 		window := uint64(3)
 
 		// Bad mux.
 		badMux := NewMux(makeConfig(kind))
-		s.SpawnBg(func() error { return ignoreDisconnect(badMux.Run(ctx, c1)) })
+		s.SpawnBg(func() error { return badMux.Run(ctx, c1) })
 		s.SpawnBg(func() error {
 			t.Log("Connect stream.")
 			stream, err := badMux.Connect(ctx, kind, 0, 0)
@@ -508,14 +503,14 @@ func TestProtocol_TooManyMsgs(t *testing.T) {
 }
 
 func TestProtocol_FrameAfterClose(t *testing.T) {
-	c1, c2 := testConn(t)
 	err := scope.Run(t.Context(), func(ctx context.Context, s scope.Scope) error {
+		c1, c2 := spawnBgPipe(ctx, s)
 		kind := StreamKind(0)
 
 		// Bad mux.
 		badMux := NewMux(makeConfig(kind))
 		maxMsgSize := uint64(10)
-		s.SpawnBg(func() error { return ignoreDisconnect(badMux.Run(ctx, c1)) })
+		s.SpawnBg(func() error { return badMux.Run(ctx, c1) })
 		s.SpawnBg(func() error {
 			t.Log("Connect stream.")
 			stream, err := badMux.Connect(ctx, kind, 0, 0)
@@ -566,13 +561,13 @@ func TestProtocol_FrameAfterClose(t *testing.T) {
 }
 
 func TestProtocol_TooManyAccepts(t *testing.T) {
-	c1, c2 := testConn(t)
 	err := scope.Run(t.Context(), func(ctx context.Context, s scope.Scope) error {
+		c1, c2 := spawnBgPipe(ctx, s)
 		kind := StreamKind(0)
 
 		// Bad mux.
 		badMux := NewMux(makeConfig(kind))
-		s.SpawnBg(func() error { return ignoreDisconnect(badMux.Run(ctx, c1)) })
+		s.SpawnBg(func() error { return badMux.Run(ctx, c1) })
 		s.SpawnBg(func() error {
 			// Artificially connect too many streams.
 			for queue, ctrl := range badMux.queue.Lock() {
@@ -599,13 +594,13 @@ func TestProtocol_TooManyAccepts(t *testing.T) {
 }
 
 func TestProtocol_UnknownStream(t *testing.T) {
-	c1, c2 := testConn(t)
 	err := scope.Run(t.Context(), func(ctx context.Context, s scope.Scope) error {
+		c1, c2 := spawnBgPipe(ctx, s)
 		kind := StreamKind(0)
 
 		// Bad mux.
 		badMux := NewMux(makeConfig(kind))
-		s.SpawnBg(func() error { return ignoreDisconnect(badMux.Run(ctx, c1)) })
+		s.SpawnBg(func() error { return badMux.Run(ctx, c1) })
 		s.SpawnBg(func() error {
 			// Artificially accept a stream without connect.
 			for queue, ctrl := range badMux.queue.Lock() {
@@ -630,14 +625,14 @@ func TestProtocol_UnknownStream(t *testing.T) {
 }
 
 func TestProtocol_UnknownKind(t *testing.T) {
-	c1, c2 := testConn(t)
 	err := scope.Run(t.Context(), func(ctx context.Context, s scope.Scope) error {
+		c1, c2 := spawnBgPipe(ctx, s)
 		kind := StreamKind(0)
 		badKind := StreamKind(1)
 
 		// Bad mux.
 		badMux := NewMux(makeConfig(kind))
-		s.SpawnBg(func() error { return ignoreDisconnect(badMux.Run(ctx, c1)) })
+		s.SpawnBg(func() error { return badMux.Run(ctx, c1) })
 		s.SpawnBg(func() error {
 			// Artificially connect a stream of unknown kind.
 			for queue, ctrl := range badMux.queue.Lock() {