Example #1
0
// handleRPCConn reads a command from the connection and executes it.
func (r *rpc) handleRPCConn(conn net.Conn) {
	defer conn.Close()
	// RPC connections should execute on the leader.  If we are not the leader,
	// proxy the connection to the leader so that clients an connect to any node
	// in the cluster.
	r.traceCluster("rpc connection from: %v", conn.RemoteAddr())

	if !r.store.IsLeader() {
		r.proxyLeader(conn.(*net.TCPConn))
		return
	}

	// Read and execute request.
	typ, resp, err := func() (internal.RPCType, proto.Message, error) {
		// Read request size.
		var sz uint64
		if err := binary.Read(conn, binary.BigEndian, &sz); err != nil {
			return internal.RPCType_Error, nil, fmt.Errorf("read size: %s", err)
		}

		if sz == 0 {
			return 0, nil, fmt.Errorf("invalid message size: %d", sz)
		}

		if sz >= MaxMessageSize {
			return 0, nil, fmt.Errorf("max message size of %d exceeded: %d", MaxMessageSize, sz)
		}

		// Read request.
		buf := make([]byte, sz)
		if _, err := io.ReadFull(conn, buf); err != nil {
			return internal.RPCType_Error, nil, fmt.Errorf("read request: %s", err)
		}

		// Determine the RPC type
		rpcType := internal.RPCType(btou64(buf[0:8]))
		buf = buf[8:]

		r.traceCluster("recv %v request on: %v", rpcType, conn.RemoteAddr())
		switch rpcType {
		case internal.RPCType_FetchData:
			var req internal.FetchDataRequest
			if err := proto.Unmarshal(buf, &req); err != nil {
				return internal.RPCType_Error, nil, fmt.Errorf("fetch request unmarshal: %v", err)
			}
			resp, err := r.handleFetchData(&req)
			return rpcType, resp, err
		case internal.RPCType_Join:
			var req internal.JoinRequest
			if err := proto.Unmarshal(buf, &req); err != nil {
				return internal.RPCType_Error, nil, fmt.Errorf("join request unmarshal: %v", err)
			}
			resp, err := r.handleJoinRequest(&req)
			return rpcType, resp, err
		default:
			return internal.RPCType_Error, nil, fmt.Errorf("unknown rpc type:%v", rpcType)
		}
	}()

	// Handle unexpected RPC errors
	if err != nil {
		resp = &internal.ErrorResponse{
			Header: &internal.ResponseHeader{
				OK: proto.Bool(false),
			},
		}
		typ = internal.RPCType_Error
	}

	// Set the status header and error message
	if reply, ok := resp.(Reply); ok {
		reply.GetHeader().OK = proto.Bool(err == nil)
		if err != nil {
			reply.GetHeader().Error = proto.String(err.Error())
		}
	}

	r.sendResponse(conn, typ, resp)
}
Example #2
0
// call sends an encoded request to the remote leader and returns
// an encoded response value.
func (r *rpc) call(dest string, req proto.Message) (proto.Message, error) {
	// Determine type of request
	var rpcType internal.RPCType
	switch t := req.(type) {
	case *internal.JoinRequest:
		rpcType = internal.RPCType_Join
	case *internal.FetchDataRequest:
		rpcType = internal.RPCType_FetchData
	default:
		return nil, fmt.Errorf("unknown rpc request type: %v", t)
	}

	// Create a connection to the leader.
	conn, err := net.DialTimeout("tcp", dest, leaderDialTimeout)
	if err != nil {
		return nil, err
	}
	defer conn.Close()

	// Write a marker byte for rpc messages.
	_, err = conn.Write([]byte{MuxRPCHeader})
	if err != nil {
		return nil, err
	}

	b, err := proto.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("rpc marshal: %v", err)
	}

	// Write request size & bytes.
	if _, err := conn.Write(r.pack(rpcType, b)); err != nil {
		return nil, fmt.Errorf("write %v rpc: %s", rpcType, err)
	}

	data, err := ioutil.ReadAll(conn)
	if err != nil {
		return nil, fmt.Errorf("read %v rpc: %v", rpcType, err)
	}

	// Should always have a size and type
	if exp := 16; len(data) < exp {
		return nil, fmt.Errorf("rpc %v failed: short read: got %v, exp %v", rpcType, len(data), exp)
	}

	sz := btou64(data[0:8])
	if len(data[8:]) != int(sz) {
		return nil, fmt.Errorf("rpc %v failed: short read: got %v, exp %v", rpcType, len(data[8:]), sz)
	}

	// See what response type we got back, could get a general error response
	rpcType = internal.RPCType(btou64(data[8:16]))
	data = data[16:]

	var resp proto.Message
	switch rpcType {
	case internal.RPCType_Join:
		resp = &internal.JoinResponse{}
	case internal.RPCType_FetchData:
		resp = &internal.FetchDataResponse{}
	case internal.RPCType_Error:
		resp = &internal.ErrorResponse{}
	default:
		return nil, fmt.Errorf("unknown rpc response type: %v", rpcType)
	}

	if err := proto.Unmarshal(data, resp); err != nil {
		return nil, fmt.Errorf("rpc unmarshal: %v", err)
	}

	if reply, ok := resp.(Reply); ok {
		if !reply.GetHeader().GetOK() {
			return nil, fmt.Errorf("rpc %v failed: %s", rpcType, reply.GetHeader().GetError())
		}
	}

	return resp, nil
}