Properly close multiplexed streams

This commit is contained in:
Elara 2023-01-03 09:15:38 -08:00
parent 94ec82c4a6
commit 86a77f6f1f

View File

@ -51,6 +51,7 @@ func (m *muxConn) Invoke(ctx context.Context, rpc string, enc drpc.Encoding, in,
}
defer conn.Close()
dconn := drpcconn.New(conn)
defer dconn.Close()
return dconn.Invoke(ctx, rpc, enc, in, out)
}
@ -59,7 +60,22 @@ func (m *muxConn) NewStream(ctx context.Context, rpc string, enc drpc.Encoding)
if err != nil {
return nil, err
}
dconn := drpcconn.New(conn)
return dconn.NewStream(ctx, rpc, enc)
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
}