Golang transport.Stream類代碼示例

// 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(dopts dialOptions, 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
	defer func() {
		if err != nil {
			if _, ok := err.(transport.ConnectionError); !ok {
				t.CloseStream(stream, err)
			}
		}
	}()
	c.headerMD, err = stream.Header()
	if err != nil {
		return err
	}
	p := &parser{r: stream}
	for {
		if err = recv(p, dopts.codec, stream, dopts.dc, reply, math.MaxInt32); err != nil {
			if err == io.EOF {
				break
			}
			return err
		}
	}
	c.trailerMD = stream.Trailer()
	return nil
}

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 {
			log.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 {
			log.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 {
		log.Printf("grpc: Server.handleStream failed to write status: %v", err)
	}
}

func recv(p *parser, c Codec, s *transport.Stream, dc Decompressor, m interface{}) error {
	// 1. 獲取1幀請求數據
	pf, d, err := p.recvMsg()
	if err != nil {
		return err
	}

	// 2. 壓縮算法在 Stream中都已經通信過?
	if err := checkRecvPayload(pf, s.RecvCompress(), dc); err != nil {
		return err
	}

	// 3. 解壓縮數據
	if pf == compressionMade {
		d, err = dc.Do(bytes.NewReader(d))
		if err != nil {
			return transport.StreamErrorf(codes.Internal, "grpc: failed to decompress the received message %v", err)
		}
	}

	// 4. 反序列化(可以使用protobuf, 也可以使用thrift等等)
	if err := c.Unmarshal(d, m); err != nil {
		return transport.StreamErrorf(codes.Internal, "grpc: failed to unmarshal the received message %v", err)
	}
	return nil
}

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) sendResponse(t transport.ServerTransport, stream *transport.Stream, msg interface{}, cp Compressor, opts *transport.Options) error {
	var (
		cbuf       *bytes.Buffer
		outPayload *stats.OutPayload
	)
	if cp != nil {
		cbuf = new(bytes.Buffer)
	}
	if stats.On() {
		outPayload = &stats.OutPayload{}
	}
	p, err := encode(s.opts.codec, msg, cp, cbuf, outPayload)
	if err != nil {
		// This typically indicates a fatal issue (e.g., memory
		// corruption or hardware faults) the application program
		// cannot handle.
		//
		// TODO(zhaoq): There exist other options also such as only closing the
		// faulty stream locally and remotely (Other streams can keep going). Find
		// the optimal option.
		grpclog.Fatalf("grpc: Server failed to encode response %v", err)
	}
	err = t.Write(stream, p, opts)
	if err == nil && outPayload != nil {
		outPayload.SentTime = time.Now()
		stats.HandleRPC(stream.Context(), outPayload)
	}
	return err
}

func recv(p *parser, c Codec, s *transport.Stream, dc Decompressor, m interface{}, maxMsgSize int, inPayload *stats.InPayload) error {
	pf, d, err := p.recvMsg(maxMsgSize)
	if err != nil {
		return err
	}
	if inPayload != nil {
		inPayload.WireLength = len(d)
	}
	if err := checkRecvPayload(pf, s.RecvCompress(), dc); err != nil {
		return err
	}
	if pf == compressionMade {
		d, err = dc.Do(bytes.NewReader(d))
		if err != nil {
			return Errorf(codes.Internal, "grpc: failed to decompress the received message %v", err)
		}
	}
	if len(d) > maxMsgSize {
		// TODO: Revisit the error code. Currently keep it consistent with java
		// implementation.
		return Errorf(codes.Internal, "grpc: received a message of %d bytes exceeding %d limit", len(d), maxMsgSize)
	}
	if err := c.Unmarshal(d, m); err != nil {
		return Errorf(codes.Internal, "grpc: failed to unmarshal the received message %v", err)
	}
	if inPayload != nil {
		inPayload.RecvTime = time.Now()
		inPayload.Payload = m
		// TODO truncate large payload.
		inPayload.Data = d
		inPayload.Length = len(d)
	}
	return nil
}

