Add connection multiplexing, fixing itgui

api/api.go

@@ -5,13 +5,14 @@ import (
 	"net"
 
 	"go.arsenm.dev/itd/internal/rpc"
+	"storj.io/drpc"
 	"storj.io/drpc/drpcconn"
 )
 
 const DefaultAddr = "/tmp/itd/socket"
 
 type Client struct {
-	conn   *drpcconn.Conn
+	conn   drpc.Conn
 	client rpc.DRPCITDClient
 }
 
@@ -20,11 +21,15 @@ func New(sockPath string) (*Client, error) {
 	if err != nil {
 		return nil, err
 	}
-	dconn := drpcconn.New(conn)
+
+	mconn, err := newMuxConn(conn)
+	if err != nil {
+		return nil, err
+	}
 
 	return &Client{
-		conn:   dconn,
-		client: rpc.NewDRPCITDClient(dconn),
+		conn:   mconn,
+		client: rpc.NewDRPCITDClient(mconn),
 	}, nil
 }
 

api/drpc.go

@@ -0,0 +1,65 @@
+package api
+
+import (
+	"context"
+	"io"
+
+	"github.com/hashicorp/yamux"
+	"storj.io/drpc"
+	"storj.io/drpc/drpcconn"
+)
+
+var _ drpc.Conn = &muxConn{}
+
+type muxConn struct {
+	conn   io.ReadWriteCloser
+	sess   *yamux.Session
+	closed chan struct{}
+}
+
+func newMuxConn(conn io.ReadWriteCloser) (*muxConn, error) {
+	sess, err := yamux.Client(conn, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	return &muxConn{
+		conn:   conn,
+		sess:   sess,
+		closed: make(chan struct{}),
+	}, nil
+}
+
+func (m *muxConn) Close() error {
+	defer close(m.closed)
+
+	err := m.sess.Close()
+	if err != nil {
+		return err
+	}
+	return m.conn.Close()
+}
+
+func (m *muxConn) Closed() <-chan struct{} {
+	return m.closed
+}
+
+func (m *muxConn) Invoke(ctx context.Context, rpc string, enc drpc.Encoding, in, out drpc.Message) error {
+	conn, err := m.sess.Open()
+	if err != nil {
+		return err
+	}
+	defer conn.Close()
+	dconn := drpcconn.New(conn)
+	return dconn.Invoke(ctx, rpc, enc, in, out)
+}
+
+func (m *muxConn) NewStream(ctx context.Context, rpc string, enc drpc.Encoding) (drpc.Stream, error) {
+	conn, err := m.sess.Open()
+	if err != nil {
+		return nil, err
+	}
+
+	dconn := drpcconn.New(conn)
+	return dconn.NewStream(ctx, rpc, enc)
+}