forge/routers/common/qos.go

// Copyright 2026 Hanzo AI, Inc. All rights reserved.
// SPDX-License-Identifier: MIT

package common

import (
	"context"
	"fmt"
	"github.com/hanzoai/git/modules/web"
	"net/http"
	"strings"

	user_model "github.com/hanzoai/git/models/user"
	"github.com/hanzoai/git/modules/log"
	"github.com/hanzoai/git/modules/setting"
	"github.com/hanzoai/git/modules/templates"
	"github.com/hanzoai/git/modules/web/middleware"
	"github.com/hanzoai/git/modules/web/routing"

	"github.com/bohde/codel"
)

const tplStatus503 templates.TplName = "status/503"

type Priority int

func (p Priority) String() string {
	switch p {
	case HighPriority:
		return "high"
	case DefaultPriority:
		return "default"
	case LowPriority:
		return "low"
	default:
		return fmt.Sprintf("%d", p)
	}
}

const (
	LowPriority     = Priority(-10)
	DefaultPriority = Priority(0)
	HighPriority    = Priority(10)
)

// QoS implements quality of service for requests, based upon whether
// or not the user is logged in. All traffic may get dropped, and
// anonymous users are deprioritized.
func QoS() func(next http.Handler) http.Handler {
	if !setting.Service.QoS.Enabled {
		return nil
	}

	maxOutstanding := setting.Service.QoS.MaxInFlightRequests
	if maxOutstanding <= 0 {
		maxOutstanding = 10
	}

	c := codel.NewPriority(codel.Options{
		// The maximum number of waiting requests.
		MaxPending: setting.Service.QoS.MaxWaitingRequests,
		// The maximum number of in-flight requests.
		MaxOutstanding: maxOutstanding,
		// The target latency that a blocked request should wait
		// for. After this, it might be dropped.
		TargetLatency: setting.Service.QoS.TargetWaitTime,
	})

	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
			ctx := req.Context()
			reqRecordInfo := routing.GetRequestRecordInfo(ctx)
			priority := requestPriority(ctx)

			// Check if the request can begin processing.
			err := c.Acquire(ctx, int(priority))
			if err != nil {
				log.Error("QoS error, dropping request of priority %s: %v", priority, err)
				renderServiceUnavailable(w, req)
				return
			}

			// Release long-polling immediately, so they don't always take up an in-flight request
			if reqRecordInfo.IsLongPolling {
				c.Release()
			} else {
				defer c.Release()
			}

			next.ServeHTTP(w, req)
		})
	}
}

// requestPriority assigns a priority value for a request based upon
// whether the user is logged in and how expensive the endpoint is
func requestPriority(ctx context.Context) Priority {
	// If the user is logged in, assign high priority.
	data := middleware.GetContextData(ctx)
	if _, ok := data[middleware.ContextDataKeySignedUser].(*user_model.User); ok {
		return HighPriority
	}

	rctx := web.GetRouteContext(ctx)
	if rctx == nil {
		return DefaultPriority
	}

	// If we're operating in the context of a repo, assign low priority
	routePattern := rctx.RoutePattern()
	if strings.HasPrefix(routePattern, "/{username}/{reponame}/") {
		return LowPriority
	}

	return DefaultPriority
}

// renderServiceUnavailable will render an HTTP 503 Service
// Unavailable page, providing HTML if the client accepts it.
func renderServiceUnavailable(w http.ResponseWriter, req *http.Request) {
	ctxData := middleware.GetContextData(req.Context())
	renderServerErrorPage(w, req, http.StatusServiceUnavailable, tplStatus503, ctxData, "Service Unavailable")
}