Golang transport.Stream类代码示例

func (s *Server) processStreamingRPC(t transport.ServerTransport, stream *transport.Stream, srv *service, sd *StreamDesc) (err error) {
	ss := &serverStream{
		t:       t,
		s:       stream,
		p:       &parser{s: stream},
		codec:   s.opts.codec,
		tracing: EnableTracing,
	}
	if ss.tracing {
		ss.traceInfo.tr = trace.New("grpc.Recv."+methodFamily(stream.Method()), stream.Method())
		ss.traceInfo.firstLine.client = false
		ss.traceInfo.tr.LazyLog(&ss.traceInfo.firstLine, false)
		defer func() {
			ss.mu.Lock()
			if err != nil && err != io.EOF {
				ss.traceInfo.tr.LazyLog(&fmtStringer{"%v", []interface{}{err}}, true)
				ss.traceInfo.tr.SetError()
			}
			ss.traceInfo.tr.Finish()
			ss.traceInfo.tr = nil
			ss.mu.Unlock()
		}()
	}
	if appErr := sd.Handler(srv.server, ss); appErr != nil {
		if err, ok := appErr.(rpcError); ok {
			ss.statusCode = err.code
			ss.statusDesc = err.desc
		} else {
			ss.statusCode = convertCode(appErr)
			ss.statusDesc = appErr.Error()
		}
	}
	return t.WriteStatus(ss.s, ss.statusCode, ss.statusDesc)

}

func (s *Server) handleStream(t transport.ServerTransport, stream *transport.Stream) {
	sm := stream.Method()
	if sm != "" && sm[0] == '/' {
		sm = sm[1:]
	}
	pos := strings.LastIndex(sm, "/")
	if pos == -1 {
		if err := t.WriteStatus(stream, codes.InvalidArgument, fmt.Sprintf("malformed method name: %q", stream.Method())); err != nil {
			grpclog.Printf("grpc: Server.handleStream failed to write status: %v", err)
		}
		return
	}
	service := sm[:pos]
	method := sm[pos+1:]
	srv, ok := s.m[service]
	if !ok {
		if err := t.WriteStatus(stream, codes.Unimplemented, fmt.Sprintf("unknown service %v", service)); err != nil {
			grpclog.Printf("grpc: Server.handleStream failed to write status: %v", err)
		}
		return
	}
	// Unary RPC or Streaming RPC?
	if md, ok := srv.md[method]; ok {
		s.processUnaryRPC(t, stream, srv, md)
		return
	}
	if sd, ok := srv.sd[method]; ok {
		s.processStreamingRPC(t, stream, srv, sd)
		return
	}
	if err := t.WriteStatus(stream, codes.Unimplemented, fmt.Sprintf("unknown method %v", method)); err != nil {
		grpclog.Printf("grpc: Server.handleStream failed to write status: %v", err)
	}
}

// recvResponse receives and parses an RPC response.
// On error, it returns the error and indicates whether the call should be retried.
//
// TODO(zhaoq): Check whether the received message sequence is valid.
func recvResponse(codec Codec, t transport.ClientTransport, c *callInfo, stream *transport.Stream, reply interface{}) error {
	// Try to acquire header metadata from the server if there is any.
	var err error
	c.headerMD, err = stream.Header()
	if err != nil {
		return err
	}
	p := &parser{s: stream}
	for {
		if err = recv(p, codec, reply); err != nil {
			if err == io.EOF {
				break
			}
			return err
		}
	}
	c.trailerMD = stream.Trailer()
	return nil
}

