func (s *testClientSuite) TestWrongMessageID(c *C) {
l := startServer(":61237", c, func(conn net.Conn, c *C) {
var msg msgpb.Message
msgID, err := util.ReadMessage(conn, &msg)
c.Assert(err, IsNil)
resp := msgpb.Message{
MsgType: msgpb.MessageType_KvResp,
KvResp: &pb.Response{
Type: msg.GetKvReq().GetType(),
},
}
// Send the request back to client, set wrong msgID for the 1st
// request.
if msgID == 1 {
err = util.WriteMessage(conn, msgID+100, &resp)
} else {
err = util.WriteMessage(conn, msgID, &resp)
}
c.Assert(err, IsNil)
})
defer l.Close()
cli := newRPCClient()
req := &pb.Request{
Type: pb.MessageType_CmdGet,
}
// Wrong ID for the first request, correct for the rests.
_, err := cli.SendKVReq(":61237", req, readTimeoutShort)
c.Assert(err, NotNil)
resp, err := cli.SendKVReq(":61237", req, readTimeoutShort)
c.Assert(err, IsNil)
c.Assert(resp.GetType(), Equals, req.GetType())
}
func readThenCloseRequest(conn net.Conn, c *C) {
c.Assert(conn, NotNil)
defer conn.Close()
var msg msgpb.Message
msgID, err := util.ReadMessage(conn, &msg)
c.Assert(err, IsNil)
c.Assert(msg.GetKvReq(), NotNil)
c.Assert(msgID, Greater, uint64(0))
}
// SendCopReq sends a Request to co-processor and receives Response.
func (c *rpcClient) SendCopReq(req *coprocessor.Request) (*coprocessor.Response, error) {
conn, err := c.pool.Get()
if err != nil {
return nil, errors.Trace(err)
}
defer conn.Close()
msg := msgpb.Message{
MsgType: msgpb.MessageType_CopReq.Enum(),
CopReq: req,
}
err = c.doSend(conn, &msg)
if err == nil {
if msg.GetMsgType() != msgpb.MessageType_CopResp || msg.GetCopResp() == nil {
err = errors.Trace(errInvalidResponse)
}
}
if err != nil {
// This connection is not valid any more, so close its underlying conn.
if poolConn, ok := conn.(*pool.PoolConn); ok {
poolConn.MarkUnusable()
}
return nil, errors.Trace(err)
}
return msg.GetCopResp(), nil
}
// SendKVReq sends a Request to kv server and receives Response.
func (c *rpcClient) SendKVReq(addr string, req *kvrpcpb.Request, timeout time.Duration) (*kvrpcpb.Response, error) {
sendReqCounter.WithLabelValues("kv").Inc()
start := time.Now()
defer func() { sendReqHistogram.WithLabelValues("kv").Observe(time.Since(start).Seconds()) }()
conn, err := c.p.GetConn(addr)
if err != nil {
return nil, errors.Trace(err)
}
defer c.p.PutConn(conn)
msg := msgpb.Message{
MsgType: msgpb.MessageType_KvReq,
KvReq: req,
}
err = c.doSend(conn, &msg, writeTimeout, timeout)
if err != nil {
conn.Close()
return nil, errors.Trace(err)
}
if msg.GetMsgType() != msgpb.MessageType_KvResp || msg.GetKvResp() == nil {
conn.Close()
return nil, errors.Trace(errInvalidResponse)
}
return msg.GetKvResp(), nil
}
func (c *rpcClient) doSend(conn *Conn, msg *msgpb.Message) error {
curMsgID := atomic.AddUint64(&c.msgID, 1)
log.Debugf("Send request msgID[%d] type[%v]", curMsgID, msg.GetMsgType())
if err := util.WriteMessage(conn, curMsgID, msg); err != nil {
return errors.Trace(err)
}
if err := conn.Flush(); err != nil {
return errors.Trace(err)
}
msgID, err := util.ReadMessage(conn.BufioReader(), msg)
if err != nil {
return errors.Trace(err)
}
if curMsgID != msgID {
log.Errorf("Sent msgID[%d] mismatches recv msgID[%d]", curMsgID, msgID)
}
log.Debugf("Receive response msgID[%d] type[%v]", msgID, msg.GetMsgType())
return nil
}
// handleRequest receive Request then send empty Response back fill with same Type.
func handleRequest(conn net.Conn, c *C) {
c.Assert(conn, NotNil)
defer conn.Close()
var msg msgpb.Message
msgID, err := util.ReadMessage(conn, &msg)
c.Assert(err, IsNil)
c.Assert(msgID, Greater, uint64(0))
c.Assert(msg.GetMsgType(), Equals, msgpb.MessageType_KvReq)
req := msg.GetKvReq()
c.Assert(req, NotNil)
var resp pb.Response
resp.Type = req.Type
msg = msgpb.Message{
MsgType: msgpb.MessageType_KvResp,
KvResp: &resp,
}
err = util.WriteMessage(conn, msgID, &msg)
c.Assert(err, IsNil)
}
func (c *rpcClient) doSend(conn net.Conn, msg *msgpb.Message) error {
curMsgID := atomic.AddUint64(&c.msgID, 1)
log.Debugf("Send request msgID[%d] type[%v]", curMsgID, msg.GetMsgType())
if err := conn.SetWriteDeadline(time.Now().Add(writeTimeout)); err != nil {
log.Warn("Set write deadline failed, it may be blocked.")
}
if err := util.WriteMessage(conn, curMsgID, msg); err != nil {
return errors.Trace(err)
}
if err := conn.SetReadDeadline(time.Now().Add(readTimeout)); err != nil {
log.Warn("Set read deadline failed, it may be blocked.")
}
msgID, err := util.ReadMessage(conn, msg)
if err != nil {
return errors.Trace(err)
}
if curMsgID != msgID {
log.Errorf("Sent msgID[%d] mismatches recv msgID[%d]", curMsgID, msgID)
}
log.Debugf("Receive response msgID[%d] type[%v]", msgID, msg.GetMsgType())
return nil
}
func (c *rpcClient) doSend(conn *Conn, msg *msgpb.Message, writeTimeout time.Duration, readTimeout time.Duration) error {
curMsgID := atomic.AddUint64(&c.msgID, 1)
log.Debugf("Send request msgID[%d] type[%v]", curMsgID, msg.GetMsgType())
conn.SetWriteDeadline(time.Now().Add(writeTimeout))
if err := util.WriteMessage(conn, curMsgID, msg); err != nil {
return errors.Trace(err)
}
if err := conn.Flush(); err != nil {
return errors.Trace(err)
}
conn.SetReadDeadline(time.Now().Add(readTimeout))
msgID, err := util.ReadMessage(conn.BufioReader(), msg)
if err != nil {
return errors.Trace(err)
}
if curMsgID != msgID {
log.Errorf("Sent msgID[%d] mismatches recv msgID[%d]", curMsgID, msgID)
return errors.Trace(errInvalidResponse)
}
log.Debugf("Receive response msgID[%d] type[%v]", msgID, msg.GetMsgType())
return nil
}
// SendKVReq sends a Request to kv server and receives Response.
func (c *rpcClient) SendKVReq(addr string, req *kvrpcpb.Request) (*kvrpcpb.Response, error) {
conn, err := c.p.GetConn(addr)
if err != nil {
return nil, errors.Trace(err)
}
defer c.p.PutConn(conn)
msg := msgpb.Message{
MsgType: msgpb.MessageType_KvReq.Enum(),
KvReq: req,
}
err = c.doSend(conn, &msg)
if err != nil {
return nil, errors.Trace(err)
}
if msg.GetMsgType() != msgpb.MessageType_KvResp || msg.GetKvResp() == nil {
return nil, errors.Trace(errInvalidResponse)
}
return msg.GetKvResp(), nil
}
// SendCopReq sends a Request to co-processor and receives Response.
func (c *rpcClient) SendCopReq(addr string, req *coprocessor.Request) (*coprocessor.Response, error) {
conn, err := c.p.GetConn(addr)
if err != nil {
return nil, errors.Trace(err)
}
defer c.p.PutConn(conn)
msg := msgpb.Message{
MsgType: msgpb.MessageType_CopReq,
CopReq: req,
}
err = c.doSend(conn, &msg)
if err != nil {
conn.Close()
return nil, errors.Trace(err)
}
if msg.GetMsgType() != msgpb.MessageType_CopResp || msg.GetCopResp() == nil {
conn.Close()
return nil, errors.Trace(errInvalidResponse)
}
return msg.GetCopResp(), nil
}