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