func (s *Server) processUnaryRPC(t transport.ServerTransport, stream *transport.Stream, srv *service, md *MethodDesc) {
	p := &parser{s: stream}
	for {
		pf, req, err := p.recvMsg()
		if err == io.EOF {
			// The entire stream is done (for unary RPC only).
			return
		}
		if err != nil {
			switch err := err.(type) {
			case transport.ConnectionError:
				// Nothing to do here.
			case transport.StreamError:
				if err := t.WriteStatus(stream, err.Code, err.Desc); err != nil {
					grpclog.Printf("grpc: Server.processUnaryRPC failed to write status: %v", err)
				}
			default:
				panic(fmt.Sprintf("grpc: Unexpected error (%T) from recvMsg: %v", err, err))
			}
			return
		}
		switch pf {
		case compressionNone:
			statusCode := codes.OK
			statusDesc := ""
			reply, appErr := md.Handler(srv.server, stream.Context(), s.opts.codec, req)
			if appErr != nil {
				if err, ok := appErr.(rpcError); ok {
					statusCode = err.code
					statusDesc = err.desc
				} else {
					statusCode = convertCode(appErr)
					statusDesc = appErr.Error()
				}
				if err := t.WriteStatus(stream, statusCode, statusDesc); err != nil {
					grpclog.Printf("grpc: Server.processUnaryRPC failed to write status: %v", err)
				}
				return
			}
			opts := &transport.Options{
				Last:  true,
				Delay: false,
			}
			if err := s.sendResponse(t, stream, reply, compressionNone, opts); err != nil {
				if _, ok := err.(transport.ConnectionError); ok {
					return
				}
				if e, ok := err.(transport.StreamError); ok {
					statusCode = e.Code
					statusDesc = e.Desc
				} else {
					statusCode = codes.Unknown
					statusDesc = err.Error()
				}
			}
			t.WriteStatus(stream, statusCode, statusDesc)
		default:
			panic(fmt.Sprintf("payload format to be supported: %d", pf))
		}
	}
}

// Invoke is called by the generated code. It sends the RPC request on the
// wire and returns after response is received.
func Invoke(ctx context.Context, method string, args, reply interface{}, cc *ClientConn, opts ...CallOption) (err error) {
	var c callInfo
	for _, o := range opts {
		if err := o.before(&c); err != nil {
			return toRPCErr(err)
		}
	}
	defer func() {
		for _, o := range opts {
			o.after(&c)
		}
	}()

	if EnableTracing {
		c.traceInfo.tr = trace.New("grpc.Sent."+methodFamily(method), method)
		defer c.traceInfo.tr.Finish()
		c.traceInfo.firstLine.client = true
		if deadline, ok := ctx.Deadline(); ok {
			c.traceInfo.firstLine.deadline = deadline.Sub(time.Now())
		}
		c.traceInfo.tr.LazyLog(&c.traceInfo.firstLine, false)
		// TODO(dsymonds): Arrange for c.traceInfo.firstLine.remoteAddr to be set.
		defer func() {
			if err != nil {
				c.traceInfo.tr.LazyLog(&fmtStringer{"%v", []interface{}{err}}, true)
				c.traceInfo.tr.SetError()
			}
		}()
	}
	callHdr := &transport.CallHdr{
		Host:   cc.authority,
		Method: method,
	}
	topts := &transport.Options{
		Last:  true,
		Delay: false,
	}
	var (
		ts      int   // track the transport sequence number
		lastErr error // record the error that happened
	)
	for {
		var (
			err    error
			t      transport.ClientTransport
			stream *transport.Stream
		)
		// TODO(zhaoq): Need a formal spec of retry strategy for non-failfast rpcs.
		if lastErr != nil && c.failFast {
			return toRPCErr(lastErr)
		}
		t, ts, err = cc.wait(ctx, ts)
		if err != nil {
			if lastErr != nil {
				// This was a retry; return the error from the last attempt.
				return toRPCErr(lastErr)
			}
			return toRPCErr(err)
		}
		if c.traceInfo.tr != nil {
			c.traceInfo.tr.LazyLog(&payload{sent: true, msg: args}, true)
		}
		stream, err = sendRequest(ctx, cc.dopts.codec, callHdr, t, args, topts)
		if err != nil {
			if _, ok := err.(transport.ConnectionError); ok {
				lastErr = err
				continue
			}
			if lastErr != nil {
				return toRPCErr(lastErr)
			}
			return toRPCErr(err)
		}
		// Receive the response
		lastErr = recvResponse(cc.dopts.codec, t, &c, stream, reply)
		if _, ok := lastErr.(transport.ConnectionError); ok {
			continue
		}
		if c.traceInfo.tr != nil {
			c.traceInfo.tr.LazyLog(&payload{sent: false, msg: reply}, true)
		}
		t.CloseStream(stream, lastErr)
		if lastErr != nil {
			return toRPCErr(lastErr)
		}
		return Errorf(stream.StatusCode(), stream.StatusDesc())
	}
}