func recv(p *parser, s *transport.Stream, dg DecompressorGenerator, m interface{}) error {
	codecType, compressType, d, err := p.recvMsg()
	if err != nil {
		return err
	}

	//decompress
	var dc Decompressor
	if compressType == COMPRESS_TYPE_GZIP && dg != nil {
		dc = dg()
	}
	if err := checkRecvPayload(codecType, compressType, s.RecvCompress(), dc); err != nil {
		return err
	}
	if compressType == COMPRESS_TYPE_GZIP {
		d, err = dc.Do(bytes.NewReader(d))
		if err != nil {
			return transport.StreamErrorf(codes.Internal, "grpc: failed to decompress the received message %v", err)
		}
	}

	//unmarshal
	c, err := GetCodec(codecType)
	if err != nil {
		return err
	}
	if err := c.Unmarshal(d, m); err != nil {
		return transport.StreamErrorf(codes.Internal, "grpc: failed to unmarshal the received message %v", err)
	}
	return nil
}

// 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(dopts dialOptions, t transport.ClientTransport, c *callInfo, stream *transport.Stream, reply interface{}) error {
	// Try to acquire header metadata from the server if there is any.

	// 如何接受請求呢?
	// 1. 和處理的Http請求類似, 首先是: 處理Header
	var err error
	c.headerMD, err = stream.Header()
	if err != nil {
		return err
	}

	// 2. 然後處理stream
	p := &parser{r: stream}
	for {
		// stream --> (parser) --> msg --> decompress --> unmarshal --> reply
		// Iter 1: 讀取到Reply
		// Iter 2: 讀取失敗，碰到EOF
		if err = recv(p, dopts.codec, stream, dopts.dc, reply); err != nil {
			// 似乎需要等待stream的EOF
			if err == io.EOF {
				break
			}
			return err
		}
	}

	// 3. 獲取Trailer?
	c.trailerMD = stream.Trailer()
	return nil
}

func (s *Server) processStreamingRPC(t transport.ServerTransport, stream *transport.Stream, srv *service, sd *StreamDesc, trInfo *traceInfo) (err error) {
	var cp Compressor
	if s.opts.cg != nil {
		cp = s.opts.cg()
		stream.SetSendCompress(cp.Type())
	}
	ss := &serverStream{
		t:      t,
		s:      stream,
		p:      &parser{s: stream},
		codec:  s.opts.codec,
		cp:     cp,
		dg:     s.opts.dg,
		trInfo: trInfo,
	}
	if cp != nil {
		ss.cbuf = new(bytes.Buffer)
	}
	if trInfo != nil {
		trInfo.tr.LazyLog(&trInfo.firstLine, false)
		defer func() {
			ss.mu.Lock()
			if err != nil && err != io.EOF {
				ss.trInfo.tr.LazyLog(&fmtStringer{"%v", []interface{}{err}}, true)
				ss.trInfo.tr.SetError()
			}
			ss.trInfo.tr.Finish()
			ss.trInfo.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 if err, ok := appErr.(transport.StreamError); ok {
			ss.statusCode = err.Code
			ss.statusDesc = err.Desc
		} else {
			ss.statusCode = convertCode(appErr)
			ss.statusDesc = appErr.Error()
		}
	}
	if trInfo != nil {
		ss.mu.Lock()
		if ss.statusCode != codes.OK {
			ss.trInfo.tr.LazyLog(stringer(ss.statusDesc), true)
			ss.trInfo.tr.SetError()
		} else {
			ss.trInfo.tr.LazyLog(stringer("OK"), false)
		}
		ss.mu.Unlock()
	}
	return t.WriteStatus(ss.s, ss.statusCode, ss.statusDesc)

}

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 {
					log.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:
			reply, appErr := md.Handler(srv.server, stream.Context(), req)
			if appErr != nil {
				if err := t.WriteStatus(stream, convertCode(appErr), appErr.Error()); err != nil {
					log.Printf("grpc: Server.processUnaryRPC failed to write status: %v", err)
				}
				return
			}
			opts := &transport.Options{
				Last:  true,
				Delay: false,
			}
			statusCode := codes.OK
			statusDesc := ""
			if err := s.sendProto(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()
				}
			}
			if err := t.WriteStatus(stream, statusCode, statusDesc); err != nil {
				log.Printf("grpc: Server.processUnaryRPC failed to write status: %v", err)
			}
		default:
			panic(fmt.Sprintf("payload format to be supported: %d", pf))
		}
	}
}

