main.go 2.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107
  1. package main
  2. import (
  3. "embed"
  4. "fmt"
  5. "github.com/gin-contrib/sessions"
  6. "github.com/gin-contrib/sessions/cookie"
  7. "github.com/gin-gonic/gin"
  8. "one-api/common"
  9. "one-api/controller"
  10. "one-api/middleware"
  11. "one-api/model"
  12. "one-api/router"
  13. "os"
  14. "strconv"
  15. )
  16. //go:embed web/build
  17. var buildFS embed.FS
  18. //go:embed web/build/index.html
  19. var indexPage []byte
  20. func main() {
  21. common.SetupLogger()
  22. common.SysLog("One API " + common.Version + " started")
  23. if os.Getenv("GIN_MODE") != "debug" {
  24. gin.SetMode(gin.ReleaseMode)
  25. }
  26. if common.DebugEnabled {
  27. common.SysLog("running in debug mode")
  28. }
  29. // Initialize SQL Database
  30. err := model.InitDB()
  31. if err != nil {
  32. common.FatalLog("failed to initialize database: " + err.Error())
  33. }
  34. defer func() {
  35. err := model.CloseDB()
  36. if err != nil {
  37. common.FatalLog("failed to close database: " + err.Error())
  38. }
  39. }()
  40. // Initialize Redis
  41. err = common.InitRedisClient()
  42. if err != nil {
  43. common.FatalLog("failed to initialize Redis: " + err.Error())
  44. }
  45. // Initialize options
  46. model.InitOptionMap()
  47. if common.RedisEnabled {
  48. // for compatibility with old versions
  49. common.MemoryCacheEnabled = true
  50. }
  51. if common.MemoryCacheEnabled {
  52. common.SysLog("memory cache enabled")
  53. common.SysError(fmt.Sprintf("sync frequency: %d seconds", common.SyncFrequency))
  54. model.InitChannelCache()
  55. }
  56. if common.MemoryCacheEnabled {
  57. go model.SyncOptions(common.SyncFrequency)
  58. go model.SyncChannelCache(common.SyncFrequency)
  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. }