offload-dispatcher/cmd/dispatcher/main.go

package main

import (
	"github.com/Depado/ginprom"
	"github.com/gin-gonic/gin"
	"github.com/mroth/weightedrand/v2"
	"github.com/sirupsen/logrus"
	ginlogrus "github.com/toorop/gin-logrus"
	"net/http"
	"strconv"
)

var log = logrus.New()

func assignTarget(c *gin.Context) {
	targets := state.GetTargets()

	string_size_hint := c.Query("SIZE_HINT")
	size_hint, err := strconv.ParseInt(string_size_hint, 10, 64)
	if err != nil {
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{
			"reason": "Invalid size hint",
		})
		return
	}

	var choices []weightedrand.Choice[ServerSpec, int]

	for _, target := range targets {
		expectedFreeSpace := target.FreeSpace - size_hint
		if expectedFreeSpace > int64(target.Target.MinimumFreeSpace) && target.Weight > 0 {
			choices = append(choices, weightedrand.NewChoice(target.Target, target.Weight))
		}
	}

	if len(choices) == 0 {
		c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{
			"reason": "No more targets available",
		})
	}

	chooser, _ := weightedrand.NewChooser(choices...)
	result := chooser.Pick()

	c.JSON(http.StatusOK, gin.H{
		"url": result.Url,
	})
}

func setupRouter() *gin.Engine {
	r := gin.New()
	r.Use(ginlogrus.Logger(log), gin.Recovery())
	p := ginprom.New(
		ginprom.Engine(r),
		ginprom.Subsystem("gin"),
		ginprom.Path("/metrics"),
	)
	r.Use(p.Instrument())

	r.GET("/ping", func(c *gin.Context) {
		c.JSON(http.StatusOK, gin.H{
			"message": "pong",
		})
	})

	r.GET("/offload_target", assignTarget)

	return r
}

func main() {
	initState()

	go state.UpdateThread()

	router := setupRouter()

	err := router.Run(state.GetListenAddr())
	log.WithError(err).Fatal("Failed to run router")
}