fix: timeout guarantee for RequestOnCnx (#492)
Signed-off-by: jonyhy96 <hy352144278@gmail.com>
Fixes #491
### Motivation
see #491
### Modifications
add timeout guarantee to RequestOnCnx
### Verifying this change
- [x] Make sure that the change passes the CI checks.
diff --git a/pulsar/internal/rpc_client.go b/pulsar/internal/rpc_client.go
index b51e4e3..1374747 100644
--- a/pulsar/internal/rpc_client.go
+++ b/pulsar/internal/rpc_client.go
@@ -21,7 +21,6 @@
"errors"
"net"
"net/url"
- "sync"
"sync/atomic"
"time"
@@ -31,6 +30,16 @@
"github.com/gogo/protobuf/proto"
)
+var (
+ // ErrRequestTimeOut happens when request not finished in given requestTimeout.
+ ErrRequestTimeOut = errors.New("request timed out")
+)
+
+type result struct {
+ *RPCResult
+ error
+}
+
type RPCResult struct {
Response *pb.BaseCommand
Cnx Connection
@@ -121,14 +130,10 @@
return nil, err
}
- type Res struct {
- *RPCResult
- error
- }
- ch := make(chan Res, 10)
+ ch := make(chan result, 1)
cnx.SendRequest(requestID, baseCommand(cmdType, message), func(response *pb.BaseCommand, err error) {
- ch <- Res{&RPCResult{
+ ch <- result{&RPCResult{
Cnx: cnx,
Response: response,
}, err}
@@ -139,29 +144,30 @@
case res := <-ch:
return res.RPCResult, res.error
case <-time.After(c.requestTimeout):
- return nil, errors.New("request timed out")
+ return nil, ErrRequestTimeOut
}
}
func (c *rpcClient) RequestOnCnx(cnx Connection, requestID uint64, cmdType pb.BaseCommand_Type,
message proto.Message) (*RPCResult, error) {
c.metrics.RPCRequestCount.Inc()
- wg := sync.WaitGroup{}
- wg.Add(1)
- rpcResult := &RPCResult{
- Cnx: cnx,
- }
+ ch := make(chan result, 1)
- var rpcErr error
cnx.SendRequest(requestID, baseCommand(cmdType, message), func(response *pb.BaseCommand, err error) {
- rpcResult.Response = response
- rpcErr = err
- wg.Done()
+ ch <- result{&RPCResult{
+ Cnx: cnx,
+ Response: response,
+ }, err}
+ close(ch)
})
- wg.Wait()
- return rpcResult, rpcErr
+ select {
+ case res := <-ch:
+ return res.RPCResult, res.error
+ case <-time.After(c.requestTimeout):
+ return nil, ErrRequestTimeOut
+ }
}
func (c *rpcClient) RequestOnCnxNoWait(cnx Connection, cmdType pb.BaseCommand_Type, message proto.Message) error {
