diff --git a/net/peer/peer.go b/net/peer/peer.go index 9c8e0ddc..879f6f9b 100644 --- a/net/peer/peer.go +++ b/net/peer/peer.go @@ -41,6 +41,7 @@ type Peer interface { AcquireDrpcConn(ctx context.Context) (drpc.Conn, error) ReleaseDrpcConn(conn drpc.Conn) + DoDrpc(ctx context.Context, do func(conn drpc.Conn) error) error IsClosed() bool @@ -95,6 +96,15 @@ func (p *peer) ReleaseDrpcConn(conn drpc.Conn) { return } +func (p *peer) DoDrpc(ctx context.Context, do func(conn drpc.Conn) error) error { + conn, err := p.AcquireDrpcConn(ctx) + if err != nil { + return err + } + defer p.ReleaseDrpcConn(conn) + return do(conn) +} + func (p *peer) acceptLoop() { var exitErr error defer func() {