feat: refactor layers registration

internal/proxy/director/layer/queue/adapter.go

@@ -0,0 +1,16 @@
+package queue
+
+import (
+	"context"
+	"time"
+)
+
+type Status struct {
+	Sessions int64
+}
+
+type Adapter interface {
+	Touch(ctx context.Context, queueName string, sessionId string) (int64, error)
+	Status(ctx context.Context, queueName string) (*Status, error)
+	Refresh(ctx context.Context, queueName string, keepAlive time.Duration) error
+}

internal/proxy/director/layer/queue/layer_options.go

@@ -0,0 +1,48 @@
+package queue
+
+import (
+	"reflect"
+	"time"
+
+	"forge.cadoles.com/cadoles/bouncer/internal/store"
+	"github.com/mitchellh/mapstructure"
+	"github.com/pkg/errors"
+)
+
+type LayerOptions struct {
+	Capacity  int64         `mapstructure:"capacity"`
+	Matchers  []string      `mapstructure:"matchers"`
+	KeepAlive time.Duration `mapstructure:"keepAlive"`
+}
+
+func fromStoreOptions(storeOptions store.LayerOptions, defaultKeepAlive time.Duration) (*LayerOptions, error) {
+	layerOptions := LayerOptions{
+		Capacity:  1000,
+		Matchers:  []string{"*"},
+		KeepAlive: defaultKeepAlive,
+	}
+
+	config := mapstructure.DecoderConfig{
+		DecodeHook: stringToDurationHook,
+		Result:     &layerOptions,
+	}
+
+	decoder, err := mapstructure.NewDecoder(&config)
+	if err != nil {
+		return nil, err
+	}
+
+	if err := decoder.Decode(storeOptions); err != nil {
+		return nil, errors.WithStack(err)
+	}
+
+	return &layerOptions, nil
+}
+
+func stringToDurationHook(f reflect.Type, t reflect.Type, data interface{}) (interface{}, error) {
+	if t == reflect.TypeOf(*new(time.Duration)) && f == reflect.TypeOf("") {
+		return time.ParseDuration(data.(string))
+	}
+
+	return data, nil
+}

internal/proxy/director/layer/queue/options.go

@@ -0,0 +1,29 @@
+package queue
+
+import "time"
+
+type Options struct {
+	TemplateDir      string
+	DefaultKeepAlive time.Duration
+}
+
+type OptionFunc func(*Options)
+
+func defaultOptions() *Options {
+	return &Options{
+		TemplateDir:      "./templates",
+		DefaultKeepAlive: time.Minute,
+	}
+}
+
+func WithTemplateDir(templateDir string) OptionFunc {
+	return func(o *Options) {
+		o.TemplateDir = templateDir
+	}
+}
+
+func WithDefaultKeepAlive(keepAlive time.Duration) OptionFunc {
+	return func(o *Options) {
+		o.DefaultKeepAlive = keepAlive
+	}
+}

internal/proxy/director/layer/queue/queue.go

@@ -0,0 +1,201 @@
+package queue
+
+import (
+	"context"
+	"fmt"
+	"html/template"
+	"net/http"
+	"path/filepath"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"forge.cadoles.com/Cadoles/go-proxy"
+	"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"
+	"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
+}
+
+// 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
+			}
+
+			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)
+
+			q.refreshQueue(queueName, options.KeepAlive)
+
+			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) 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
+	}
+
+	templateData := struct {
+		QueueName       string
+		LayerOptions    *LayerOptions
+		Rank            int64
+		CurrentSessions int64
+		MaxSessions     int64
+		RefreshRate     int64
+	}{
+		QueueName:       queueName,
+		LayerOptions:    options,
+		Rank:            rank + 1,
+		CurrentSessions: status.Sessions,
+		MaxSessions:     options.Capacity,
+		RefreshRate:     5,
+	}
+
+	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(queueName string, keepAlive time.Duration) {
+	if !atomic.CompareAndSwapUint32(&q.refreshJobRunning, 0, 1) {
+		return
+	}
+
+	go func() {
+		defer atomic.StoreUint32(&q.refreshJobRunning, 0)
+
+		ctx, cancel := context.WithTimeout(context.Background(), keepAlive*2)
+		defer cancel()
+
+		if err := q.adapter.Refresh(ctx, queueName, keepAlive); err != nil {
+			logger.Error(ctx, "could not refresh queue",
+				logger.E(errors.WithStack(err)),
+				logger.F("queue", queueName),
+			)
+		}
+	}()
+}
+
+func (q *Queue) getCookieName(layerName store.LayerName) string {
+	return fmt.Sprintf("_%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,
+	}
+}
+
+var _ director.MiddlewareLayer = &Queue{}

internal/proxy/director/layer/queue/redis/adapter.go

