Track push inflight requests via grpc server handlers.

[pstibrany]

What this PR does

This PR moves tracking of inflight push requests to gRPC handlers. We install 1) tap handle and 2) stats handler, and they work together to check if another push request is allowed, and track number of ongoing push requests.

Tap handle is called as soon as request headers are received, but before request body is read into memory. At that point we can check and reject request early, or accept it and increase ongoing requests. At this point no interceptor or stats handler has run yet, so if we reject request, the only record of request will be in cortex_ingester_instance_rejected_requests_total metric.

If tap handle allows the request to proceed, only then stats handlers and interceptors are called. Stats handler is also called after request has finished -- that is where we decrease number of inflight requests again.

Big benefit of this change is that we can reject requests early, before they are read into memory. This protects the ingester from crashing when it receives many push requests.

Since we increase number of ongoing requests in Tap handle, are we guaranteed that StatsHandler is called when each time when request finishes?

Tap handle is called in operateHeaders:

mimir/vendor/google.golang.org/grpc/internal/transport/http2_server.go

Lines 567 to 632 in 611e6ba

	if t.inTapHandle != nil {
	var err error
	if s.ctx, err = t.inTapHandle(s.ctx, &tap.Info{FullMethodName: s.method}); err != nil {
	t.mu.Unlock()
	if t.logger.V(logLevel) {
	t.logger.Infof("Aborting the stream early due to InTapHandle failure: %v", err)
	}
	stat, ok := status.FromError(err)
	if !ok {
	stat = status.New(codes.PermissionDenied, err.Error())
	}
	t.controlBuf.put(&earlyAbortStream{
	httpStatus: 200,
	streamID: s.id,
	contentSubtype: s.contentSubtype,
	status: stat,
	rst: !frame.StreamEnded(),
	})
	return nil
	}
	}
	t.activeStreams[streamID] = s
	if len(t.activeStreams) == 1 {
	t.idle = time.Time{}
	}
	t.mu.Unlock()
	if channelz.IsOn() {
	atomic.AddInt64(&t.czData.streamsStarted, 1)
	atomic.StoreInt64(&t.czData.lastStreamCreatedTime, time.Now().UnixNano())
	}
	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{
	FullMethod: s.method,
	RemoteAddr: t.remoteAddr,
	LocalAddr: t.localAddr,
	Compression: s.recvCompress,
	WireLength: int(frame.Header().Length),
	Header: mdata.Copy(),
	}
	sh.HandleRPC(s.ctx, inHeader)
	}
	s.ctxDone = s.ctx.Done()
	s.wq = newWriteQuota(defaultWriteQuota, s.ctxDone)
	s.trReader = &transportReader{
	reader: &recvBufferReader{
	ctx: s.ctx,
	ctxDone: s.ctxDone,
	recv: s.buf,
	freeBuffer: t.bufferPool.put,
	},
	windowHandler: func(n int) {
	t.updateWindow(s, uint32(n))
	},
	}
	// Register the stream with loopy.
	t.controlBuf.put(&registerStream{
	streamID: s.id,
	wq: s.wq,
	})
	handle(s)
	return nil

If request is not rejected, eventually handle function is called at the end of operateHeaders. handle is really a function passed to HandleStreams:

mimir/vendor/google.golang.org/grpc/server.go