// Invoke is called by the generated code. It sends the RPC request on the
// wire and returns after response is received.
func Invoke(ctx context.Context, method string, args, reply interface{}, cc *ClientConn, opts ...CallOption) error {
	var c callInfo
	for _, o := range opts {
		if err := o.before(&c); err != nil {
			return toRPCErr(err)
		}
	}
	defer func() {
		for _, o := range opts {
			o.after(&c)
		}
	}()
	callHdr := &transport.CallHdr{
		Host:   cc.authority,
		Method: method,
	}
	topts := &transport.Options{
		Last:  true,
		Delay: false,
	}
	var (
		ts      int   // track the transport sequence number
		lastErr error // record the error that happened
	)
	for {
		var (
			err    error
			t      transport.ClientTransport
			stream *transport.Stream
		)
		// TODO(zhaoq): Need a formal spec of retry strategy for non-failfast rpcs.
		if lastErr != nil && c.failFast {
			return toRPCErr(lastErr)
		}
		t, ts, err = cc.wait(ctx, ts)
		if err != nil {
			if lastErr != nil {
				// This was a retry; return the error from the last attempt.
				return toRPCErr(lastErr)
			}
			return toRPCErr(err)
		}
		stream, err = sendRequest(ctx, cc.dopts.codec, callHdr, t, args, topts)
		if err != nil {
			if _, ok := err.(transport.ConnectionError); ok {
				lastErr = err
				continue
			}
			if lastErr != nil {
				return toRPCErr(lastErr)
			}
			return toRPCErr(err)
		}
		// Receive the response
		lastErr = recvResponse(cc.dopts.codec, t, &c, stream, reply)
		if _, ok := lastErr.(transport.ConnectionError); ok {
			continue
		}
		t.CloseStream(stream, lastErr)
		if lastErr != nil {
			return toRPCErr(lastErr)
		}
		return Errorf(stream.StatusCode(), stream.StatusDesc())
	}
}

func (s *Server) processUnaryRPC(t transport.ServerTransport, stream *transport.Stream, srv *service, md *MethodDesc) (err error) {
	var traceInfo traceInfo
	if EnableTracing {
		traceInfo.tr = trace.New("grpc.Recv."+methodFamily(stream.Method()), stream.Method())
		defer traceInfo.tr.Finish()
		traceInfo.firstLine.client = false
		traceInfo.tr.LazyLog(&traceInfo.firstLine, false)
		defer func() {
			if err != nil && err != io.EOF {
				traceInfo.tr.LazyLog(&fmtStringer{"%v", []interface{}{err}}, true)
				traceInfo.tr.SetError()
			}
		}()
	}
	p := &parser{s: stream}
	for {
		pf, req, err := p.recvMsg()
		if err == io.EOF {
			// The entire stream is done (for unary RPC only).
			return err
		}
		if err != nil {
			switch err := err.(type) {
			case transport.ConnectionError:
				// Nothing to do here.
			case transport.StreamError:
				if err := t.WriteStatus(stream, err.Code, err.Desc); err != nil {
					grpclog.Printf("grpc: Server.processUnaryRPC failed to write status: %v", err)
				}
			default:
				panic(fmt.Sprintf("grpc: Unexpected error (%T) from recvMsg: %v", err, err))
			}
			return err
		}
		if traceInfo.tr != nil {
			traceInfo.tr.LazyLog(&payload{sent: false, msg: req}, true)
		}
		switch pf {
		case compressionNone:
			statusCode := codes.OK
			statusDesc := ""
			reply, appErr := md.Handler(srv.server, stream.Context(), s.opts.codec, req)
			if appErr != nil {
				if err, ok := appErr.(rpcError); ok {
					statusCode = err.code
					statusDesc = err.desc
				} else {
					statusCode = convertCode(appErr)
					statusDesc = appErr.Error()
				}
				if err := t.WriteStatus(stream, statusCode, statusDesc); err != nil {
					grpclog.Printf("grpc: Server.processUnaryRPC failed to write status: %v", err)
					return err
				}
				return nil
			}
			opts := &transport.Options{
				Last:  true,
				Delay: false,
			}
			if err := s.sendResponse(t, stream, reply, compressionNone, opts); err != nil {
				switch err := err.(type) {
				case transport.ConnectionError:
					// Nothing to do here.
				case transport.StreamError:
					statusCode = err.Code
					statusDesc = err.Desc
				default:
					statusCode = codes.Unknown
					statusDesc = err.Error()
				}
				return err
			}
			if traceInfo.tr != nil {
				traceInfo.tr.LazyLog(&payload{sent: true, msg: reply}, true)
			}
			return t.WriteStatus(stream, statusCode, statusDesc)
		default:
			panic(fmt.Sprintf("payload format to be supported: %d", pf))
		}
	}
}