@@ -0,0 +1,167 @@
+package redis
+
+import (
+	"context"
+	"strconv"
+	"strings"
+	"time"
+
+	"forge.cadoles.com/cadoles/bouncer/internal/proxy/director/layer/queue"
+	"github.com/pkg/errors"
+	"github.com/redis/go-redis/v9"
+)
+
+const (
+	keyPrefixQueue = "queue"
+)
+
+type Adapter struct {
+	client     redis.UniversalClient
+	txMaxRetry int
+}
+
+// Refresh implements queue.Adapter
+func (a *Adapter) Refresh(ctx context.Context, queueName string, keepAlive time.Duration) error {
+	lastSeenKey := lastSeenKey(queueName)
+	rankKey := rankKey(queueName)
+
+	err := withTx(ctx, a.client, func(ctx context.Context, tx *redis.Tx) error {
+		expires := time.Now().UTC().Add(-keepAlive)
+
+		cmd := tx.ZRangeByScore(ctx, lastSeenKey, &redis.ZRangeBy{
+			Min: "0",
+			Max: strconv.FormatInt(expires.Unix(), 10),
+		})
+
+		members, err := cmd.Result()
+		if err != nil {
+			return errors.WithStack(err)
+		}
+
+		if len(members) == 0 {
+			return nil
+		}
+
+		anyMembers := make([]any, len(members))
+		for i, m := range members {
+			anyMembers[i] = m
+		}
+
+		if err := tx.ZRem(ctx, rankKey, anyMembers...).Err(); err != nil {
+			return errors.WithStack(err)
+		}
+
+		if err := tx.ZRem(ctx, lastSeenKey, anyMembers...).Err(); err != nil {
+			return errors.WithStack(err)
+		}
+
+		return nil
+	}, rankKey, lastSeenKey)
+	if err != nil {
+		return errors.WithStack(err)
+	}
+
+	return nil
+}
+
+// Touch implements queue.Adapter
+func (a *Adapter) Touch(ctx context.Context, queueName string, sessionId string) (int64, error) {
+	lastSeenKey := lastSeenKey(queueName)
+	rankKey := rankKey(queueName)
+
+	var rank int64
+
+	retry := a.txMaxRetry
+
+	for retry > 0 {
+		err := withTx(ctx, a.client, func(ctx context.Context, tx *redis.Tx) error {
+			now := time.Now().UTC().Unix()
+
+			err := tx.ZAddNX(ctx, rankKey, redis.Z{Score: float64(now), Member: sessionId}).Err()
+			if err != nil {
+				return errors.WithStack(err)
+			}
+
+			err = tx.ZAdd(ctx, lastSeenKey, redis.Z{Score: float64(now), Member: sessionId}).Err()
+			if err != nil {
+				return errors.WithStack(err)
+			}
+
+			val, err := tx.ZRank(ctx, rankKey, sessionId).Result()
+			if err != nil {
+				return errors.WithStack(err)
+			}
+
+			rank = val
+
+			return nil
+		}, rankKey, lastSeenKey)
+		if err != nil {
+			if errors.Is(err, redis.Nil) && retry > 0 {
+				retry--
+
+				continue
+			}
+
+			return 0, errors.WithStack(err)
+		}
+
+		break
+	}
+
+	return rank, nil
+}
+
+// Status implements queue.Adapter
+func (a *Adapter) Status(ctx context.Context, queueName string) (*queue.Status, error) {
+	rankKey := rankKey(queueName)
+
+	status := &queue.Status{}
+
+	cmd := a.client.ZCard(ctx, rankKey)
+	if err := cmd.Err(); err != nil {
+		return nil, errors.WithStack(err)
+	}
+
+	status.Sessions = cmd.Val()
+
+	return status, nil
+}
+
+func NewAdapter(client redis.UniversalClient, txMaxRetry int) *Adapter {
+	return &Adapter{
+		client:     client,
+		txMaxRetry: txMaxRetry,
+	}
+}
+
+var _ queue.Adapter = &Adapter{}
+
+func key(parts ...string) string {
+	return strings.Join(parts, ":")
+}
+
+func rankKey(queueName string) string {
+	return key(keyPrefixQueue, queueName, "rank")
+}
+
+func lastSeenKey(queueName string) string {
+	return key(keyPrefixQueue, queueName, "last_seen")
+}
+
+func withTx(ctx context.Context, client redis.UniversalClient, fn func(ctx context.Context, tx *redis.Tx) error, keys ...string) error {
+	txf := func(tx *redis.Tx) error {
+		if err := fn(ctx, tx); err != nil {
+			return errors.WithStack(err)
+		}
+
+		return nil
+	}
+
+	err := client.Watch(ctx, txf, keys...)
+	if err != nil {
+		return errors.WithStack(err)
+	}
+
+	return nil
+}

internal/proxy/director/layer/queue/schema.go

@@ -0,0 +1,8 @@
+package queue
+
+import (
+	_ "embed"
+)
+
+//go:embed schema/layer-options.json
+var RawLayerOptionsSchema []byte

internal/proxy/director/layer/queue/schema/layer-options.json

@@ -0,0 +1,15 @@
+{
+    "$id": "https://forge.cadoles.com/cadoles/bouncer/schemas/queue-layer-options",
+    "title": "Queue layer options",
+    "type": "object",
+    "properties": {
+        "capacity": {
+            "type": "number",
+            "minimum": 0
+        },
+        "keepAlive": {
+            "type": "string"
+        }
+    },
+    "additionalProperties": false
+}