itd/api/drpc.go
Arsen Musayelyan 0cf36f220d
All checks were successful
ci/woodpecker/push/woodpecker Pipeline was successful
Properly close multiplexed streams
2023-01-03 09:15:38 -08:00

82 lines
1.3 KiB
Go

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)
defer dconn.Close()
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)
go func() {
<-dconn.Closed()
conn.Close()
}()
s, err := dconn.NewStream(ctx, rpc, enc)
if err != nil {
return nil, err
}
go func() {
<-s.Context().Done()
dconn.Close()
}()
return s, nil
}