grpc: Change server stream context handling (#6598)

internal/transport/handler_server.go

@@ -344,7 +344,7 @@ func (ht *serverHandlerTransport) WriteHeader(s *Stream, md metadata.MD) error {
 	return err
 }
 
-func (ht *serverHandlerTransport) HandleStreams(startStream func(*Stream), traceCtx func(context.Context, string) context.Context) {
+func (ht *serverHandlerTransport) HandleStreams(startStream func(*Stream)) {
 	// With this transport type there will be exactly 1 stream: this HTTP request.
 
 	ctx := ht.req.Context()

internal/transport/handler_server_test.go

@@ -315,7 +315,6 @@ func (s) TestHandlerTransport_HandleStreams(t *testing.T) {
 	}
 	st.ht.HandleStreams(
 		func(s *Stream) { go handleStream(s) },
-		func(ctx context.Context, method string) context.Context { return ctx },
 	)
 	wantHeader := http.Header{
 		"Date":          nil,
@@ -349,7 +348,6 @@ func handleStreamCloseBodyTest(t *testing.T, statusCode codes.Code, msg string)
 	}
 	st.ht.HandleStreams(
 		func(s *Stream) { go handleStream(s) },
-		func(ctx context.Context, method string) context.Context { return ctx },
 	)
 	wantHeader := http.Header{
 		"Date":         nil,
@@ -399,7 +397,6 @@ func (s) TestHandlerTransport_HandleStreams_Timeout(t *testing.T) {
 	}
 	ht.HandleStreams(
 		func(s *Stream) { go runStream(s) },
-		func(ctx context.Context, method string) context.Context { return ctx },
 	)
 	wantHeader := http.Header{
 		"Date":         nil,
@@ -452,7 +449,6 @@ func testHandlerTransportHandleStreams(t *testing.T, handleStream func(st *handl
 	st := newHandleStreamTest(t)
 	st.ht.HandleStreams(
 		func(s *Stream) { go handleStream(st, s) },
-		func(ctx context.Context, method string) context.Context { return ctx },
 	)
 }
 
@@ -486,7 +482,6 @@ func (s) TestHandlerTransport_HandleStreams_ErrDetails(t *testing.T) {
 	}
 	hst.ht.HandleStreams(
 		func(s *Stream) { go handleStream(s) },
-		func(ctx context.Context, method string) context.Context { return ctx },
 	)
 	wantHeader := http.Header{
 		"Date":         nil,

internal/transport/http2_server.go

@@ -347,7 +347,7 @@ func NewServerTransport(conn net.Conn, config *ServerConfig) (_ ServerTransport,
 
 // operateHeaders takes action on the decoded headers. Returns an error if fatal
 // error encountered and transport needs to close, otherwise returns nil.
-func (t *http2Server) operateHeaders(frame *http2.MetaHeadersFrame, handle func(*Stream), traceCtx func(context.Context, string) context.Context) error {
+func (t *http2Server) operateHeaders(frame *http2.MetaHeadersFrame, handle func(*Stream)) error {
 	// Acquire max stream ID lock for entire duration
 	t.maxStreamMu.Lock()
 	defer t.maxStreamMu.Unlock()
@@ -597,7 +597,6 @@ func (t *http2Server) operateHeaders(frame *http2.MetaHeadersFrame, handle func(
 	s.requestRead = func(n int) {
 		t.adjustWindow(s, uint32(n))
 	}
-	s.ctx = traceCtx(s.ctx, s.method)
 	for _, sh := range t.stats {
 		s.ctx = sh.TagRPC(s.ctx, &stats.RPCTagInfo{FullMethodName: s.method})
 		inHeader := &stats.InHeader{
@@ -635,7 +634,7 @@ func (t *http2Server) operateHeaders(frame *http2.MetaHeadersFrame, handle func(
 // HandleStreams receives incoming streams using the given handler. This is
 // typically run in a separate goroutine.
 // traceCtx attaches trace to ctx and returns the new context.
-func (t *http2Server) HandleStreams(handle func(*Stream), traceCtx func(context.Context, string) context.Context) {
+func (t *http2Server) HandleStreams(handle func(*Stream)) {
 	defer close(t.readerDone)
 	for {
 		t.controlBuf.throttle()
@@ -670,7 +669,7 @@ func (t *http2Server) HandleStreams(handle func(*Stream), traceCtx func(context.
 		}
 		switch frame := frame.(type) {
 		case *http2.MetaHeadersFrame:
-			if err := t.operateHeaders(frame, handle, traceCtx); err != nil {
+			if err := t.operateHeaders(frame, handle); err != nil {
 				t.Close(err)
 				break
 			}

internal/transport/transport.go

@@ -698,7 +698,7 @@ type ClientTransport interface {
 // Write methods for a given Stream will be called serially.
 type ServerTransport interface {
 	// HandleStreams receives incoming streams using the given handler.
-	HandleStreams(func(*Stream), func(context.Context, string) context.Context)
+	HandleStreams(func(*Stream))
 
 	// WriteHeader sends the header metadata for the given stream.
 	// WriteHeader may not be called on all streams.

internal/transport/transport_test.go

@@ -353,32 +353,20 @@ func (s *server) start(t *testing.T, port int, serverConfig *ServerConfig, ht hT
 		s.mu.Unlock()
 		switch ht {
 		case notifyCall:
-			go transport.HandleStreams(h.handleStreamAndNotify,
-				func(ctx context.Context, _ string) context.Context {
-					return ctx
-				})
+			go transport.HandleStreams(h.handleStreamAndNotify)
 		case suspended:
-			go transport.HandleStreams(func(*Stream) {}, // Do nothing to handle the stream.
-				func(ctx context.Context, method string) context.Context {
-					return ctx
-				})
+			go transport.HandleStreams(func(*Stream) {})
 		case misbehaved:
 			go transport.HandleStreams(func(s *Stream) {
 				go h.handleStreamMisbehave(t, s)
-			}, func(ctx context.Context, method string) context.Context {
-				return ctx
 			})
 		case encodingRequiredStatus:
 			go transport.HandleStreams(func(s *Stream) {
 				go h.handleStreamEncodingRequiredStatus(s)
-			}, func(ctx context.Context, method string) context.Context {
-				return ctx
 			})
 		case invalidHeaderField:
 			go transport.HandleStreams(func(s *Stream) {
 				go h.handleStreamInvalidHeaderField(s)
-			}, func(ctx context.Context, method string) context.Context {
-				return ctx
 			})
 		case delayRead:
 			h.notify = make(chan struct{})
@@ -388,20 +376,14 @@ func (s *server) start(t *testing.T, port int, serverConfig *ServerConfig, ht hT
 			s.mu.Unlock()
 			go transport.HandleStreams(func(s *Stream) {
 				go h.handleStreamDelayRead(t, s)
-			}, func(ctx context.Context, method string) context.Context {
-				return ctx
 			})
 		case pingpong:
 			go transport.HandleStreams(func(s *Stream) {
 				go h.handleStreamPingPong(t, s)
-			}, func(ctx context.Context, method string) context.Context {
-				return ctx
 			})
 		default:
 			go transport.HandleStreams(func(s *Stream) {
 				go h.handleStream(t, s)
-			}, func(ctx context.Context, method string) context.Context {
-				return ctx
 			})
 		}
 	}