Lines 969 to 984 in 611e6ba

	st.HandleStreams(func(stream *transport.Stream) {
	wg.Add(1)
	if s.opts.numServerWorkers > 0 {
	data := &serverWorkerData{st: st, wg: &wg, stream: stream}
	select {
	case s.serverWorkerChannel <- data:
	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))
	}()
	}, func(ctx context.Context, method string) context.Context {

This function will either pass request to worker, or start new goroutine. Either way, processing of the request moves to handleStream method:

mimir/vendor/google.golang.org/grpc/server.go

Lines 1707 to 1770 in 611e6ba

	func (s *Server) handleStream(t transport.ServerTransport, stream *transport.Stream, trInfo *traceInfo) {
	sm := stream.Method()
	if sm != "" && sm[0] == '/' {
	sm = sm[1:]
	}
	pos := strings.LastIndex(sm, "/")
	if pos == -1 {
	if trInfo != nil {
	trInfo.tr.LazyLog(&fmtStringer{"Malformed method name %q", []interface{}{sm}}, true)
	trInfo.tr.SetError()
	}
	errDesc := fmt.Sprintf("malformed method name: %q", stream.Method())
	if err := t.WriteStatus(stream, status.New(codes.Unimplemented, errDesc)); err != nil {
	if trInfo != nil {
	trInfo.tr.LazyLog(&fmtStringer{"%v", []interface{}{err}}, true)
	trInfo.tr.SetError()
	}
	channelz.Warningf(logger, s.channelzID, "grpc: Server.handleStream failed to write status: %v", err)
	}
	if trInfo != nil {
	trInfo.tr.Finish()
	}
	return
	}
	service := sm[:pos]
	method := sm[pos+1:]
	srv, knownService := s.services[service]
	if knownService {
	if md, ok := srv.methods[method]; ok {
	s.processUnaryRPC(t, stream, srv, md, trInfo)
	return
	}
	if sd, ok := srv.streams[method]; ok {
	s.processStreamingRPC(t, stream, srv, sd, trInfo)
	return
	}
	}
	// Unknown service, or known server unknown method.
	if unknownDesc := s.opts.unknownStreamDesc; unknownDesc != nil {
	s.processStreamingRPC(t, stream, nil, unknownDesc, trInfo)
	return
	}
	var errDesc string
	if !knownService {
	errDesc = fmt.Sprintf("unknown service %v", service)
	} else {
	errDesc = fmt.Sprintf("unknown method %v for service %v", method, service)
	}
	if trInfo != nil {
	trInfo.tr.LazyPrintf("%s", errDesc)
	trInfo.tr.SetError()
	}
	if err := t.WriteStatus(stream, status.New(codes.Unimplemented, errDesc)); err != nil {
	if trInfo != nil {
	trInfo.tr.LazyLog(&fmtStringer{"%v", []interface{}{err}}, true)
	trInfo.tr.SetError()
	}
	channelz.Warningf(logger, s.channelzID, "grpc: Server.handleStream failed to write status: %v", err)
	}
	if trInfo != nil {
	trInfo.tr.Finish()
	}
	}

This method can return early with error, without calling stats handler in several cases:

1. method name is malformed. This should not be a problem in our usage, as tap handle only increases inflight push requests for specific and correct method name.
2. when service is unknown. This should also not happen in our case, since tap handle requires Mimir ingester component to be initialized, and it registers itself into gRPC server.
3. if the method is unknown, request would be processed as streaming RCP, not unary. Our unit test covers this. (However in this case, stats handler is actually called).

Once processing moves to processUnaryRPC or processStreamingRPC, we are guaranteed that stats handlers are called when request is finished via defer statement. (Push requests are unary)

TODO:

• make this feature optional, and revert back to previous way of tracking when new tracking is not enabled
• update ingester unit tests (not needed for now, since this new limiting is disabled by default)
• add integration test checking for both ways of limiting push requests

Which issue(s) this PR fixes or relates to

Related to

• WIP: track ingester inflight push requests before unmarshalling the request and interceptors #5962
• Check if ingester can accept push request before reading full request into memory. #5921
• WIP: Add limit for total size of concurrent requests in ingesters's memory, and reject additional requests. #5958

Checklist

• Tests updated
• Documentation added
• CHANGELOG.md updated - the order of entries should be [CHANGE], [FEATURE], [ENHANCEMENT], [BUGFIX]

[dimitarvdimitrov]

LGTM

[jhalterman]

LGTM

[pstibrany]

Merging, will add integration test in separate PR.

[pstibrany]

grafana/dskit#377 moves this code into dskit, and introduces global gRPC server inflight limit + a way of limiting individual methods (like we do in this PR, where only Ingester's Push method is limited).

[pracucci]

YAML parameter and CLI flag don't match. Can you fix it, please?

[pstibrany]

Sure

[pstibrany]

Will be fixed in #6073 (draft for now)

[pracucci]

Why was DoNotLogError{} removed?

[pstibrany]

It was moved from StartPushRequest to PushWithCleanup.

When StartPushRequest is called from gRPC tap handle, no log messages will be logged anyway, because we do logging from interceptor, but interceptor will not run at all.

[pracucci]

When StartPushRequest is called from gRPC tap handle, no log messages will be logged anyway, because we do logging from interceptor, but interceptor will not run at all.

I see. I think it's a bit obscure. Maybe we can improve the code comments to explain it?

[pstibrany]

Will be fixed in #6073 (draft for now)