William Petit
22fda4ed54
All checks were successful
Cadoles/bouncer/pipeline/pr-develop This commit looks good
275 lines
7.3 KiB
Go
275 lines
7.3 KiB
Go
package queue
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"html/template"
|
|
"math/rand"
|
|
"net/http"
|
|
"path/filepath"
|
|
"strconv"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"forge.cadoles.com/Cadoles/go-proxy"
|
|
"forge.cadoles.com/Cadoles/go-proxy/wildcard"
|
|
"forge.cadoles.com/cadoles/bouncer/internal/proxy/director"
|
|
"forge.cadoles.com/cadoles/bouncer/internal/store"
|
|
"github.com/Masterminds/sprig/v3"
|
|
"github.com/google/uuid"
|
|
"github.com/pkg/errors"
|
|
"github.com/prometheus/client_golang/prometheus"
|
|
"gitlab.com/wpetit/goweb/logger"
|
|
)
|
|
|
|
const LayerType store.LayerType = "queue"
|
|
|
|
type Queue struct {
|
|
adapter Adapter
|
|
|
|
defaultKeepAlive time.Duration
|
|
|
|
templateDir string
|
|
loadOnce sync.Once
|
|
tmpl *template.Template
|
|
|
|
refreshJobRunning uint32
|
|
updateMetricsJobRunning uint32
|
|
postKeepAliveDebouncer *DebouncerMap
|
|
}
|
|
|
|
// LayerType implements director.MiddlewareLayer
|
|
func (q *Queue) LayerType() store.LayerType {
|
|
return LayerType
|
|
}
|
|
|
|
// Middleware implements director.MiddlewareLayer
|
|
func (q *Queue) Middleware(layer *store.Layer) proxy.Middleware {
|
|
return func(h http.Handler) http.Handler {
|
|
fn := func(w http.ResponseWriter, r *http.Request) {
|
|
ctx := r.Context()
|
|
|
|
options, err := fromStoreOptions(layer.Options, q.defaultKeepAlive)
|
|
if err != nil {
|
|
logger.Error(ctx, "could not parse layer options", logger.E(errors.WithStack(err)))
|
|
http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
|
|
|
|
return
|
|
}
|
|
|
|
matches := wildcard.MatchAny(r.URL.String(), options.MatchURLs...)
|
|
if !matches {
|
|
h.ServeHTTP(w, r)
|
|
|
|
return
|
|
}
|
|
|
|
defer q.updateMetrics(ctx, layer.Proxy, layer.Name, options)
|
|
|
|
cookieName := q.getCookieName(layer.Name)
|
|
|
|
cookie, err := r.Cookie(cookieName)
|
|
if err != nil && !errors.Is(err, http.ErrNoCookie) {
|
|
logger.Error(ctx, "could not retrieve cookie", logger.E(errors.WithStack(err)))
|
|
}
|
|
|
|
if cookie == nil {
|
|
cookie = &http.Cookie{
|
|
Name: cookieName,
|
|
Value: uuid.NewString(),
|
|
Path: "/",
|
|
}
|
|
|
|
w.Header().Add("Set-Cookie", cookie.String())
|
|
}
|
|
|
|
sessionID := cookie.Value
|
|
queueName := string(layer.Name)
|
|
|
|
rank, err := q.adapter.Touch(ctx, queueName, sessionID)
|
|
if err != nil {
|
|
logger.Error(ctx, "could not retrieve session rank", logger.E(errors.WithStack(err)))
|
|
http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
|
|
|
|
return
|
|
}
|
|
|
|
if rank >= options.Capacity {
|
|
q.renderQueuePage(w, r, queueName, options, rank)
|
|
|
|
return
|
|
}
|
|
|
|
ctx = logger.With(ctx,
|
|
logger.F("queueSessionId", sessionID),
|
|
logger.F("queueName", queueName),
|
|
logger.F("queueSessionRank", rank),
|
|
)
|
|
r = r.WithContext(ctx)
|
|
|
|
h.ServeHTTP(w, r)
|
|
}
|
|
|
|
return http.HandlerFunc(fn)
|
|
}
|
|
}
|
|
|
|
func (q *Queue) updateSessionsMetric(ctx context.Context, proxyName store.ProxyName, layerName store.LayerName) {
|
|
if !atomic.CompareAndSwapUint32(&q.updateMetricsJobRunning, 0, 1) {
|
|
return
|
|
}
|
|
|
|
defer atomic.StoreUint32(&q.updateMetricsJobRunning, 0)
|
|
|
|
queueName := string(layerName)
|
|
|
|
status, err := q.adapter.Status(ctx, queueName)
|
|
if err != nil {
|
|
logger.Error(ctx, "could not retrieve queue status", logger.E(errors.WithStack(err)))
|
|
|
|
return
|
|
}
|
|
|
|
metricQueueSessions.With(
|
|
prometheus.Labels{
|
|
metricLabelLayer: string(layerName),
|
|
metricLabelProxy: string(proxyName),
|
|
},
|
|
).Set(float64(status.Sessions))
|
|
}
|
|
|
|
func (q *Queue) renderQueuePage(w http.ResponseWriter, r *http.Request, queueName string, options *LayerOptions, rank int64) {
|
|
ctx := r.Context()
|
|
|
|
status, err := q.adapter.Status(ctx, queueName)
|
|
if err != nil {
|
|
logger.Error(ctx, "could not retrieve queue status", logger.E(errors.WithStack(err)))
|
|
http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
|
|
|
|
return
|
|
}
|
|
|
|
q.loadOnce.Do(func() {
|
|
pattern := filepath.Join(q.templateDir, "*.gohtml")
|
|
|
|
logger.Info(ctx, "loading queue page templates", logger.F("pattern", pattern))
|
|
|
|
tmpl, err := template.New("").Funcs(sprig.FuncMap()).ParseGlob(pattern)
|
|
if err != nil {
|
|
logger.Error(ctx, "could not load queue templates", logger.E(errors.WithStack(err)))
|
|
|
|
return
|
|
}
|
|
|
|
q.tmpl = tmpl
|
|
})
|
|
|
|
if q.tmpl == nil {
|
|
logger.Error(ctx, "queue page templates not loaded", logger.E(errors.WithStack(err)))
|
|
http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
|
|
|
|
return
|
|
}
|
|
|
|
refreshRate := time.Duration(int64(options.KeepAlive.Seconds()/2)) * time.Second
|
|
|
|
templateData := struct {
|
|
QueueName string
|
|
LayerOptions *LayerOptions
|
|
Rank int64
|
|
CurrentSessions int64
|
|
MaxSessions int64
|
|
RefreshRate time.Duration
|
|
}{
|
|
QueueName: queueName,
|
|
LayerOptions: options,
|
|
Rank: rank + 1,
|
|
CurrentSessions: status.Sessions,
|
|
MaxSessions: options.Capacity,
|
|
RefreshRate: refreshRate,
|
|
}
|
|
|
|
w.Header().Add("Cache-Control", "no-cache")
|
|
w.Header().Add("Retry-After", strconv.FormatInt(int64(refreshRate.Seconds()), 10))
|
|
w.WriteHeader(http.StatusServiceUnavailable)
|
|
|
|
if err := q.tmpl.ExecuteTemplate(w, "queue", templateData); err != nil {
|
|
logger.Error(ctx, "could not render queue page", logger.E(errors.WithStack(err)))
|
|
http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
|
|
|
|
return
|
|
}
|
|
}
|
|
|
|
func (q *Queue) refreshQueue(ctx context.Context, layerName store.LayerName, keepAlive time.Duration) {
|
|
if !atomic.CompareAndSwapUint32(&q.refreshJobRunning, 0, 1) {
|
|
return
|
|
}
|
|
|
|
defer atomic.StoreUint32(&q.refreshJobRunning, 0)
|
|
|
|
if err := q.adapter.Refresh(ctx, string(layerName), keepAlive); err != nil {
|
|
logger.Error(ctx, "could not refresh queue",
|
|
logger.E(errors.WithStack(err)),
|
|
logger.F("queue", layerName),
|
|
)
|
|
}
|
|
}
|
|
|
|
func (q *Queue) updateMetrics(ctx context.Context, proxyName store.ProxyName, layerName store.LayerName, options *LayerOptions) {
|
|
// Update queue capacity metric
|
|
metricQueueCapacity.With(
|
|
prometheus.Labels{
|
|
metricLabelLayer: string(layerName),
|
|
metricLabelProxy: string(proxyName),
|
|
},
|
|
).Set(float64(options.Capacity))
|
|
|
|
// Refresh queue data and metrics
|
|
q.refreshQueue(ctx, layerName, options.KeepAlive)
|
|
q.updateSessionsMetric(ctx, proxyName, layerName)
|
|
|
|
// (Re)schedule an update job after session ttl + semi-random time padding
|
|
// to update metrics after last session expiration
|
|
randDuration := rand.Int63n(int64(options.KeepAlive))
|
|
timePadding := options.KeepAlive/2 + time.Duration(randDuration)
|
|
after := options.KeepAlive + timePadding
|
|
|
|
debouncingKey := fmt.Sprintf("%s/%s", proxyName, layerName)
|
|
|
|
q.postKeepAliveDebouncer.Do(debouncingKey, after, func() {
|
|
ctx := logger.With(
|
|
context.Background(),
|
|
logger.F("proxy", proxyName),
|
|
logger.F("layer", layerName),
|
|
logger.F("after", after),
|
|
)
|
|
|
|
logger.Info(ctx, "running post keep alive refresh job")
|
|
|
|
q.refreshQueue(ctx, layerName, options.KeepAlive)
|
|
q.updateSessionsMetric(ctx, proxyName, layerName)
|
|
})
|
|
}
|
|
|
|
func (q *Queue) getCookieName(layerName store.LayerName) string {
|
|
return fmt.Sprintf("_bouncer_%s_%s", LayerType, layerName)
|
|
}
|
|
|
|
func New(adapter Adapter, funcs ...OptionFunc) *Queue {
|
|
opts := defaultOptions()
|
|
for _, fn := range funcs {
|
|
fn(opts)
|
|
}
|
|
|
|
return &Queue{
|
|
adapter: adapter,
|
|
templateDir: opts.TemplateDir,
|
|
defaultKeepAlive: opts.DefaultKeepAlive,
|
|
postKeepAliveDebouncer: NewDebouncerMap(),
|
|
}
|
|
}
|
|
|
|
var _ director.MiddlewareLayer = &Queue{}
|