Merge pull request #384 from mwitkow-io/bugfix/fix_stream_codec_errors_take2
fix difference between unitary and stream codec error handling
diff --git a/.travis.yml b/.travis.yml
index 2503560..3f83776 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -1,4 +1,14 @@
language: go
+before_install:
+ - go get github.com/axw/gocov/gocov
+ - go get github.com/mattn/goveralls
+ - go get golang.org/x/tools/cmd/cover
+
+install:
+ - mkdir -p "$GOPATH/src/google.golang.org"
+ - mv "$TRAVIS_BUILD_DIR" "$GOPATH/src/google.golang.org/grpc"
+
script:
- make test testrace
+ - make coverage
diff --git a/Makefile b/Makefile
index 0dc225f..5bc38be 100644
--- a/Makefile
+++ b/Makefile
@@ -45,3 +45,6 @@
clean:
go clean google.golang.org/grpc/...
+
+coverage: testdeps
+ goveralls -v google.golang.org/grpc/...
diff --git a/README.md b/README.md
index 94dc739..37b05f0 100644
--- a/README.md
+++ b/README.md
@@ -1,6 +1,6 @@
#gRPC-Go
-[](https://travis-ci.org/grpc/grpc-go) [](https://godoc.org/google.golang.org/grpc) [](https://coveralls.io/github/grpc/grpc-go?branch=master)
+[](https://travis-ci.org/grpc/grpc-go) [](https://godoc.org/google.golang.org/grpc)
The Go implementation of [gRPC](http://www.grpc.io/): A high performance, open source, general RPC framework that puts mobile and HTTP/2 first. For more information see the [gRPC Quick Start](http://www.grpc.io/docs/) guide.
@@ -20,11 +20,11 @@
Constraints
-----------
-The grpc package should only depend on standard Go packages and a short list of exceptions. A new addition to the list requires a discussion with gRPC-Go authors and consultants.
+The grpc package should only depend on standard Go packages and a small number of exceptions. If your contribution introduces new dependencies which are NOT in the [list](http://godoc.org/google.golang.org/grpc?imports), you need a discussion with gRPC-Go authors and consultants.
Documentation
-------------
-You can find more detailed documentation and examples in the [examples directory](examples/).
+See [API documentation](https://godoc.org/google.golang.org/grpc) for package and API descriptions and find examples in the [examples directory](examples/).
Status
------
diff --git a/benchmark/grpc_testing/test.pb.go b/benchmark/grpc_testing/test.pb.go
index 619c450..74e13c9 100644
--- a/benchmark/grpc_testing/test.pb.go
+++ b/benchmark/grpc_testing/test.pb.go
@@ -419,9 +419,9 @@
s.RegisterService(&_TestService_serviceDesc, srv)
}
-func _TestService_UnaryCall_Handler(srv interface{}, ctx context.Context, codec grpc.Codec, buf []byte) (interface{}, error) {
+func _TestService_UnaryCall_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error) (interface{}, error) {
in := new(SimpleRequest)
- if err := codec.Unmarshal(buf, in); err != nil {
+ if err := dec(in); err != nil {
return nil, err
}
out, err := srv.(TestServiceServer).UnaryCall(ctx, in)
diff --git a/clientconn.go b/clientconn.go
index ea3ccd0..f5b28d8 100644
--- a/clientconn.go
+++ b/clientconn.go
@@ -290,7 +290,7 @@
c.Close()
return
}
- go c.transportMonitor()
+ c.transportMonitor()
}()
}
return c, nil
diff --git a/examples/helloworld/helloworld/helloworld.pb.go b/examples/helloworld/helloworld/helloworld.pb.go
index 1ff931a..366b23b 100644
--- a/examples/helloworld/helloworld/helloworld.pb.go
+++ b/examples/helloworld/helloworld/helloworld.pb.go
@@ -84,9 +84,9 @@
s.RegisterService(&_Greeter_serviceDesc, srv)
}
-func _Greeter_SayHello_Handler(srv interface{}, ctx context.Context, codec grpc.Codec, buf []byte) (interface{}, error) {
+func _Greeter_SayHello_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error) (interface{}, error) {
in := new(HelloRequest)
- if err := codec.Unmarshal(buf, in); err != nil {
+ if err := dec(in); err != nil {
return nil, err
}
out, err := srv.(GreeterServer).SayHello(ctx, in)
diff --git a/examples/route_guide/routeguide/route_guide.pb.go b/examples/route_guide/routeguide/route_guide.pb.go
index fcf5c74..9ac9029 100644
--- a/examples/route_guide/routeguide/route_guide.pb.go
+++ b/examples/route_guide/routeguide/route_guide.pb.go
@@ -310,9 +310,9 @@
s.RegisterService(&_RouteGuide_serviceDesc, srv)
}
-func _RouteGuide_GetFeature_Handler(srv interface{}, ctx context.Context, codec grpc.Codec, buf []byte) (interface{}, error) {
+func _RouteGuide_GetFeature_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error) (interface{}, error) {
in := new(Point)
- if err := codec.Unmarshal(buf, in); err != nil {
+ if err := dec(in); err != nil {
return nil, err
}
out, err := srv.(RouteGuideServer).GetFeature(ctx, in)
diff --git a/health/grpc_health_v1alpha/health.pb.go b/health/grpc_health_v1alpha/health.pb.go
index c333a97..96eba6f 100644
--- a/health/grpc_health_v1alpha/health.pb.go
+++ b/health/grpc_health_v1alpha/health.pb.go
@@ -108,9 +108,9 @@
s.RegisterService(&_Health_serviceDesc, srv)
}
-func _Health_Check_Handler(srv interface{}, ctx context.Context, codec grpc.Codec, buf []byte) (interface{}, error) {
+func _Health_Check_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error) (interface{}, error) {
in := new(HealthCheckRequest)
- if err := codec.Unmarshal(buf, in); err != nil {
+ if err := dec(in); err != nil {
return nil, err
}
out, err := srv.(HealthServer).Check(ctx, in)
diff --git a/interop/grpc_testing/test.pb.go b/interop/grpc_testing/test.pb.go
index b25e98b..bd492fe 100755
--- a/interop/grpc_testing/test.pb.go
+++ b/interop/grpc_testing/test.pb.go
@@ -539,9 +539,9 @@
s.RegisterService(&_TestService_serviceDesc, srv)
}
-func _TestService_EmptyCall_Handler(srv interface{}, ctx context.Context, codec grpc.Codec, buf []byte) (interface{}, error) {
+func _TestService_EmptyCall_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error) (interface{}, error) {
in := new(Empty)
- if err := codec.Unmarshal(buf, in); err != nil {
+ if err := dec(in); err != nil {
return nil, err
}
out, err := srv.(TestServiceServer).EmptyCall(ctx, in)
@@ -551,9 +551,9 @@
return out, nil
}
-func _TestService_UnaryCall_Handler(srv interface{}, ctx context.Context, codec grpc.Codec, buf []byte) (interface{}, error) {
+func _TestService_UnaryCall_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error) (interface{}, error) {
in := new(SimpleRequest)
- if err := codec.Unmarshal(buf, in); err != nil {
+ if err := dec(in); err != nil {
return nil, err
}
out, err := srv.(TestServiceServer).UnaryCall(ctx, in)
diff --git a/server.go b/server.go
index ba68a21..5c20dd4 100644
--- a/server.go
+++ b/server.go
@@ -42,6 +42,7 @@
"runtime"
"strings"
"sync"
+ "time"
"golang.org/x/net/context"
"golang.org/x/net/trace"
@@ -52,7 +53,7 @@
"google.golang.org/grpc/transport"
)
-type methodHandler func(srv interface{}, ctx context.Context, codec Codec, buf []byte) (interface{}, error)
+type methodHandler func(srv interface{}, ctx context.Context, dec func(interface{}) error) (interface{}, error)
// MethodDesc represents an RPC service's method specification.
type MethodDesc struct {
@@ -258,8 +259,24 @@
s.mu.Unlock()
go func() {
- st.HandleStreams(func(stream *transport.Stream) {
- s.handleStream(st, stream)
+ st.HandleStreams(func(stream *transport.Stream, wg *sync.WaitGroup) {
+ var trInfo *traceInfo
+ if EnableTracing {
+ trInfo = &traceInfo{
+ tr: trace.New("grpc.Recv."+methodFamily(stream.Method()), stream.Method()),
+ }
+ trInfo.firstLine.client = false
+ trInfo.firstLine.remoteAddr = st.RemoteAddr()
+ stream.TraceContext(trInfo.tr)
+ if dl, ok := stream.Context().Deadline(); ok {
+ trInfo.firstLine.deadline = dl.Sub(time.Now())
+ }
+ }
+ wg.Add(1)
+ go func() {
+ s.handleStream(st, stream, trInfo)
+ wg.Done()
+ }()
})
s.mu.Lock()
delete(s.conns, st)
@@ -283,20 +300,15 @@
return t.Write(stream, p, opts)
}
-func (s *Server) processUnaryRPC(t transport.ServerTransport, stream *transport.Stream, srv *service, md *MethodDesc) (err error) {
- ctx, cancel := context.WithCancel(stream.Context())
- defer cancel()
- 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)
- ctx = trace.NewContext(ctx, traceInfo.tr)
+func (s *Server) processUnaryRPC(t transport.ServerTransport, stream *transport.Stream, srv *service, md *MethodDesc, trInfo *traceInfo) (err error) {
+ if trInfo != nil {
+ defer trInfo.tr.Finish()
+ trInfo.firstLine.client = false
+ trInfo.tr.LazyLog(&trInfo.firstLine, false)
defer func() {
if err != nil && err != io.EOF {
- traceInfo.tr.LazyLog(&fmtStringer{"%v", []interface{}{err}}, true)
- traceInfo.tr.SetError()
+ trInfo.tr.LazyLog(&fmtStringer{"%v", []interface{}{err}}, true)
+ trInfo.tr.SetError()
}
}()
}
@@ -320,16 +332,20 @@
}
return err
}
- if traceInfo.tr != nil {
- // TODO: set payload.msg to something that
- // prints usefully with %s; req is a []byte.
- traceInfo.tr.LazyLog(&payload{sent: false}, true)
- }
switch pf {
case compressionNone:
statusCode := codes.OK
statusDesc := ""
- reply, appErr := md.Handler(srv.server, ctx, s.opts.codec, req)
+ df := func(v interface{}) error {
+ if err := s.opts.codec.Unmarshal(req, v); err != nil {
+ return err
+ }
+ if trInfo != nil {
+ trInfo.tr.LazyLog(&payload{sent: false, msg: v}, true)
+ }
+ return nil
+ }
+ reply, appErr := md.Handler(srv.server, stream.Context(), df)
if appErr != nil {
if err, ok := appErr.(rpcError); ok {
statusCode = err.code
@@ -338,9 +354,9 @@
statusCode = convertCode(appErr)
statusDesc = appErr.Error()
}
- if traceInfo.tr != nil && statusCode != codes.OK {
- traceInfo.tr.LazyLog(stringer(statusDesc), true)
- traceInfo.tr.SetError()
+ if trInfo != nil && statusCode != codes.OK {
+ trInfo.tr.LazyLog(stringer(statusDesc), true)
+ trInfo.tr.SetError()
}
if err := t.WriteStatus(stream, statusCode, statusDesc); err != nil {
@@ -349,8 +365,8 @@
}
return nil
}
- if traceInfo.tr != nil {
- traceInfo.tr.LazyLog(stringer("OK"), false)
+ if trInfo != nil {
+ trInfo.tr.LazyLog(stringer("OK"), false)
}
opts := &transport.Options{
Last: true,
@@ -369,8 +385,8 @@
}
return err
}
- if traceInfo.tr != nil {
- traceInfo.tr.LazyLog(&payload{sent: true, msg: reply}, true)
+ if trInfo != nil {
+ trInfo.tr.LazyLog(&payload{sent: true, msg: reply}, true)
}
return t.WriteStatus(stream, statusCode, statusDesc)
default:
@@ -379,30 +395,24 @@
}
}
-func (s *Server) processStreamingRPC(t transport.ServerTransport, stream *transport.Stream, srv *service, sd *StreamDesc) (err error) {
- ctx, cancel := context.WithCancel(stream.Context())
- defer cancel()
+func (s *Server) processStreamingRPC(t transport.ServerTransport, stream *transport.Stream, srv *service, sd *StreamDesc, trInfo *traceInfo) (err error) {
ss := &serverStream{
- t: t,
- s: stream,
- ctx: ctx,
- p: &parser{s: stream},
- codec: s.opts.codec,
- tracing: EnableTracing,
+ t: t,
+ s: stream,
+ p: &parser{s: stream},
+ codec: s.opts.codec,
+ trInfo: trInfo,
}
- 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)
- ss.ctx = trace.NewContext(ss.ctx, ss.traceInfo.tr)
+ if trInfo != nil {
+ trInfo.tr.LazyLog(&trInfo.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()
+ trInfo.tr.LazyLog(&fmtStringer{"%v", []interface{}{err}}, true)
+ trInfo.tr.SetError()
}
- ss.traceInfo.tr.Finish()
- ss.traceInfo.tr = nil
+ trInfo.tr.Finish()
+ trInfo.tr = nil
ss.mu.Unlock()
}()
}
@@ -415,13 +425,13 @@
ss.statusDesc = appErr.Error()
}
}
- if ss.tracing {
+ if trInfo != nil {
ss.mu.Lock()
if ss.statusCode != codes.OK {
- ss.traceInfo.tr.LazyLog(stringer(ss.statusDesc), true)
- ss.traceInfo.tr.SetError()
+ trInfo.tr.LazyLog(stringer(ss.statusDesc), true)
+ trInfo.tr.SetError()
} else {
- ss.traceInfo.tr.LazyLog(stringer("OK"), false)
+ trInfo.tr.LazyLog(stringer("OK"), false)
}
ss.mu.Unlock()
}
@@ -429,7 +439,7 @@
}
-func (s *Server) handleStream(t transport.ServerTransport, stream *transport.Stream) {
+func (s *Server) handleStream(t transport.ServerTransport, stream *transport.Stream, trInfo *traceInfo) {
sm := stream.Method()
if sm != "" && sm[0] == '/' {
sm = sm[1:]
@@ -452,11 +462,11 @@
}
// Unary RPC or Streaming RPC?
if md, ok := srv.md[method]; ok {
- s.processUnaryRPC(t, stream, srv, md)
+ s.processUnaryRPC(t, stream, srv, md, trInfo)
return
}
if sd, ok := srv.sd[method]; ok {
- s.processStreamingRPC(t, stream, srv, sd)
+ s.processStreamingRPC(t, stream, srv, sd, trInfo)
return
}
if err := t.WriteStatus(stream, codes.Unimplemented, fmt.Sprintf("unknown method %v", method)); err != nil {
diff --git a/stream.go b/stream.go
index 34774f0..21c51e8 100644
--- a/stream.go
+++ b/stream.go
@@ -126,13 +126,13 @@
tracing: EnableTracing,
}
if cs.tracing {
- cs.traceInfo.tr = trace.New("grpc.Sent."+methodFamily(method), method)
- cs.traceInfo.firstLine.client = true
+ cs.trInfo.tr = trace.New("grpc.Sent."+methodFamily(method), method)
+ cs.trInfo.firstLine.client = true
if deadline, ok := ctx.Deadline(); ok {
- cs.traceInfo.firstLine.deadline = deadline.Sub(time.Now())
+ cs.trInfo.firstLine.deadline = deadline.Sub(time.Now())
}
- cs.traceInfo.tr.LazyLog(&cs.traceInfo.firstLine, false)
- ctx = trace.NewContext(ctx, cs.traceInfo.tr)
+ cs.trInfo.tr.LazyLog(&cs.trInfo.firstLine, false)
+ ctx = trace.NewContext(ctx, cs.trInfo.tr)
}
s, err := t.NewStream(ctx, callHdr)
if err != nil {
@@ -154,10 +154,10 @@
tracing bool // set to EnableTracing when the clientStream is created.
- mu sync.Mutex // protects traceInfo
- // traceInfo.tr is set when the clientStream is created (if EnableTracing is true),
+ mu sync.Mutex // protects trInfo.tr
+ // trInfo.tr is set when the clientStream is created (if EnableTracing is true),
// and is set to nil when the clientStream's finish method is called.
- traceInfo traceInfo
+ trInfo traceInfo
}
func (cs *clientStream) Context() context.Context {
@@ -181,8 +181,8 @@
func (cs *clientStream) SendMsg(m interface{}) (err error) {
if cs.tracing {
cs.mu.Lock()
- if cs.traceInfo.tr != nil {
- cs.traceInfo.tr.LazyLog(&payload{sent: true, msg: m}, true)
+ if cs.trInfo.tr != nil {
+ cs.trInfo.tr.LazyLog(&payload{sent: true, msg: m}, true)
}
cs.mu.Unlock()
}
@@ -213,8 +213,8 @@
if err == nil {
if cs.tracing {
cs.mu.Lock()
- if cs.traceInfo.tr != nil {
- cs.traceInfo.tr.LazyLog(&payload{sent: false, msg: m}, true)
+ if cs.trInfo.tr != nil {
+ cs.trInfo.tr.LazyLog(&payload{sent: false, msg: m}, true)
}
cs.mu.Unlock()
}
@@ -266,15 +266,15 @@
}
cs.mu.Lock()
defer cs.mu.Unlock()
- if cs.traceInfo.tr != nil {
+ if cs.trInfo.tr != nil {
if err == nil || err == io.EOF {
- cs.traceInfo.tr.LazyPrintf("RPC: [OK]")
+ cs.trInfo.tr.LazyPrintf("RPC: [OK]")
} else {
- cs.traceInfo.tr.LazyPrintf("RPC: [%v]", err)
- cs.traceInfo.tr.SetError()
+ cs.trInfo.tr.LazyPrintf("RPC: [%v]", err)
+ cs.trInfo.tr.SetError()
}
- cs.traceInfo.tr.Finish()
- cs.traceInfo.tr = nil
+ cs.trInfo.tr.Finish()
+ cs.trInfo.tr = nil
}
}
@@ -294,22 +294,17 @@
type serverStream struct {
t transport.ServerTransport
s *transport.Stream
- ctx context.Context // provides trace.FromContext when tracing
p *parser
codec Codec
statusCode codes.Code
statusDesc string
+ trInfo *traceInfo
- tracing bool // set to EnableTracing when the serverStream is created.
-
- mu sync.Mutex // protects traceInfo
- // traceInfo.tr is set when the serverStream is created (if EnableTracing is true),
- // and is set to nil when the serverStream's finish method is called.
- traceInfo traceInfo
+ mu sync.Mutex // protects trInfo.tr after the service handler runs.
}
func (ss *serverStream) Context() context.Context {
- return ss.ctx
+ return ss.s.Context()
}
func (ss *serverStream) SendHeader(md metadata.MD) error {
@@ -326,13 +321,15 @@
func (ss *serverStream) SendMsg(m interface{}) (err error) {
defer func() {
- if ss.tracing {
+ if ss.trInfo != nil {
ss.mu.Lock()
- if err == nil {
- ss.traceInfo.tr.LazyLog(&payload{sent: true, msg: m}, true)
- } else {
- ss.traceInfo.tr.LazyLog(&fmtStringer{"%v", []interface{}{err}}, true)
- ss.traceInfo.tr.SetError()
+ if ss.trInfo.tr != nil {
+ if err == nil {
+ ss.trInfo.tr.LazyLog(&payload{sent: true, msg: m}, true)
+ } else {
+ ss.trInfo.tr.LazyLog(&fmtStringer{"%v", []interface{}{err}}, true)
+ ss.trInfo.tr.SetError()
+ }
}
ss.mu.Unlock()
}
@@ -347,13 +344,15 @@
func (ss *serverStream) RecvMsg(m interface{}) (err error) {
defer func() {
- if ss.tracing {
+ if ss.trInfo != nil {
ss.mu.Lock()
- if err == nil {
- ss.traceInfo.tr.LazyLog(&payload{sent: false, msg: m}, true)
- } else if err != io.EOF {
- ss.traceInfo.tr.LazyLog(&fmtStringer{"%v", []interface{}{err}}, true)
- ss.traceInfo.tr.SetError()
+ if ss.trInfo.tr != nil {
+ if err == nil {
+ ss.trInfo.tr.LazyLog(&payload{sent: false, msg: m}, true)
+ } else if err != io.EOF {
+ ss.trInfo.tr.LazyLog(&fmtStringer{"%v", []interface{}{err}}, true)
+ ss.trInfo.tr.SetError()
+ }
}
ss.mu.Unlock()
}
diff --git a/test/grpc_testing/test.pb.go b/test/grpc_testing/test.pb.go
index b25e98b..bd492fe 100644
--- a/test/grpc_testing/test.pb.go
+++ b/test/grpc_testing/test.pb.go
@@ -539,9 +539,9 @@
s.RegisterService(&_TestService_serviceDesc, srv)
}
-func _TestService_EmptyCall_Handler(srv interface{}, ctx context.Context, codec grpc.Codec, buf []byte) (interface{}, error) {
+func _TestService_EmptyCall_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error) (interface{}, error) {
in := new(Empty)
- if err := codec.Unmarshal(buf, in); err != nil {
+ if err := dec(in); err != nil {
return nil, err
}
out, err := srv.(TestServiceServer).EmptyCall(ctx, in)
@@ -551,9 +551,9 @@
return out, nil
}
-func _TestService_UnaryCall_Handler(srv interface{}, ctx context.Context, codec grpc.Codec, buf []byte) (interface{}, error) {
+func _TestService_UnaryCall_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error) (interface{}, error) {
in := new(SimpleRequest)
- if err := codec.Unmarshal(buf, in); err != nil {
+ if err := dec(in); err != nil {
return nil, err
}
out, err := srv.(TestServiceServer).UnaryCall(ctx, in)
diff --git a/transport/http2_server.go b/transport/http2_server.go
index 057d936..ed8fde0 100644
--- a/transport/http2_server.go
+++ b/transport/http2_server.go
@@ -115,15 +115,15 @@
}
var buf bytes.Buffer
t := &http2Server{
- conn: conn,
- authInfo: authInfo,
- framer: framer,
- hBuf: &buf,
- hEnc: hpack.NewEncoder(&buf),
- maxStreams: maxStreams,
- controlBuf: newRecvBuffer(),
- fc: &inFlow{limit: initialConnWindowSize},
- sendQuotaPool: newQuotaPool(defaultWindowSize),
+ conn: conn,
+ authInfo: authInfo,
+ framer: framer,
+ hBuf: &buf,
+ hEnc: hpack.NewEncoder(&buf),
+ maxStreams: maxStreams,
+ controlBuf: newRecvBuffer(),
+ fc: &inFlow{limit: initialConnWindowSize},
+ sendQuotaPool: newQuotaPool(defaultWindowSize),
state: reachable,
writableChan: make(chan int, 1),
shutdownChan: make(chan struct{}),
@@ -138,7 +138,7 @@
// operateHeader takes action on the decoded headers. It returns the current
// stream if there are remaining headers on the wire (in the following
// Continuation frame).
-func (t *http2Server) operateHeaders(hDec *hpackDecoder, s *Stream, frame headerFrame, endStream bool, handle func(*Stream), wg *sync.WaitGroup) (pendingStream *Stream) {
+func (t *http2Server) operateHeaders(hDec *hpackDecoder, s *Stream, frame headerFrame, endStream bool, handle func(*Stream, *sync.WaitGroup), wg *sync.WaitGroup) (pendingStream *Stream) {
defer func() {
if pendingStream == nil {
hDec.state = decodeState{}
@@ -202,18 +202,13 @@
recv: s.buf,
}
s.method = hDec.state.method
-
- wg.Add(1)
- go func() {
- handle(s)
- wg.Done()
- }()
+ handle(s, wg)
return nil
}
// HandleStreams receives incoming streams using the given handler. This is
// typically run in a separate goroutine.
-func (t *http2Server) HandleStreams(handle func(*Stream)) {
+func (t *http2Server) HandleStreams(handle func(*Stream, *sync.WaitGroup)) {
// Check the validity of client preface.
preface := make([]byte, len(clientPreface))
if _, err := io.ReadFull(t.conn, preface); err != nil {
@@ -689,3 +684,7 @@
// other goroutines.
s.cancel()
}
+
+func (t *http2Server) RemoteAddr() net.Addr {
+ return t.conn.RemoteAddr()
+}
diff --git a/transport/transport.go b/transport/transport.go
index 2dd38a8..c319a5f 100644
--- a/transport/transport.go
+++ b/transport/transport.go
@@ -47,6 +47,7 @@
"time"
"golang.org/x/net/context"
+ "golang.org/x/net/trace"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/credentials"
"google.golang.org/grpc/metadata"
@@ -232,6 +233,11 @@
return s.ctx
}
+// TraceContext recreates the context of s with a trace.Trace.
+func (s *Stream) TraceContext(tr trace.Trace) {
+ s.ctx = trace.NewContext(s.ctx, tr)
+}
+
// Method returns the method for the stream.
func (s *Stream) Method() string {
return s.method
@@ -385,11 +391,13 @@
// WriteHeader sends the header metedata for the given stream.
WriteHeader(s *Stream, md metadata.MD) error
// HandleStreams receives incoming streams using the given handler.
- HandleStreams(func(*Stream))
+ HandleStreams(func(*Stream, *sync.WaitGroup))
// Close tears down the transport. Once it is called, the transport
// should not be accessed any more. All the pending streams and their
// handlers will be terminated asynchronously.
Close() error
+ // RemoteAddr returns the remote network address.
+ RemoteAddr() net.Addr
}
// StreamErrorf creates an StreamError with the specified error code and description.
diff --git a/transport/transport_test.go b/transport/transport_test.go
index 07ba005..ba1d66a 100644
--- a/transport/transport_test.go
+++ b/transport/transport_test.go
@@ -77,7 +77,8 @@
misbehaved
)
-func (h *testStreamHandler) handleStream(t *testing.T, s *Stream) {
+func (h *testStreamHandler) handleStream(t *testing.T, s *Stream, wg *sync.WaitGroup) {
+ defer wg.Done()
req := expectedRequest
resp := expectedResponse
if s.Method() == "foo.Large" {
@@ -99,11 +100,16 @@
}
// handleStreamSuspension blocks until s.ctx is canceled.
-func (h *testStreamHandler) handleStreamSuspension(s *Stream) {
- <-s.ctx.Done()
+func (h *testStreamHandler) handleStreamSuspension(s *Stream, wg *sync.WaitGroup) {
+ wg.Add(1)
+ go func() {
+ <-s.ctx.Done()
+ wg.Done()
+ }()
}
-func (h *testStreamHandler) handleStreamMisbehave(t *testing.T, s *Stream) {
+func (h *testStreamHandler) handleStreamMisbehave(t *testing.T, s *Stream, wg *sync.WaitGroup) {
+ defer wg.Done()
conn, ok := s.ServerTransport().(*http2Server)
if !ok {
t.Fatalf("Failed to convert %v to *http2Server", s.ServerTransport())
@@ -167,12 +173,14 @@
case suspended:
go transport.HandleStreams(h.handleStreamSuspension)
case misbehaved:
- go transport.HandleStreams(func(s *Stream) {
- h.handleStreamMisbehave(t, s)
+ go transport.HandleStreams(func(s *Stream, wg *sync.WaitGroup) {
+ wg.Add(1)
+ go h.handleStreamMisbehave(t, s, wg)
})
default:
- go transport.HandleStreams(func(s *Stream) {
- h.handleStream(t, s)
+ go transport.HandleStreams(func(s *Stream, wg *sync.WaitGroup) {
+ wg.Add(1)
+ go h.handleStream(t, s, wg)
})
}
}