func (g *grpcServer) processStream(t transport.ServerTransport, stream *transport.Stream, service *service, mtype *methodType, codec grpc.Codec, ct string, ctx context.Context) (err error) {
	opts := g.opts

	r := &rpcRequest{
		service:     opts.Name,
		contentType: ct,
		method:      stream.Method(),
		stream:      true,
	}

	ss := &rpcStream{
		request:    r,
		t:          t,
		s:          stream,
		p:          &parser{r: stream},
		codec:      codec,
		maxMsgSize: defaultMaxMsgSize,
	}

	function := mtype.method.Func
	var returnValues []reflect.Value

	// Invoke the method, providing a new value for the reply.
	fn := func(ctx context.Context, req server.Request, stream interface{}) error {
		returnValues = function.Call([]reflect.Value{service.rcvr, mtype.prepareContext(ctx), reflect.ValueOf(stream)})
		if err := returnValues[0].Interface(); err != nil {
			return err.(error)
		}

		return nil
	}

	for i := len(opts.HdlrWrappers); i > 0; i-- {
		fn = opts.HdlrWrappers[i-1](fn)
	}

	appErr := fn(ctx, r, ss)
	if appErr != nil {
		if err, ok := appErr.(*rpcError); ok {
			ss.statusCode = err.code
			ss.statusDesc = err.desc
		} else if err, ok := appErr.(transport.StreamError); 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) 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 = stream.Trace()
		ss.traceInfo.firstLine.client = false
		ss.traceInfo.firstLine.remoteAddr = t.RemoteAddr()
		if dl, ok := stream.Context().Deadline(); ok {
			ss.traceInfo.firstLine.deadline = dl.Sub(time.Now())
		}
		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()
		}
	}
	if ss.tracing {
		ss.mu.Lock()
		if ss.statusCode != codes.OK {
			ss.traceInfo.tr.LazyLog(stringer(ss.statusDesc), true)
			ss.traceInfo.tr.SetError()
		} else {
			ss.traceInfo.tr.LazyLog(stringer("OK"), false)
		}
		ss.mu.Unlock()
	}
	return t.WriteStatus(ss.s, ss.statusCode, ss.statusDesc)

}

func (s *Proxy) handleStream(frontTrans transport.ServerTransport, frontStream *transport.Stream) {
	sm := frontStream.Method()
	if sm != "" && sm[0] == '/' {
		sm = sm[1:]
	}
	pos := strings.LastIndex(sm, "/")
	if pos == -1 {
		if err := frontTrans.WriteStatus(frontStream, codes.InvalidArgument, fmt.Sprintf("malformed method name: %q", frontStream.Method())); err != nil {
			s.logger.Printf("proxy: Proxy.handleStream failed to write status: %v", err)
		}
		return
	}
	ProxyStream(s.director, s.logger, frontTrans, frontStream)

}

// traceInfo returns a traceInfo and associates it with stream, if tracing is enabled.
// If tracing is not enabled, it returns nil.
func (s *Server) traceInfo(st transport.ServerTransport, stream *transport.Stream) (trInfo *traceInfo) {
	tr, ok := trace.FromContext(stream.Context())
	if !ok {
		return nil
	}

	trInfo = &traceInfo{
		tr: tr,
	}
	trInfo.firstLine.client = false
	trInfo.firstLine.remoteAddr = st.RemoteAddr()

	if dl, ok := stream.Context().Deadline(); ok {
		trInfo.firstLine.deadline = dl.Sub(time.Now())
	}
	return trInfo
}

// recv 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 recv(t transport.ClientTransport, c *callInfo, stream *transport.Stream, reply proto.Message) 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 = recvProto(p, reply); err != nil {
			if err == io.EOF {
				break
			}
			return err
		}
	}
	c.trailerMD = stream.Trailer()
	return nil
}