// rpcCallTablet wil execute the RPC on the remote server. func (client *GoRpcTabletManagerClient) rpcCallTablet(ctx context.Context, tablet *topo.TabletInfo, name string, args, reply interface{}) error { // create the RPC client, using ctx.Deadline if set, or no timeout. var connectTimeout time.Duration deadline, ok := ctx.Deadline() if ok { connectTimeout = deadline.Sub(time.Now()) if connectTimeout < 0 { return fmt.Errorf("Timeout connecting to TabletManager.%v on %v", name, tablet.Alias) } } rpcClient, err := bsonrpc.DialHTTP("tcp", tablet.Addr(), connectTimeout, nil) if err != nil { return fmt.Errorf("RPC error for %v: %v", tablet.Alias, err.Error()) } defer rpcClient.Close() // use the context Done() channel. Will handle context timeout. call := rpcClient.Go(ctx, "TabletManager."+name, args, reply, nil) select { case <-ctx.Done(): return fmt.Errorf("Timeout waiting for TabletManager.%v to %v", name, tablet.Alias) case <-call.Done: if call.Error != nil { return fmt.Errorf("Remote error for %v: %v", tablet.Alias, call.Error.Error()) } else { return nil } } }
// getConn reuses an existing connection if possible. Otherwise // it returns a connection which it will save for future reuse. // If it returns an error, retry will tell you if getConn can be retried. // If the context has a deadline and exceeded, it returns error and no-retry immediately. func (sdc *ShardConn) getConn(ctx context.Context) (conn tabletconn.TabletConn, endPoint topo.EndPoint, err error, retry bool) { sdc.mu.Lock() defer sdc.mu.Unlock() // fail-fast if deadline exceeded deadline, ok := ctx.Deadline() if ok { if time.Now().After(deadline) { return nil, topo.EndPoint{}, tabletconn.OperationalError("vttablet: deadline exceeded"), false } } if sdc.conn != nil { return sdc.conn, sdc.conn.EndPoint(), nil, false } endPoint, err = sdc.balancer.Get() if err != nil { return nil, topo.EndPoint{}, err, false } conn, err = tabletconn.GetDialer()(ctx, endPoint, sdc.keyspace, sdc.shard, sdc.timeout) if err != nil { sdc.balancer.MarkDown(endPoint.Uid, err.Error()) return nil, endPoint, err, true } sdc.conn = conn return sdc.conn, endPoint, nil, false }
// Call sends the request, waits for it to complete, and returns its error status. func (client *Client) Call(ctx context.Context, req proto.Message, resp proto.Message) error { // Setup a new call call := &Call{ Req: req, Resp: resp, cancelc: make(chan interface{}), } // Make sure a service name is passed if serviceName, ok := ServiceNameFromContext(ctx); !ok { return ErrServiceNameMissing } else { call.Service = serviceName } // If a manual deadline is not passed, setup with default timeout if _, ok := ctx.Deadline(); !ok { ctx, _ = context.WithDeadline(ctx, time.Now().Add(RequestTimeout)) } // Run the request in a goroutine and write the reponse to c c := make(chan error, 1) go func() { c <- client.send(ctx, call) }() select { // Use context done to manually trigger cancel case <-ctx.Done(): // Cancel the request in flight call.cancelc <- true <-c // Wait for the request to finish if ctx.Err() == context.DeadlineExceeded { return ErrTimeout } return ctx.Err() // Request finished case err := <-c: return err } }
func (client *Client) send(ctx context.Context, call *Call) error { if client.shutdown || client.closing { return ErrShutdown } // Every call gets its own req socket sock, err := client.ctx.NewSocket(zmq.REQ) if err != nil { return err } defer sock.Close() // Connect it to the router if err := sock.Connect(RouterURL); err != nil { return err } // Marshal the outgoing message msgBytes, err := proto.Marshal(call.Req) if err != nil { return err } // Envelope the message reqID := uuid.NewV4() envelope := &Request{ UUID: reqID.Bytes(), Path: proto.String(fmt.Sprintf("zrpc://%s/%s", call.Service, util.GetMessageName(call.Req))), Payload: msgBytes, } // If request has a timeout, send it in the request envelope d, ok := ctx.Deadline() if ok { envelope.Expires = proto.Int64(d.Unix()) } // Marshal the outgoing envelope envBytes, err := proto.Marshal(envelope) if err != nil { return err } // Send it if _, err := sock.SendBytes(envBytes, 0); err != nil { return err } handleResponse := func(state zmq.State) error { respBytes, err := sock.RecvBytes(0) if err != nil { return err } // Unmarshal the response envelope resp := &Response{} if err := proto.Unmarshal(respBytes, resp); err != nil { return err } // Make sure the same message ID was received respID, err := uuid.FromBytes(resp.UUID) if err != nil || !uuid.Equal(reqID, respID) { glog.Errorf("Mismatching message IDs, sent '%s', got '%s'", reqID, respID) return ErrMessageMismatch } // Check if there is an error if resp.Error != nil { return NewClientError(resp.Error.GetMessage(), int(resp.GetStatusCode())) } // Decode the actual message (if exists and wanted) if call.Resp != nil && len(resp.Payload) > 0 { if err := proto.Unmarshal(resp.Payload, call.Resp); err != nil { return err } } // This is insane, but the reactor runs until an error is returned... return errDone } // Use a reactor to be able to utilize the calls cancelc reactor := zmq.NewReactor() reactor.AddSocket(sock, zmq.POLLIN, handleResponse) reactor.AddChannel(call.cancelc, 1, func(interface{}) error { return ErrCancel }) // Poll for a short interval to be able to return to the channel handling if err := reactor.Run(time.Millisecond * 50); err != errDone { return err } return nil }