aiproxy/core/main.go at ef9153152759d8a6f34df37ddf25dc37f8e9d6ee · labring/aiproxy · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
package main

import (
	"context"
	"flag"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/labring/aiproxy/core/common"
	"github.com/labring/aiproxy/core/common/config"
	"github.com/labring/aiproxy/core/common/consume"
	"github.com/labring/aiproxy/core/controller"
	"github.com/labring/aiproxy/core/model"
	"github.com/labring/aiproxy/core/task"
	log "github.com/sirupsen/logrus"
)

var (
	listen    string
	pprofPort int
)

func init() {
	flag.StringVar(&listen, "listen", "0.0.0.0:3000", "http server listen")
	flag.IntVar(&pprofPort, "pprof-port", 15000, "pport http server port")
}

// Swagger godoc
//
//	@title						AI Proxy Swagger API
//	@version					1.0
//	@securityDefinitions.apikey	ApiKeyAuth
//	@in							header
//	@name						Authorization
func main() {
	flag.Parse()

	loadEnv()

	config.ReloadEnv()

	common.InitLog(log.StandardLogger(), config.DebugEnabled)

	printLoadedEnvFiles()

	if err := initializeServices(pprofPort); err != nil {
		log.Fatal("failed to initialize services: " + err.Error())
	}

	defer func() {
		if err := model.CloseDB(); err != nil {
			log.Fatal("failed to close database: " + err.Error())
		}
	}()

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	var wg sync.WaitGroup
	startSyncServices(ctx, &wg)

	srv, _ := setupHTTPServer(listen)

	log.Info("auto test banned models task started")

	go task.AutoTestBannedModelsTask(ctx)

	log.Info("clean log task started")

	go task.CleanLogTask(ctx)

	log.Info("detect ip groups task started")

	go task.DetectIPGroupsTask(ctx)

	log.Info("usage alert task started")

	go task.UsageAlertTask(ctx)

	log.Info("async usage poll task started")

	go task.AsyncUsagePollTask(ctx)

	if common.RedisEnabled {
		log.Info("redis health check task started")

		go task.RedisHealthCheckTask(ctx)
	}

	if task.AdminKeyCacheEnabled() {
		log.Info("admin key cache task started")

		go task.AdminKeyCacheTask(ctx)
	}

	log.Info("update channels balance task started")

	go controller.UpdateChannelsBalance(time.Minute * 10)

	batchProcessorCtx, batchProcessorCancel := context.WithCancel(context.Background())

	wg.Add(1)

	go model.StartBatchProcessorSummary(batchProcessorCtx, &wg)

	log.Infof("server started on http://%s", srv.Addr)
	log.Infof("swagger started on http://%s/swagger/index.html", srv.Addr)

	go listenAndServe(srv)

	<-ctx.Done()

	shutdownSrvCtx, shutdownSrvCancel := context.WithTimeout(context.Background(), 600*time.Second)
	defer shutdownSrvCancel()

	log.Info("shutting down http server...")
	log.Info("max wait time: 600s")

	if err := srv.Shutdown(shutdownSrvCtx); err != nil {
		log.Error("server forced to shutdown: " + err.Error())
	} else {
		log.Info("server shutdown successfully")
	}

	log.Info("shutting down consumer...")
	consume.Wait()

	batchProcessorCancel()

	log.Info("shutting down sync services...")
	wg.Wait()

	log.Info("shutting down batch summary...")
	log.Info("max wait time: 600s")

	cleanCtx, cleanCancel := context.WithTimeout(context.Background(), 600*time.Second)
	defer cleanCancel()

	model.CleanBatchUpdatesSummary(cleanCtx)

	log.Info("server exiting")
}