server: prohibit more than MaxConcurrentStreams handlers from running at once (#6703) (#6706)

Author: dfawley

internal/transport/http2_server.go

@@ -171,15 +171,10 @@ func NewServerTransport(conn net.Conn, config *ServerConfig) (_ ServerTransport,
 ID: http2.SettingMaxFrameSize,
 Val: http2MaxFrameLen,
 }}
-// TODO(zhaoq): Have a better way to signal "no limit" because 0 is
-// permitted in the HTTP2 spec.
-maxStreams := config.MaxStreams
-if maxStreams == 0 {
-maxStreams = math.MaxUint32
-} else {
+if config.MaxStreams != math.MaxUint32 {
 isettings = append(isettings, http2.Setting{
 ID: http2.SettingMaxConcurrentStreams,
-Val: maxStreams,
+Val: config.MaxStreams,
 })
 }
 dynamicWindow := true
@@ -258,7 +253,7 @@ func NewServerTransport(conn net.Conn, config *ServerConfig) (_ ServerTransport,
 framer: framer,
 readerDone: make(chan struct{}),
 writerDone: make(chan struct{}),
-maxStreams: maxStreams,
+maxStreams: config.MaxStreams,
 inTapHandle: config.InTapHandle,
 fc: &trInFlow{limit: uint32(icwz)},
 state: reachable,

internal/transport/transport_test.go

@@ -337,6 +337,9 @@ func (s *server) start(t *testing.T, port int, serverConfig *ServerConfig, ht hT
 return
 }
 rawConn := conn
+if serverConfig.MaxStreams == 0 {
+serverConfig.MaxStreams = math.MaxUint32
+}
 transport, err := NewServerTransport(conn, serverConfig)
 if err != nil {
 return
@@ -443,8 +446,8 @@ func setUpServerOnly(t *testing.T, port int, sc *ServerConfig, ht hType) *server
 return server
 }
 
-func setUp(t *testing.T, port int, maxStreams uint32, ht hType) (*server, *http2Client, func()) {
-return setUpWithOptions(t, port, &ServerConfig{MaxStreams: maxStreams}, ht, ConnectOptions{})
+func setUp(t *testing.T, port int, ht hType) (*server, *http2Client, func()) {
+return setUpWithOptions(t, port, &ServerConfig{}, ht, ConnectOptions{})
 }
 
 func setUpWithOptions(t *testing.T, port int, sc *ServerConfig, ht hType, copts ConnectOptions) (*server, *http2Client, func()) {
@@ -539,7 +542,7 @@ func (s) TestInflightStreamClosing(t *testing.T) {
 
 // Tests that when streamID > MaxStreamId, the current client transport drains.
 func (s) TestClientTransportDrainsAfterStreamIDExhausted(t *testing.T) {
-server, ct, cancel := setUp(t, 0, math.MaxUint32, normal)
+server, ct, cancel := setUp(t, 0, normal)
 defer cancel()
 defer server.stop()
 callHdr := &CallHdr{
@@ -584,7 +587,7 @@ func (s) TestClientTransportDrainsAfterStreamIDExhausted(t *testing.T) {
 }
 
 func (s) TestClientSendAndReceive(t *testing.T) {
-server, ct, cancel := setUp(t, 0, math.MaxUint32, normal)
+server, ct, cancel := setUp(t, 0, normal)
 defer cancel()
 callHdr := &CallHdr{
 Host: "localhost",
@@ -624,7 +627,7 @@ func (s) TestClientSendAndReceive(t *testing.T) {
 }
 
 func (s) TestClientErrorNotify(t *testing.T) {
-server, ct, cancel := setUp(t, 0, math.MaxUint32, normal)
+server, ct, cancel := setUp(t, 0, normal)
 defer cancel()
 go server.stop()
 // ct.reader should detect the error and activate ct.Error().
@@ -658,7 +661,7 @@ func performOneRPC(ct ClientTransport) {
 }
 
 func (s) TestClientMix(t *testing.T) {
-s, ct, cancel := setUp(t, 0, math.MaxUint32, normal)
+s, ct, cancel := setUp(t, 0, normal)
 defer cancel()
 time.AfterFunc(time.Second, s.stop)
 go func(ct ClientTransport) {
@@ -672,7 +675,7 @@ func (s) TestClientMix(t *testing.T) {
 }
 
 func (s) TestLargeMessage(t *testing.T) {
-server, ct, cancel := setUp(t, 0, math.MaxUint32, normal)
+server, ct, cancel := setUp(t, 0, normal)
 defer cancel()
 callHdr := &CallHdr{
 Host: "localhost",
@@ -807,7 +810,7 @@ func (s) TestLargeMessageWithDelayRead(t *testing.T) {
 // proceed until they complete naturally, while not allowing creation of new
 // streams during this window.
 func (s) TestGracefulClose(t *testing.T) {
-server, ct, cancel := setUp(t, 0, math.MaxUint32, pingpong)
+server, ct, cancel := setUp(t, 0, pingpong)
 defer cancel()
 defer func() {
 // Stop the server's listener to make the server's goroutines terminate
@@ -873,7 +876,7 @@ func (s) TestGracefulClose(t *testing.T) {
 }
 
 func (s) TestLargeMessageSuspension(t *testing.T) {
-server, ct, cancel := setUp(t, 0, math.MaxUint32, suspended)
+server, ct, cancel := setUp(t, 0, suspended)
 defer cancel()
 callHdr := &CallHdr{
 Host: "localhost",
@@ -981,7 +984,7 @@ func (s) TestMaxStreams(t *testing.T) {
 }
 
 func (s) TestServerContextCanceledOnClosedConnection(t *testing.T) {
-server, ct, cancel := setUp(t, 0, math.MaxUint32, suspended)
+server, ct, cancel := setUp(t, 0, suspended)
 defer cancel()
 callHdr := &CallHdr{
 Host: "localhost",
@@ -1453,7 +1456,7 @@ func (s) TestClientWithMisbehavedServer(t *testing.T) {
 var encodingTestStatus = status.New(codes.Internal, "\n")
 
 func (s) TestEncodingRequiredStatus(t *testing.T) {
-server, ct, cancel := setUp(t, 0, math.MaxUint32, encodingRequiredStatus)
+server, ct, cancel := setUp(t, 0, encodingRequiredStatus)
 defer cancel()
 callHdr := &CallHdr{
 Host: "localhost",
@@ -1481,7 +1484,7 @@ func (s) TestEncodingRequiredStatus(t *testing.T) {
 }
 
 func (s) TestInvalidHeaderField(t *testing.T) {
-server, ct, cancel := setUp(t, 0, math.MaxUint32, invalidHeaderField)
+server, ct, cancel := setUp(t, 0, invalidHeaderField)
 defer cancel()
 callHdr := &CallHdr{
 Host: "localhost",
@@ -1503,7 +1506,7 @@ func (s) TestInvalidHeaderField(t *testing.T) {
 }
 
 func (s) TestHeaderChanClosedAfterReceivingAnInvalidHeader(t *testing.T) {
-server, ct, cancel := setUp(t, 0, math.MaxUint32, invalidHeaderField)
+server, ct, cancel := setUp(t, 0, invalidHeaderField)
 defer cancel()
 defer server.stop()
 defer ct.Close(fmt.Errorf("closed manually by test"))
@@ -2171,7 +2174,7 @@ func (s) TestPingPong1MB(t *testing.T) {
 
 // This is a stress-test of flow control logic.
 func runPingPongTest(t *testing.T, msgSize int) {
-server, client, cancel := setUp(t, 0, 0, pingpong)
+server, client, cancel := setUp(t, 0, pingpong)
 defer cancel()
 defer server.stop()
 defer client.Close(fmt.Errorf("closed manually by test"))
@@ -2253,7 +2256,7 @@ func (s) TestHeaderTblSize(t *testing.T) {
 }
 }()
 
-server, ct, cancel := setUp(t, 0, math.MaxUint32, normal)
+server, ct, cancel := setUp(t, 0, normal)
 defer cancel()
 defer ct.Close(fmt.Errorf("closed manually by test"))
 defer server.stop()
@@ -2612,7 +2615,7 @@ func TestConnectionError_Unwrap(t *testing.T) {
 
 func (s) TestPeerSetInServerContext(t *testing.T) {
 // create client and server transports.
-server, client, cancel := setUp(t, 0, math.MaxUint32, normal)
+server, client, cancel := setUp(t, 0, normal)
 defer cancel()
 defer server.stop()
 defer client.Close(fmt.Errorf("closed manually by test"))

server.go

@@ -115,12 +115,6 @@ type serviceInfo struct {
 mdata interface{}
 }
 
-type serverWorkerData struct {
-st transport.ServerTransport
-wg *sync.WaitGroup
-stream *transport.Stream
-}
-
 // Server is a gRPC server to serve RPC requests.
 type Server struct {
 opts serverOptions
@@ -145,7 +139,7 @@ type Server struct {
 channelzID *channelz.Identifier
 czData *channelzData
 
-serverWorkerChannel chan *serverWorkerData
+serverWorkerChannel chan func()
 }
 
 type serverOptions struct {
@@ -178,6 +172,7 @@ type serverOptions struct {
 }
 
 var defaultServerOptions = serverOptions{
+maxConcurrentStreams: math.MaxUint32,
 maxReceiveMessageSize: defaultServerMaxReceiveMessageSize,
 maxSendMessageSize: defaultServerMaxSendMessageSize,
 connectionTimeout: 120 * time.Second,
@@ -389,6 +384,9 @@ func MaxSendMsgSize(m int) ServerOption {
 // MaxConcurrentStreams returns a ServerOption that will apply a limit on the number
 // of concurrent streams to each ServerTransport.
 func MaxConcurrentStreams(n uint32) ServerOption {
+if n == 0 {
+n = math.MaxUint32
+}
 return newFuncServerOption(func(o *serverOptions) {
 o.maxConcurrentStreams = n
 })
@@ -590,24 +588,19 @@ const serverWorkerResetThreshold = 1 << 16
 // [1] https://github.com/golang/go/issues/18138
 func (s *Server) serverWorker() {
 for completed := 0; completed < serverWorkerResetThreshold; completed++ {
-data, ok := <-s.serverWorkerChannel
+f, ok := <-s.serverWorkerChannel
 if !ok {
 return
 }
-s.handleSingleStream(data)
+f()
 }
 go s.serverWorker()
 }
 
-func (s *Server) handleSingleStream(data *serverWorkerData) {
-defer data.wg.Done()
-s.handleStream(data.st, data.stream, s.traceInfo(data.st, data.stream))
-}
-
 // initServerWorkers creates worker goroutines and a channel to process incoming
 // connections to reduce the time spent overall on runtime.morestack.
 func (s *Server) initServerWorkers() {
-s.serverWorkerChannel = make(chan *serverWorkerData)
+s.serverWorkerChannel = make(chan func())
 for i := uint32(0); i < s.opts.numServerWorkers; i++ {
 go s.serverWorker()
 }
@@ -966,21 +959,26 @@ func (s *Server) serveStreams(st transport.ServerTransport) {
 defer st.Close(errors.New("finished serving streams for the server transport"))
 var wg sync.WaitGroup
 
+streamQuota := newHandlerQuota(s.opts.maxConcurrentStreams)
 st.HandleStreams(func(stream *transport.Stream) {
 wg.Add(1)
+
+streamQuota.acquire()
+f := func() {
+defer streamQuota.release()
+defer wg.Done()
+s.handleStream(st, stream, s.traceInfo(st, stream))
+}
+
 if s.opts.numServerWorkers > 0 {
-data := &serverWorkerData{st: st, wg: &wg, stream: stream}
 select {
-case s.serverWorkerChannel <- data:
+case s.serverWorkerChannel <- f:
 return
 default:
 // If all stream workers are busy, fallback to the default code path.
 }
 }
-go func() {
-defer wg.Done()
-s.handleStream(st, stream, s.traceInfo(st, stream))
-}()
+go f()
 }, func(ctx context.Context, method string) context.Context {
 if !EnableTracing {
 return ctx
@@ -2075,3 +2073,32 @@ func validateSendCompressor(name, clientCompressors string) error {
 }
 return fmt.Errorf("client does not support compressor %q", name)
 }
+
+// atomicSemaphore implements a blocking, counting semaphore. acquire should be
+// called synchronously; release may be called asynchronously.
+type atomicSemaphore struct {
+n int64 // accessed atomically
+wait chan struct{}
+}
+
+func (q *atomicSemaphore) acquire() {
+if atomic.AddInt64(&q.n, -1) < 0 {
+// We ran out of quota. Block until a release happens.
+<-q.wait
+}
+}
+
+func (q *atomicSemaphore) release() {
+// N.B. the "<= 0" check below should allow for this to work with multiple
+// concurrent calls to acquire, but also note that with synchronous calls to
+// acquire, as our system does, n will never be less than -1. There are
+// fairness issues (queuing) to consider if this was to be generalized.
+if atomic.AddInt64(&q.n, -1) <= 0 {
+// An acquire was waiting on us. Unblock it.
+q.wait <- struct{}{}
+}
+}
+
+func newHandlerQuota(n uint32) *atomicSemaphore {
+return &atomicSemaphore{n: int64(n), wait: make(chan struct{}, 1)}
+}