Refactor APF handler in preparation for dynamic retryAfter

staging/src/k8s.io/apiserver/pkg/server/filters/maxinflight.go

@@ -34,7 +34,6 @@ import (
 
 const (
 	// Constant for the retry-after interval on rate limiting.
-	// TODO: maybe make this dynamic? or user-adjustable?
 	retryAfter = "1"
 
 	// How often inflight usage metric should be updated. Because
@@ -210,7 +209,7 @@ func WithMaxInFlightLimit(
 				// We need to split this data between buckets used for throttling.
 				metrics.RecordDroppedRequest(r, requestInfo, metrics.APIServerComponent, isMutatingRequest)
 				metrics.RecordRequestTermination(r, requestInfo, metrics.APIServerComponent, http.StatusTooManyRequests)
-				tooManyRequests(r, w)
+				tooManyRequests(r, w, retryAfter)
 			}
 		}
 	})
@@ -221,9 +220,3 @@ func WithMaxInFlightLimit(
 func StartMaxInFlightWatermarkMaintenance(stopCh <-chan struct{}) {
 	startWatermarkMaintenance(watermark, stopCh)
 }
-
-func tooManyRequests(req *http.Request, w http.ResponseWriter) {
-	// Return a 429 status indicating "Too Many Requests"
-	w.Header().Set("Retry-After", retryAfter)
-	http.Error(w, "Too many requests, please try again later.", http.StatusTooManyRequests)
-}

staging/src/k8s.io/apiserver/pkg/server/filters/priority-and-fairness.go

@@ -67,26 +67,14 @@ func truncateLogField(s string) string {
 
 var initAPFOnce sync.Once
 
-// WithPriorityAndFairness limits the number of in-flight
+type priorityAndFairnessHandler struct {
-// requests in a fine-grained way.
+	handler                 http.Handler
-func WithPriorityAndFairness(
+	longRunningRequestCheck apirequest.LongRunningRequestCheck
-	handler http.Handler,
+	fcIfc                   utilflowcontrol.Interface
-	longRunningRequestCheck apirequest.LongRunningRequestCheck,
+	workEstimator           flowcontrolrequest.WorkEstimatorFunc
-	fcIfc utilflowcontrol.Interface,
+}
-	workEstimator flowcontrolrequest.WorkEstimatorFunc,
+
-) http.Handler {
+func (h *priorityAndFairnessHandler) Handle(w http.ResponseWriter, r *http.Request) {
-	if fcIfc == nil {
-		klog.Warningf("priority and fairness support not found, skipping")
-		return handler
-	}
-	initAPFOnce.Do(func() {
-		initMaxInFlight(0, 0)
-		// Fetching these gauges is delayed until after their underlying metric has been registered
-		// so that this latches onto the efficient implementation.
-		waitingMark.readOnlyObserver = fcmetrics.GetWaitingReadonlyConcurrency()
-		waitingMark.mutatingObserver = fcmetrics.GetWaitingMutatingConcurrency()
-	})
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 	ctx := r.Context()
 	requestInfo, ok := apirequest.RequestInfoFrom(ctx)
 	if !ok {
@@ -102,9 +90,9 @@ func WithPriorityAndFairness(
 	isWatchRequest := watchVerbs.Has(requestInfo.Verb)
 
 	// Skip tracking long running non-watch requests.
-		if longRunningRequestCheck != nil && longRunningRequestCheck(r, requestInfo) && !isWatchRequest {
+	if h.longRunningRequestCheck != nil && h.longRunningRequestCheck(r, requestInfo) && !isWatchRequest {
 		klog.V(6).Infof("Serving RequestInfo=%#+v, user.Info=%#+v as longrunning\n", requestInfo, user)
-			handler.ServeHTTP(w, r)
+		h.handler.ServeHTTP(w, r)
 		return
 	}
 
@@ -114,7 +102,8 @@ func WithPriorityAndFairness(
 			FlowSchemaName:    fs.Name,
 			FlowSchemaUID:     fs.UID,
 			PriorityLevelName: pl.Name,
-				PriorityLevelUID:  pl.UID}
+			PriorityLevelUID:  pl.UID,
+		}
 
 		httplog.AddKeyValue(ctx, "apf_pl", truncateLogField(pl.Name))
 		httplog.AddKeyValue(ctx, "apf_fs", truncateLogField(fs.Name))
@@ -126,11 +115,10 @@ func WithPriorityAndFairness(
 			// the request has completed, we should never be here though.
 			klog.ErrorS(fmt.Errorf("workEstimator is being invoked before classification of the request has completed"),
 				"Using empty FlowSchema and PriorityLevelConfiguration name", "verb", r.Method, "URI", r.RequestURI)
-
+			return h.workEstimator(r, "", "")
-				return workEstimator(r, "", "")
 		}
 
-			workEstimate := workEstimator(r, classification.FlowSchemaName, classification.PriorityLevelName)
+		workEstimate := h.workEstimator(r, classification.FlowSchemaName, classification.PriorityLevelName)
 
 		fcmetrics.ObserveWorkEstimatedSeats(classification.PriorityLevelName, classification.FlowSchemaName, workEstimate.MaxSeats())
 		httplog.AddKeyValue(ctx, "apf_iseats", workEstimate.InitialSeats)
@@ -190,7 +178,7 @@ func WithPriorityAndFairness(
 			}
 			// Forget the watcher if it was registered.
 			//
-				// // This is race-free because by this point, one of the following occurred:
+			// This is race-free because by this point, one of the following occurred:
 			// case <-shouldStartWatchCh: execute() completed the assignment to forgetWatch
 			// case <-resultCh: Handle() completed, and Handle() does not return
 			//   while execute() is running
@@ -209,7 +197,7 @@ func WithPriorityAndFairness(
 			served = true
 			setResponseHeaders(classification, w)
 
-				forgetWatch = fcIfc.RegisterWatch(r)
+			forgetWatch = h.fcIfc.RegisterWatch(r)
 
 			// Notify the main thread that we're ready to start the watch.
 			close(shouldStartWatchCh)
@@ -261,14 +249,14 @@ func WithPriorityAndFairness(
 			// Note that Handle will return irrespective of whether the request
 			// executes or is rejected. In the latter case, the function will return
 			// without calling the passed `execute` function.
-				fcIfc.Handle(handleCtx, digest, noteFn, estimateWork, queueNote, execute)
+			h.fcIfc.Handle(handleCtx, digest, noteFn, estimateWork, queueNote, execute)
 		}()
 
 		select {
 		case <-shouldStartWatchCh:
 			watchCtx := utilflowcontrol.WithInitializationSignal(ctx, watchInitializationSignal)
 			watchReq = r.WithContext(watchCtx)
-				handler.ServeHTTP(w, watchReq)
+			h.handler.ServeHTTP(w, watchReq)
 			// Protect from the situation when request will not reach storage layer
 			// and the initialization signal will not be send.
 			// It has to happen before waiting on the resultCh below.
@@ -289,10 +277,10 @@ func WithPriorityAndFairness(
 			served = true
 			setResponseHeaders(classification, w)
 
-				handler.ServeHTTP(w, r)
+			h.handler.ServeHTTP(w, r)
 		}
 
-			fcIfc.Handle(ctx, digest, noteFn, estimateWork, queueNote, execute)
+		h.fcIfc.Handle(ctx, digest, noteFn, estimateWork, queueNote, execute)
 	}
 
 	if !served {
@@ -300,9 +288,37 @@ func WithPriorityAndFairness(
 
 		epmetrics.RecordDroppedRequest(r, requestInfo, epmetrics.APIServerComponent, isMutatingRequest)
 		epmetrics.RecordRequestTermination(r, requestInfo, epmetrics.APIServerComponent, http.StatusTooManyRequests)
-			tooManyRequests(r, w)
+		tooManyRequests(r, w, retryAfter)
 	}
+}
+
+// WithPriorityAndFairness limits the number of in-flight
+// requests in a fine-grained way.
+func WithPriorityAndFairness(
+	handler http.Handler,
+	longRunningRequestCheck apirequest.LongRunningRequestCheck,
+	fcIfc utilflowcontrol.Interface,
+	workEstimator flowcontrolrequest.WorkEstimatorFunc,
+) http.Handler {
+	if fcIfc == nil {
+		klog.Warningf("priority and fairness support not found, skipping")
+		return handler
+	}
+	initAPFOnce.Do(func() {
+		initMaxInFlight(0, 0)
+		// Fetching these gauges is delayed until after their underlying metric has been registered
+		// so that this latches onto the efficient implementation.
+		waitingMark.readOnlyObserver = fcmetrics.GetWaitingReadonlyConcurrency()
+		waitingMark.mutatingObserver = fcmetrics.GetWaitingMutatingConcurrency()
 	})
+
+	priorityAndFairnessHandler := &priorityAndFairnessHandler{
+		handler:                 handler,
+		longRunningRequestCheck: longRunningRequestCheck,
+		fcIfc:                   fcIfc,
+		workEstimator:           workEstimator,
+	}
+	return http.HandlerFunc(priorityAndFairnessHandler.Handle)
 }
 
 // StartPriorityAndFairnessWatermarkMaintenance starts the goroutines to observe and maintain watermarks for
@@ -323,3 +339,9 @@ func setResponseHeaders(classification *PriorityAndFairnessClassification, w htt
 	w.Header().Set(flowcontrol.ResponseHeaderMatchedPriorityLevelConfigurationUID, string(classification.PriorityLevelUID))
 	w.Header().Set(flowcontrol.ResponseHeaderMatchedFlowSchemaUID, string(classification.FlowSchemaUID))
 }
+
+func tooManyRequests(req *http.Request, w http.ResponseWriter, retryAfter string) {
+	// Return a 429 status indicating "Too Many Requests"
+	w.Header().Set("Retry-After", retryAfter)
+	http.Error(w, "Too many requests, please try again later.", http.StatusTooManyRequests)
+}