From 2cecf7f9c7af111b6878d08521ec81e213269cb5 Mon Sep 17 00:00:00 2001
From: Sam Kleinman <garen@tychoish.com>
Date: Mon, 16 May 2022 10:31:55 -0400
Subject: [PATCH] p2p: cache libp2p streams (#8562)

---
 internal/p2p/channel.go | 51 ++++++++++++++++++++++++++++++++++-------
 1 file changed, 43 insertions(+), 8 deletions(-)

diff --git a/internal/p2p/channel.go b/internal/p2p/channel.go
index 47664e36e..d167ffe55 100644
--- a/internal/p2p/channel.go
+++ b/internal/p2p/channel.go
@@ -403,13 +403,13 @@ func (ch *libp2pChannelImpl) Send(ctx context.Context, e Envelope) error {
 		e.Message = msg
 	}
 
-	e.From = types.NodeID(ch.host.ID())
-	bz, err := proto.Marshal(e.Message)
-	if err != nil {
-		return err
-	}
-
 	if e.Broadcast {
+		e.From = types.NodeID(ch.host.ID())
+		bz, err := proto.Marshal(e.Message)
+		if err != nil {
+			return err
+		}
+
 		return ch.topic.Publish(ctx, bz)
 	}
 
@@ -418,12 +418,13 @@ func (ch *libp2pChannelImpl) Send(ctx context.Context, e Envelope) error {
 		return fmt.Errorf("cannot connect to %q", e.To)
 	default:
 		// TODO: should we should attempt to reuse between peers?
-		stream, err := ch.host.NewStream(ctx, peer.ID(e.To), protocol.ID(ch.canonicalizedTopicName()))
+		stream, err := ch.getStream(ctx, peer.ID(e.To))
 		if err != nil {
 			return err
 		}
 
-		_, err = stream.Write(bz)
+		writer := protoio.NewDelimitedWriter(bufio.NewWriter(stream))
+		_, err = writer.WriteMsg(e.Message)
 		if err != nil {
 			return err
 		}
@@ -431,6 +432,40 @@ func (ch *libp2pChannelImpl) Send(ctx context.Context, e Envelope) error {
 	return nil
 }
 
+func (ch *libp2pChannelImpl) getStream(ctx context.Context, peer peer.ID) (network.Stream, error) {
+	conns := ch.host.Network().ConnsToPeer(peer)
+	pid := protocol.ID(ch.canonicalizedTopicName())
+	if len(conns) > 0 {
+		for cidx := range conns {
+			streams := conns[cidx].GetStreams()
+			for sidx := range streams {
+				stream := streams[sidx]
+				if stream.Protocol() == pid && stream.Stat().Direction == network.DirOutbound {
+					return stream, nil
+				}
+			}
+		}
+
+	}
+
+	// TODO: these context should be a broader context rather than
+	// the one from send to avoid shutting down the stream when
+	// the context that's sending the message closes. Right now
+	// they're functionally synonymous so it's safe, but it'd be
+	// possible to introduce an unintended bug here.
+	conn, err := ch.host.Network().DialPeer(ctx, peer)
+	if err != nil {
+		return nil, err
+	}
+
+	stream, err := ch.host.NewStream(ctx, conn.RemotePeer(), pid)
+	if err != nil {
+		return nil, err
+	}
+
+	return stream, nil
+}
+
 func (ch *libp2pChannelImpl) SendError(ctx context.Context, pe PeerError) error {
 	if errors.Is(pe.Err, context.Canceled) || errors.Is(pe.Err, context.DeadlineExceeded) || ctx.Err() != nil {
 		return nil