main.go 2.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107
  1. package main
  2. import (
  3. "embed"
  4. "github.com/gin-contrib/sessions"
  5. "github.com/gin-contrib/sessions/cookie"
  6. "github.com/gin-gonic/gin"
  7. "one-api/common"
  8. "one-api/controller"
  9. "one-api/middleware"
  10. "one-api/model"
  11. "one-api/router"
  12. "os"
  13. "strconv"
  14. )
  15. //go:embed web/build
  16. var buildFS embed.FS
  17. //go:embed web/build/index.html
  18. var indexPage []byte
  19. func main() {
  20. common.SetupGinLog()
  21. common.SysLog("One API " + common.Version + " started")
  22. if os.Getenv("GIN_MODE") != "debug" {
  23. gin.SetMode(gin.ReleaseMode)
  24. }
  25. if common.DebugEnabled {
  26. common.SysLog("running in debug mode")
  27. }
  28. // Initialize SQL Database
  29. err := model.InitDB()
  30. if err != nil {
  31. common.FatalLog("failed to initialize database: " + err.Error())
  32. }
  33. defer func() {
  34. err := model.CloseDB()
  35. if err != nil {
  36. common.FatalLog("failed to close database: " + err.Error())
  37. }
  38. }()
  39. // Initialize Redis
  40. err = common.InitRedisClient()
  41. if err != nil {
  42. common.FatalLog("failed to initialize Redis: " + err.Error())
  43. }
  44. // Initialize options
  45. model.InitOptionMap()
  46. if common.RedisEnabled {
  47. model.InitChannelCache()
  48. }
  49. if os.Getenv("SYNC_FREQUENCY") != "" {
  50. frequency, err := strconv.Atoi(os.Getenv("SYNC_FREQUENCY"))
  51. if err != nil {
  52. common.FatalLog("failed to parse SYNC_FREQUENCY: " + err.Error())
  53. }
  54. common.SyncFrequency = frequency
  55. go model.SyncOptions(frequency)
  56. if common.RedisEnabled {
  57. go model.SyncChannelCache(frequency)
  58. }
  59. }
  60. if os.Getenv("CHANNEL_UPDATE_FREQUENCY") != "" {
  61. frequency, err := strconv.Atoi(os.Getenv("CHANNEL_UPDATE_FREQUENCY"))
  62. if err != nil {
  63. common.FatalLog("failed to parse CHANNEL_UPDATE_FREQUENCY: " + err.Error())
  64. }
  65. go controller.AutomaticallyUpdateChannels(frequency)
  66. }
  67. if os.Getenv("CHANNEL_TEST_FREQUENCY") != "" {
  68. frequency, err := strconv.Atoi(os.Getenv("CHANNEL_TEST_FREQUENCY"))
  69. if err != nil {
  70. common.FatalLog("failed to parse CHANNEL_TEST_FREQUENCY: " + err.Error())
  71. }
  72. go controller.AutomaticallyTestChannels(frequency)
  73. }
  74. if os.Getenv("BATCH_UPDATE_ENABLED") == "true" {
  75. common.BatchUpdateEnabled = true
  76. common.SysLog("batch update enabled with interval " + strconv.Itoa(common.BatchUpdateInterval) + "s")
  77. model.InitBatchUpdater()
  78. }
  79. controller.InitTokenEncoders()
  80. // Initialize HTTP server
  81. server := gin.New()
  82. server.Use(gin.Recovery())
  83. // This will cause SSE not to work!!!
  84. //server.Use(gzip.Gzip(gzip.DefaultCompression))
  85. server.Use(middleware.RequestId())
  86. middleware.SetUpLogger(server)
  87. // Initialize session store
  88. store := cookie.NewStore([]byte(common.SessionSecret))
  89. server.Use(sessions.Sessions("session", store))
  90. router.SetRouter(server, buildFS, indexPage)
  91. var port = os.Getenv("PORT")
  92. if port == "" {
  93. port = strconv.Itoa(*common.Port)
  94. }
  95. err = server.Run(":" + port)
  96. if err != nil {
  97. common.FatalLog("failed to start HTTP server: " + err.Error())
  98. }
  99. }