mirror of
https://github.com/coder/coder.git
synced 2025-07-09 11:45:56 +00:00
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.
130 lines
2.8 KiB
Go
130 lines
2.8 KiB
Go
package drpc
|
|
|
|
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
|
|
}
|