chore: move drpc transport tools to codersdk/drpc (#11224)

Part of #10532

DRPC transport over yamux and in-mem pipes was previously only used on the provisioner APIs, but now will also be used in tailnet.  Moved to subpackage of codersdk to avoid import loops.

provisionersdk/serve_test.go

@@ -10,6 +10,7 @@ import (
 	"go.uber.org/goleak"
 	"storj.io/drpc/drpcconn"
 
+	"github.com/coder/coder/v2/codersdk/drpc"
 	"github.com/coder/coder/v2/provisionersdk"
 	"github.com/coder/coder/v2/provisionersdk/proto"
 	"github.com/coder/coder/v2/testutil"
@@ -23,7 +24,7 @@ func TestProvisionerSDK(t *testing.T) {
 	t.Parallel()
 	t.Run("ServeListener", func(t *testing.T) {
 		t.Parallel()
-		client, server := provisionersdk.MemTransportPipe()
+		client, server := drpc.MemTransportPipe()
 		defer client.Close()
 		defer server.Close()
 
@@ -65,7 +66,7 @@ func TestProvisionerSDK(t *testing.T) {
 
 	t.Run("ServeClosedPipe", func(t *testing.T) {
 		t.Parallel()
-		client, server := provisionersdk.MemTransportPipe()
+		client, server := drpc.MemTransportPipe()
 		_ = client.Close()
 		_ = server.Close()
 

provisionersdk/transport.go

@@ -1,129 +0,0 @@
-package provisionersdk
-
-import (
-	"context"
-	"net"
-	"sync"
-
-	"github.com/hashicorp/yamux"
-	"github.com/valyala/fasthttp/fasthttputil"
-	"storj.io/drpc"
-	"storj.io/drpc/drpcconn"
-
-	"github.com/coder/coder/v2/coderd/tracing"
-)
-
-const (
-	// MaxMessageSize is the maximum payload size that can be
-	// transported without error.
-	MaxMessageSize = 4 << 20
-)
-
-// MultiplexedConn returns a multiplexed dRPC connection from a yamux Session.
-func MultiplexedConn(session *yamux.Session) drpc.Conn {
-	return &multiplexedDRPC{session}
-}
-
-// Allows concurrent requests on a single dRPC connection.
-// Required for calling functions concurrently.
-type multiplexedDRPC struct {
-	session *yamux.Session
-}
-
-func (m *multiplexedDRPC) Close() error {
-	return m.session.Close()
-}
-
-func (m *multiplexedDRPC) Closed() <-chan struct{} {
-	return m.session.CloseChan()
-}
-
-func (m *multiplexedDRPC) Invoke(ctx context.Context, rpc string, enc drpc.Encoding, inMessage, outMessage drpc.Message) error {
-	conn, err := m.session.Open()
-	if err != nil {
-		return err
-	}
-	dConn := drpcconn.New(conn)
-	defer func() {
-		_ = dConn.Close()
-	}()
-	return dConn.Invoke(ctx, rpc, enc, inMessage, outMessage)
-}
-
-func (m *multiplexedDRPC) NewStream(ctx context.Context, rpc string, enc drpc.Encoding) (drpc.Stream, error) {
-	conn, err := m.session.Open()
-	if err != nil {
-		return nil, err
-	}
-	dConn := drpcconn.New(conn)
-	stream, err := dConn.NewStream(ctx, rpc, enc)
-	if err == nil {
-		go func() {
-			<-stream.Context().Done()
-			_ = dConn.Close()
-		}()
-	}
-	return stream, err
-}
-
-func MemTransportPipe() (drpc.Conn, net.Listener) {
-	m := &memDRPC{
-		closed: make(chan struct{}),
-		l:      fasthttputil.NewInmemoryListener(),
-	}
-
-	return m, m.l
-}
-
-type memDRPC struct {
-	closeOnce sync.Once
-	closed    chan struct{}
-	l         *fasthttputil.InmemoryListener
-}
-
-func (m *memDRPC) Close() error {
-	err := m.l.Close()
-	m.closeOnce.Do(func() { close(m.closed) })
-	return err
-}
-
-func (m *memDRPC) Closed() <-chan struct{} {
-	return m.closed
-}
-
-func (m *memDRPC) Invoke(ctx context.Context, rpc string, enc drpc.Encoding, inMessage, outMessage drpc.Message) error {
-	conn, err := m.l.Dial()
-	if err != nil {
-		return err
-	}
-
-	dConn := &tracing.DRPCConn{Conn: drpcconn.New(conn)}
-	defer func() {
-		_ = dConn.Close()
-		_ = conn.Close()
-	}()
-	return dConn.Invoke(ctx, rpc, enc, inMessage, outMessage)
-}
-
-func (m *memDRPC) NewStream(ctx context.Context, rpc string, enc drpc.Encoding) (drpc.Stream, error) {
-	conn, err := m.l.Dial()
-	if err != nil {
-		return nil, err
-	}
-	dConn := &tracing.DRPCConn{Conn: drpcconn.New(conn)}
-	stream, err := dConn.NewStream(ctx, rpc, enc)
-	if err != nil {
-		_ = dConn.Close()
-		_ = conn.Close()
-		return nil, err
-	}
-	go func() {
-		select {
-		case <-stream.Context().Done():
-		case <-m.closed:
-		}
-		_ = dConn.Close()
-		_ = conn.Close()
-	}()
-	return stream, nil
-}