From 1d47d3721d77e82720c885abfa92ca8818052ecc Mon Sep 17 00:00:00 2001 From: yandu Date: Tue, 11 Nov 2025 17:19:22 -0500 Subject: [PATCH] We should hold onto requests if there are too many in-flight requests going to any single sandbox and we can/want to scale up --- internal/controller/controller.go | 26 ++++ internal/controller/scale.go | 26 ++-- internal/controller/server.go | 203 +++++++++++++++++++++++++---- internal/controller/server_test.go | 142 ++++++++++++++++++++ internal/function/heartbeat.go | 7 +- internal/router/router.go | 84 ++++++++---- internal/router/router_test.go | 12 +- 7 files changed, 436 insertions(+), 64 deletions(-) diff --git a/internal/controller/controller.go b/internal/controller/controller.go index 829c0297..ce8b0740 100644 --- a/internal/controller/controller.go +++ b/internal/controller/controller.go @@ -32,6 +32,8 @@ import ( kubernetesmetrics "k8s.io/metrics/pkg/client/clientset/versioned" ) +const scaleActivityWindow = 5 * time.Second + type namespaceLister struct { podIndexer cache.Indexer podLister listerv1.PodLister @@ -51,6 +53,7 @@ type Controller struct { scaleMu *xsync.Map[function.Function, *sync.Mutex] routerHeartbeats *xsync.Map[function.Function, RouterHeartbeats] stabilizationWindows *xsync.Map[function.Function, *StabilizationWindow] + scaleActivity *xsync.Map[function.Function, time.Time] } func New(newClientFunc NewClientFunc, kubernetes kubernetes.Interface, kubernetesMetrics kubernetesmetrics.Interface) *Controller { @@ -64,6 +67,7 @@ func New(newClientFunc NewClientFunc, kubernetes kubernetes.Interface, kubernete scaleMu: xsync.NewMap[function.Function, *sync.Mutex](), routerHeartbeats: xsync.NewMap[function.Function, RouterHeartbeats](), stabilizationWindows: xsync.NewMap[function.Function, *StabilizationWindow](), + scaleActivity: xsync.NewMap[function.Function, time.Time](), } } @@ -99,6 +103,28 @@ func (ctrl *Controller) getControllerClient(ip string) Client { return controllerClient } +func (ctrl *Controller) markScaleActivity(fn function.Function) { + if ctrl.scaleActivity == nil { + return + } + ctrl.scaleActivity.Store(fn, time.Now()) +} + +func (ctrl *Controller) lastScaleActivity(fn function.Function) (time.Time, bool) { + if ctrl.scaleActivity == nil { + return time.Time{}, false + } + return ctrl.scaleActivity.Load(fn) +} + +func (ctrl *Controller) isRecentlyScaling(fn function.Function) bool { + last, ok := ctrl.lastScaleActivity(fn) + if !ok { + return false + } + return time.Since(last) <= scaleActivityWindow +} + func (ctrl *Controller) startInformers(ctx context.Context) error { ctx, span := telemetry.Trace(ctx, "controller.start_informers") defer span.End() diff --git a/internal/controller/scale.go b/internal/controller/scale.go index ff3dd7ca..498c30cc 100644 --- a/internal/controller/scale.go +++ b/internal/controller/scale.go @@ -46,28 +46,28 @@ var ( Subsystem: "controller", Name: "waiting_for_unassigned_pods", Help: "The number of functions that are waiting for an unassigned pod", - }, []string{"function_deployment"}) + }, []string{"function_deployment", "function_tenant"}) assignmentsTotal = promauto.NewCounterVec(prometheus.CounterOpts{ Namespace: "skipper", Subsystem: "controller", Name: "assignments_total", Help: "The number of times the controller has assigned a pod to a function", - }, []string{"function_deployment"}) + }, []string{"function_deployment", "function_tenant"}) scaleUpsTotal = promauto.NewCounterVec(prometheus.CounterOpts{ Namespace: "skipper", Subsystem: "controller", Name: "scale_ups_total", Help: "The number of times the controller has scaled up a function", - }, []string{"function_deployment"}) + }, []string{"function_deployment", "function_tenant"}) scaleDownsTotal = promauto.NewCounterVec(prometheus.CounterOpts{ Namespace: "skipper", Subsystem: "controller", Name: "scale_downs_total", Help: "The number of times the controller has scaled down a function", - }, []string{"function_deployment"}) + }, []string{"function_deployment", "function_tenant"}) ) func (ctrl *Controller) scaleNamespace(ctx context.Context, namespace string) error { @@ -317,7 +317,8 @@ func (ctrl *Controller) scale(ctx context.Context, fn function.Function, decisio } log.Info(ctx, "scaling function up") - scaleUpsTotal.WithLabelValues(fn.Deployment).Add(float64(decision.DesiredInstances - len(readyInstances))) + ctrl.markScaleActivity(fn) + scaleUpsTotal.WithLabelValues(fn.Deployment, fn.Tenant).Add(float64(decision.DesiredInstances - len(readyInstances))) for range decision.DesiredInstances - len(readyInstances) { instance, err := ctrl.assignPod(ctx, fn) @@ -329,7 +330,7 @@ func (ctrl *Controller) scale(ctx context.Context, fn function.Function, decisio } else { // we either need to scale down or we're already at the desired number of instances but have extra unready instances log.Info(ctx, "scaling function down") - scaleDownsTotal.WithLabelValues(fn.Deployment).Add(float64(len(readyInstances) + len(unreadyInstances) - decision.DesiredInstances)) + scaleDownsTotal.WithLabelValues(fn.Deployment, fn.Tenant).Add(float64(len(readyInstances) + len(unreadyInstances) - decision.DesiredInstances)) // delete all unready instances for _, unreadyInstance := range unreadyInstances { @@ -360,7 +361,7 @@ func (ctrl *Controller) assignPod(ctx context.Context, fn function.Function) (in ctx, span := telemetry.Trace(ctx, "controller.assign_pod") defer span.End() - assignmentsTotal.WithLabelValues(fn.Deployment).Inc() + assignmentsTotal.WithLabelValues(fn.Deployment, fn.Tenant).Inc() GET_UNASSIGNED_POD: var pod *v1.Pod @@ -469,8 +470,8 @@ func (ctrl *Controller) getUnassignedPod(ctx context.Context, fn function.Functi ctx, span := telemetry.Trace(ctx, "controller.get_unassigned_pod") defer span.End() - waitingForUnassignedPods.WithLabelValues(fn.Deployment).Inc() - defer waitingForUnassignedPods.WithLabelValues(fn.Deployment).Dec() + waitingForUnassignedPods.WithLabelValues(fn.Deployment, fn.Tenant).Inc() + defer waitingForUnassignedPods.WithLabelValues(fn.Deployment, fn.Tenant).Dec() return timer.Poll(ctx, 250*time.Millisecond, func(ctx context.Context) (*v1.Pod, error) { unassignedPods, err := ctrl.getUnassignedPods(fn) @@ -745,12 +746,19 @@ type RouterHeartbeats map[string]function.Heartbeat // Combined returns a heartbeat that is the sum of all the heartbeats from all the routers func (r RouterHeartbeats) Combined() function.Heartbeat { var combined function.Heartbeat + combined.InFlightPerInstance = make(map[string]int) for _, heartbeat := range r { combined.Function = heartbeat.Function combined.InFlightRequests += heartbeat.InFlightRequests if combined.Timestamp.Before(heartbeat.Timestamp) { combined.Timestamp = heartbeat.Timestamp } + for instance, count := range heartbeat.InFlightPerInstance { + combined.InFlightPerInstance[instance] += count + } + } + if len(combined.InFlightPerInstance) == 0 { + combined.InFlightPerInstance = map[string]int{} } return combined } diff --git a/internal/controller/server.go b/internal/controller/server.go index f9577b39..9f9f3f4d 100644 --- a/internal/controller/server.go +++ b/internal/controller/server.go @@ -2,6 +2,7 @@ package controller import ( "context" + "math" "math/rand" "net/http" "slices" @@ -19,12 +20,19 @@ import ( "go.opentelemetry.io/otel/attribute" ) +const ( + instanceOverloadHeadroom = 0.25 // 25% headroom for instance overload + instanceThrottleInterval = 100 * time.Millisecond + instanceThrottleMaxWait = 2 * time.Second + instanceThrottleDeadlineBuffer = 250 * time.Millisecond +) + var heartbeatsCounter = promauto.NewCounterVec(prometheus.CounterOpts{ Namespace: "skipper", Subsystem: "controller", Name: "heartbeats_total", Help: "The number of heartbeats received by the controller", -}, []string{"function_deployment"}) +}, []string{"function_deployment", "function_tenant"}) func (ctrl *Controller) Handler() http.Handler { mux := http.NewServeMux() @@ -54,35 +62,63 @@ func (ctrl *Controller) handleInstance(rw http.ResponseWriter, req *http.Request ctx = log.With(ctx, key.Function.Field(fn)) ctx = telemetry.WithPropagatedAttributes(ctx, key.Function.Attributes(fn)...) - instances, err := ctrl.getReadyInstances(fn) - if err != nil { - log.Error(ctx, "failed to get instances", key.Error.Field(err)) - http.Error(rw, err.Error(), http.StatusInternalServerError) - return - } - - telemetry.SetAttributes(ctx, attribute.Bool("has_instances", len(instances) > 0)) + selectionStart := time.Now() + var ( + instances []*function.Instance + loads map[string]int + ) - for len(instances) == 0 { - if instances, err = ctrl.scale(ctx, fn, ScalingDecision{ - DesiredInstances: 1, - UnclampedDesiredInstances: 1, - Reason: "no ready instances", - }); err != nil { - log.Error(ctx, "failed to scale function", key.Error.Field(err)) + for { + instances, err = ctrl.getReadyInstances(fn) + if err != nil { + log.Error(ctx, "failed to get instances", key.Error.Field(err)) http.Error(rw, err.Error(), http.StatusInternalServerError) return } + + telemetry.SetAttributes(ctx, attribute.Bool("has_instances", len(instances) > 0)) + + if len(instances) == 0 { + if _, err = ctrl.scale(ctx, fn, ScalingDecision{ + DesiredInstances: 1, + UnclampedDesiredInstances: 1, + Reason: "no ready instances", + }); err != nil { + log.Error(ctx, "failed to scale function", key.Error.Field(err)) + http.Error(rw, err.Error(), http.StatusInternalServerError) + return + } + continue + } + + if len(instances) > fn.Scale.MaxInstances { + // sort instances by assigned at in descending order (newest first) + slices.SortFunc(instances, func(a, b *function.Instance) int { return b.AssignedAt.Compare(a.AssignedAt) }) + // keep the newest instances + instances = instances[:fn.Scale.MaxInstances] + } + + loads = ctrl.inFlightPerInstance(fn) + + if ctrl.shouldThrottle(ctx, fn, instances, loads, selectionStart) { + totalLoad := totalInFlight(loads) + log.Debug(ctx, "delaying instance selection while scaling", key.InFlightRequests.Field(totalLoad), key.Count.Field(len(instances))) + if !ctrl.waitForCapacity(ctx, selectionStart) { + log.Debug(ctx, "unable to wait longer for additional capacity", key.InFlightRequests.Field(totalLoad)) + break + } + continue + } + + break } - if len(instances) > fn.Scale.MaxInstances { - // sort instances by assigned at in descending order (newest first) - slices.SortFunc(instances, func(a, b *function.Instance) int { return b.AssignedAt.Compare(a.AssignedAt) }) - // keep the newest instances - instances = instances[:fn.Scale.MaxInstances] + if len(instances) == 0 { + http.Error(rw, "no ready instances", http.StatusServiceUnavailable) + return } - instance := instances[rand.Intn(len(instances))] + instance := ctrl.chooseLeastBusyInstance(fn, instances, loads) rw.Header().Set("Content-Type", "application/json") rw.WriteHeader(http.StatusOK) @@ -91,6 +127,127 @@ func (ctrl *Controller) handleInstance(rw http.ResponseWriter, req *http.Request } } +func (ctrl *Controller) inFlightPerInstance(fn function.Function) map[string]int { + routerHeartbeats, ok := ctrl.routerHeartbeats.Load(fn) + if !ok { + return map[string]int{} + } + + combined := routerHeartbeats.Combined() + if len(combined.InFlightPerInstance) == 0 { + return map[string]int{} + } + + loads := make(map[string]int, len(combined.InFlightPerInstance)) + for instance, count := range combined.InFlightPerInstance { + loads[instance] = count + } + return loads +} + +func (ctrl *Controller) shouldThrottle(ctx context.Context, fn function.Function, instances []*function.Instance, loads map[string]int, start time.Time) bool { + if fn.Scale.TargetInFlightRequests <= 0 { + return false + } + if len(instances) == 0 { + return false + } + + threshold := int(math.Ceil(float64(fn.Scale.TargetInFlightRequests) * (1 + instanceOverloadHeadroom))) + if threshold < 0 { + return false + } + + for _, instance := range instances { + if loads[instance.Name] <= threshold { + return false + } + } + + if !ctrl.isRecentlyScaling(fn) { + return false + } + + return ctrl.canThrottle(ctx, start) +} + +func (ctrl *Controller) canThrottle(ctx context.Context, start time.Time) bool { + if time.Since(start) >= instanceThrottleMaxWait { + return false + } + + if deadline, ok := ctx.Deadline(); ok { + if time.Until(deadline) <= instanceThrottleDeadlineBuffer { + return false + } + } + + return true +} + +func (ctrl *Controller) waitForCapacity(ctx context.Context, start time.Time) bool { + elapsed := time.Since(start) + if elapsed >= instanceThrottleMaxWait { + return false + } + + wait := instanceThrottleInterval + if remaining := instanceThrottleMaxWait - elapsed; remaining < wait { + wait = remaining + } + if wait <= 0 { + return false + } + + select { + case <-ctx.Done(): + return false + case <-time.After(wait): + return true + } +} + +func (ctrl *Controller) chooseLeastBusyInstance(fn function.Function, instances []*function.Instance, loads map[string]int) *function.Instance { + if len(instances) == 0 { + return nil + } + + var candidates []*function.Instance + + // If no threshold is set, use all instances + if fn.Scale.TargetInFlightRequests <= 0 { + candidates = instances + } else { + // Calculate threshold: instances above this are considered overloaded + threshold := int(math.Ceil(float64(fn.Scale.TargetInFlightRequests) * (1 + instanceOverloadHeadroom))) + + // Filter out instances above the threshold + candidates = make([]*function.Instance, 0, len(instances)) + for _, instance := range instances { + load := loads[instance.Name] + if load <= threshold { + candidates = append(candidates, instance) + } + } + + // If no instances are below threshold, fall back to all instances + if len(candidates) == 0 { + candidates = instances + } + } + + // Randomly select from the acceptable instances + return candidates[rand.Intn(len(candidates))] +} + +func totalInFlight(loads map[string]int) int { + total := 0 + for _, count := range loads { + total += count + } + return total +} + func (ctrl *Controller) handleScale(rw http.ResponseWriter, req *http.Request) { ctx := req.Context() fn, err := function.FromHeader(req) @@ -147,7 +304,7 @@ func (ctrl *Controller) handleHeartbeat(rw http.ResponseWriter, req *http.Reques } for _, heartbeat := range heartbeats { - heartbeatsCounter.WithLabelValues(heartbeat.Function.Deployment).Inc() + heartbeatsCounter.WithLabelValues(heartbeat.Function.Deployment, heartbeat.Function.Tenant).Inc() ctrl.routerHeartbeats.Compute(heartbeat.Function, func(routerHeartbeats RouterHeartbeats, loaded bool) (RouterHeartbeats, xsync.ComputeOp) { if !loaded { diff --git a/internal/controller/server_test.go b/internal/controller/server_test.go index 9407f49a..48c5b320 100644 --- a/internal/controller/server_test.go +++ b/internal/controller/server_test.go @@ -262,3 +262,145 @@ func TestHandleHeartbeat(t *testing.T) { }) } } + +func TestShouldThrottle(t *testing.T) { + ctrl := New(nil, fake.NewClientset(), nil) + fn := fixture.NewFunction() + fn.Scale.TargetInFlightRequests = 4 + + instances := []*function.Instance{ + {Function: fn, Name: "a"}, + {Function: fn, Name: "b"}, + } + + loads := map[string]int{ + "a": 6, + "b": 6, + } + + ctrl.markScaleActivity(fn) + must.True(t, ctrl.shouldThrottle(context.Background(), fn, instances, loads, time.Now())) +} + +func TestShouldThrottleRespectsDeadline(t *testing.T) { + ctrl := New(nil, fake.NewClientset(), nil) + fn := fixture.NewFunction() + fn.Scale.TargetInFlightRequests = 4 + + instances := []*function.Instance{ + {Function: fn, Name: "a"}, + {Function: fn, Name: "b"}, + } + + loads := map[string]int{ + "a": 6, + "b": 6, + } + + ctrl.markScaleActivity(fn) + + ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(instanceThrottleDeadlineBuffer/2)) + defer cancel() + + must.False(t, ctrl.shouldThrottle(ctx, fn, instances, loads, time.Now())) +} + +func TestChooseLeastBusyInstance(t *testing.T) { + ctrl := New(nil, fake.NewClientset(), nil) + fn := fixture.NewFunction() + fn.Scale.TargetInFlightRequests = 50 // threshold will be 50 * 1.25 = 62.5 -> 63 + + t.Run("filters out overloaded instances", func(t *testing.T) { + instances := []*function.Instance{ + {Function: fn, Name: "a"}, + {Function: fn, Name: "b"}, + {Function: fn, Name: "c"}, + } + + loads := map[string]int{ + "a": 100, // above threshold (63) + "b": 1, // below threshold (63) + "c": 1, // below threshold (63) + } + + selected := ctrl.chooseLeastBusyInstance(fn, instances, loads) + // Only b and c should be selected (both below threshold) + must.True(t, selected.Name == "b" || selected.Name == "c") + must.NotEq(t, "a", selected.Name) + }) + + t.Run("all instances below threshold", func(t *testing.T) { + instances := []*function.Instance{ + {Function: fn, Name: "a"}, + {Function: fn, Name: "b"}, + {Function: fn, Name: "c"}, + } + + loads := map[string]int{ + "a": 5, // below threshold (63) + "b": 1, // below threshold (63) + "c": 1, // below threshold (63) + } + + selected := ctrl.chooseLeastBusyInstance(fn, instances, loads) + // All instances are below threshold, so any can be selected + must.True(t, selected.Name == "a" || selected.Name == "b" || selected.Name == "c") + }) + + t.Run("all instances above threshold falls back to all", func(t *testing.T) { + instances := []*function.Instance{ + {Function: fn, Name: "a"}, + {Function: fn, Name: "b"}, + } + + loads := map[string]int{ + "a": 100, // above threshold (63) + "b": 200, // above threshold (63) + } + + selected := ctrl.chooseLeastBusyInstance(fn, instances, loads) + // All instances are above threshold, so fall back to all instances + must.True(t, selected.Name == "a" || selected.Name == "b") + }) + + t.Run("no threshold set uses all instances", func(t *testing.T) { + fnNoThreshold := fixture.NewFunction() + fnNoThreshold.Scale.TargetInFlightRequests = 0 // no threshold + + instances := []*function.Instance{ + {Function: fnNoThreshold, Name: "a"}, + {Function: fnNoThreshold, Name: "b"}, + {Function: fnNoThreshold, Name: "c"}, + } + + loads := map[string]int{ + "a": 1000, // any load is acceptable when no threshold + "b": 5000, + "c": 100, + } + + selected := ctrl.chooseLeastBusyInstance(fnNoThreshold, instances, loads) + // No threshold means all instances are candidates + must.True(t, selected.Name == "a" || selected.Name == "b" || selected.Name == "c") + }) +} + +func TestInFlightPerInstanceAggregates(t *testing.T) { + ctrl := New(nil, fake.NewClientset(), nil) + fn := fixture.NewFunction() + + ctrl.routerHeartbeats.Store(fn, RouterHeartbeats{ + "router-1": { + Function: fn, + InFlightPerInstance: map[string]int{"a": 2}, + }, + "router-2": { + Function: fn, + InFlightPerInstance: map[string]int{"a": 3, "b": 1}, + }, + }) + + loads := ctrl.inFlightPerInstance(fn) + must.Eq(t, 5, loads["a"]) + must.Eq(t, 1, loads["b"]) +} diff --git a/internal/function/heartbeat.go b/internal/function/heartbeat.go index 0fdfd1dd..79aae675 100644 --- a/internal/function/heartbeat.go +++ b/internal/function/heartbeat.go @@ -9,9 +9,10 @@ import ( ) type Heartbeat struct { - Function Function `json:"function"` - Timestamp time.Time `json:"timestamp"` - InFlightRequests int `json:"in_flight_requests"` + Function Function `json:"function"` + Timestamp time.Time `json:"timestamp"` + InFlightRequests int `json:"in_flight_requests"` + InFlightPerInstance map[string]int `json:"in_flight_per_instance,omitempty"` } func (h Heartbeat) Fields() []slog.Attr { diff --git a/internal/router/router.go b/internal/router/router.go index cc35155b..9546b1d7 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -31,21 +31,21 @@ var ( Subsystem: "router", Name: "requests_total", Help: "The number of requests handled by the router", - }, []string{"function_deployment"}) + }, []string{"function_deployment", "function_tenant"}) requestsInFlight = promauto.NewGaugeVec(prometheus.GaugeOpts{ Namespace: "skipper", Subsystem: "router", Name: "requests_in_flight", Help: "The number of requests that are currently being handled by the router", - }, []string{"function_deployment"}) + }, []string{"function_deployment", "function_tenant"}) heartbeatsTotal = promauto.NewCounterVec(prometheus.CounterOpts{ Namespace: "skipper", Subsystem: "router", Name: "heartbeats_total", Help: "The number of heartbeats sent by the router", - }, []string{"function_deployment"}) + }, []string{"function_deployment", "function_tenant"}) ) type Router struct { @@ -92,7 +92,7 @@ func (r *Router) Start(ctx context.Context) { r.heartbeats.Delete(fn) // remove the heartbeat if it hasn't been updated in 3 intervals } else { heartbeats = append(heartbeats, heartbeat) // otherwise, send the heartbeat - heartbeatsTotal.WithLabelValues(fn.Deployment).Inc() + heartbeatsTotal.WithLabelValues(fn.Deployment, fn.Tenant).Inc() } return true }) @@ -119,7 +119,7 @@ func (r *Router) ServeHTTP(rw http.ResponseWriter, req *http.Request) { ctx := log.With(req.Context(), key.Function.Field(fn)) ctx = telemetry.WithPropagatedAttributes(ctx, key.Function.Attributes(fn)...) - requestsTotal.WithLabelValues(fn.Deployment).Inc() + requestsTotal.WithLabelValues(fn.Deployment, fn.Tenant).Inc() // continuously update the heartbeat timestamp for this function while the request is in flight go timer.Loop(ctx, FlagHeartbeatInterval.Value(), func(ctx context.Context) error { @@ -131,24 +131,6 @@ func (r *Router) ServeHTTP(rw http.ResponseWriter, req *http.Request) { return nil }) - // increment the in-flight requests for this function - r.heartbeats.Compute(fn, func(heartbeat function.Heartbeat, _ bool) (function.Heartbeat, xsync.ComputeOp) { - heartbeat.Function = fn - heartbeat.Timestamp = time.Now() - heartbeat.InFlightRequests++ - requestsInFlight.WithLabelValues(fn.Deployment).Inc() - return heartbeat, xsync.UpdateOp - }) - - // decrement the in-flight requests for this function when the request is complete - defer r.heartbeats.Compute(fn, func(heartbeat function.Heartbeat, _ bool) (function.Heartbeat, xsync.ComputeOp) { - heartbeat.Function = fn - heartbeat.Timestamp = time.Now() - heartbeat.InFlightRequests-- - requestsInFlight.WithLabelValues(fn.Deployment).Dec() - return heartbeat, xsync.UpdateOp - }) - r.reverseProxy.ServeHTTP(rw, req.WithContext(function.With(ctx, fn))) } @@ -190,18 +172,19 @@ func (r *Router) RoundTrip(req *http.Request) (*http.Response, error) { log.Warn(ctx, "failed to get instance for function", key.Error.Field(err)) continue } + getInstanceDuration += time.Since(getInstanceStart) ctx = log.With(ctx, key.Instance.Field(instance)) ctx = telemetry.WithPropagatedAttributes(ctx, key.Instance.Attributes(instance)...) - req := req.WithContext(ctx) - req.URL.Scheme = "http" - req.URL.Host = instance.Addr + forwardReq := req.WithContext(ctx) + forwardReq.URL.Scheme = "http" + forwardReq.URL.Host = instance.Addr log.Info(ctx, "forwarding request") start := time.Now() - res, err := r.roundTripper.RoundTrip(req) + res, err := r.forward(forwardReq, fn, instance) duration := time.Since(start) ctx = log.With(ctx, key.Response.Field(res), key.Duration.Field(duration)) @@ -230,6 +213,53 @@ func (r *Router) RoundTrip(req *http.Request) (*http.Response, error) { } } +func (r *Router) forward(req *http.Request, fn function.Function, instance *function.Instance) (*http.Response, error) { + r.incrementInFlight(fn, instance) + defer r.decrementInFlight(fn, instance) + return r.roundTripper.RoundTrip(req) +} + +func (r *Router) incrementInFlight(fn function.Function, instance *function.Instance) { + requestsInFlight.WithLabelValues(fn.Deployment, fn.Tenant).Inc() + + r.heartbeats.Compute(fn, func(heartbeat function.Heartbeat, _ bool) (function.Heartbeat, xsync.ComputeOp) { + heartbeat.Function = fn + heartbeat.Timestamp = time.Now() + heartbeat.InFlightRequests++ + + if heartbeat.InFlightPerInstance == nil { + heartbeat.InFlightPerInstance = make(map[string]int) + } + heartbeat.InFlightPerInstance[instance.Name]++ + + return heartbeat, xsync.UpdateOp + }) +} + +func (r *Router) decrementInFlight(fn function.Function, instance *function.Instance) { + requestsInFlight.WithLabelValues(fn.Deployment, fn.Tenant).Dec() + + r.heartbeats.Compute(fn, func(heartbeat function.Heartbeat, _ bool) (function.Heartbeat, xsync.ComputeOp) { + heartbeat.Function = fn + heartbeat.Timestamp = time.Now() + if heartbeat.InFlightRequests > 0 { + heartbeat.InFlightRequests-- + } + + if heartbeat.InFlightPerInstance != nil && instance != nil { + if count, ok := heartbeat.InFlightPerInstance[instance.Name]; ok { + if count <= 1 { + delete(heartbeat.InFlightPerInstance, instance.Name) + } else { + heartbeat.InFlightPerInstance[instance.Name] = count - 1 + } + } + } + + return heartbeat, xsync.UpdateOp + }) +} + func rewriteRequestHeaders(pr *httputil.ProxyRequest) { function.RemoveHeader(pr.Out) diff --git a/internal/router/router_test.go b/internal/router/router_test.go index 808da679..8297bbb6 100644 --- a/internal/router/router_test.go +++ b/internal/router/router_test.go @@ -305,13 +305,16 @@ func TestHeartbeats(t *testing.T) { done := make(chan struct{}) defer close(done) + var instanceName string mcc := fixture.NewMockControllerClient(t) mcc.HandleInstance(func(ctx context.Context, fn function.Function) (*function.Instance, error) { - return fixture.NewInstance(t, fn, func(rw http.ResponseWriter, req *http.Request) { + instance := fixture.NewInstance(t, fn, func(rw http.ResponseWriter, req *http.Request) { rw.WriteHeader(http.StatusOK) rw.Write([]byte("Hello, " + fn.Tenant)) <-done - }), nil + }) + instanceName = instance.Name + return instance, nil }) mcc.HandleHeartbeat(func(ctx context.Context, routerIP string, heartbeats []function.Heartbeat, forwardedFor ...string) error { if len(heartbeats) == 0 { @@ -328,6 +331,10 @@ func TestHeartbeats(t *testing.T) { must.True(t, heartbeat.Timestamp.After(testStartTime)) if heartbeat.InFlightRequests > 0 { once.Do(func() { + if instanceName != "" && heartbeat.InFlightPerInstance != nil { + must.MapContainsKey(t, heartbeat.InFlightPerInstance, instanceName) + must.True(t, heartbeat.InFlightPerInstance[instanceName] > 0) + } done <- struct{}{} }) } @@ -349,6 +356,7 @@ func TestHeartbeats(t *testing.T) { must.Eq(t, fn, heartbeat.Function) must.True(t, heartbeat.Timestamp.After(testStartTime)) must.Eq(t, 0, heartbeat.InFlightRequests) // ensure the number of in-flight requests is 0 now that the request is complete + must.True(t, len(heartbeat.InFlightPerInstance) == 0) } func TestRetries(t *testing.T) {