|
- 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")
- }
|