134 lines
3.6 KiB
Go
134 lines
3.6 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"log"
|
|
"net/http"
|
|
"os"
|
|
"os/signal"
|
|
"syscall"
|
|
"time"
|
|
|
|
"git.misaka.ren/M1saka/token_thief/config"
|
|
"git.misaka.ren/M1saka/token_thief/db"
|
|
"git.misaka.ren/M1saka/token_thief/logger"
|
|
"git.misaka.ren/M1saka/token_thief/proxy"
|
|
)
|
|
|
|
func main() {
|
|
log.SetFlags(log.LstdFlags | log.Lmicroseconds)
|
|
if err := run(); err != nil {
|
|
log.Printf("[main] fatal: %v", err)
|
|
os.Exit(1)
|
|
}
|
|
}
|
|
|
|
type shutdownStep struct {
|
|
name string
|
|
run func(context.Context) error
|
|
}
|
|
|
|
func waitForShutdown(stop <-chan struct{}, serverErr <-chan error) error {
|
|
select {
|
|
case <-stop:
|
|
log.Printf("[main] shutdown signal received")
|
|
return nil
|
|
case err := <-serverErr:
|
|
if errors.Is(err, http.ErrServerClosed) {
|
|
return nil
|
|
}
|
|
return fmt.Errorf("server: %w", err)
|
|
}
|
|
}
|
|
|
|
func shutdownAll(ctx context.Context, cancelRoot context.CancelFunc, steps ...shutdownStep) error {
|
|
defer cancelRoot()
|
|
var shutdownErr error
|
|
for _, step := range steps {
|
|
done := make(chan error, 1)
|
|
started := make(chan struct{})
|
|
go func() {
|
|
close(started)
|
|
done <- step.run(ctx)
|
|
}()
|
|
<-started
|
|
var err error
|
|
select {
|
|
case err = <-done:
|
|
case <-ctx.Done():
|
|
err = ctx.Err()
|
|
}
|
|
if err != nil {
|
|
log.Printf("[main] %s shutdown: %v", step.name, err)
|
|
shutdownErr = errors.Join(shutdownErr, fmt.Errorf("%s shutdown: %w", step.name, err))
|
|
}
|
|
}
|
|
return shutdownErr
|
|
}
|
|
|
|
func run() error {
|
|
cfg, err := config.Load()
|
|
if err != nil {
|
|
return fmt.Errorf("config: %w", err)
|
|
}
|
|
if _, err := db.ClickHouseOptions(cfg.ClickHouseURL); err != nil {
|
|
return err
|
|
}
|
|
|
|
filter, err := config.LoadFilter(cfg.FilterFile)
|
|
if err != nil {
|
|
return fmt.Errorf("load filter: %w", err)
|
|
}
|
|
log.Printf("[main] filter mode=%s patterns=%d", filter.Mode, len(filter.Patterns))
|
|
|
|
rootCtx, stopSignals := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
|
|
defer stopSignals()
|
|
rootCtx, cancel := context.WithCancel(rootCtx)
|
|
defer cancel()
|
|
|
|
pool := db.NewPool(rootCtx, cfg.ClickHouseURL, cfg.DBReconnectInterval)
|
|
|
|
queue := logger.NewQueue(pool, cfg.LogQueueSize, cfg.LogBatchSize, cfg.LogWorkers, cfg.LogBatchInterval, cfg.LogQueueBytes)
|
|
queue.Start(rootCtx)
|
|
|
|
h := proxy.NewWithOptions(cfg.UpstreamURL, filter, queue, cfg.MaxBodyBytes, proxy.Options{
|
|
UpstreamTimeout: cfg.UpstreamTimeout,
|
|
ResponseTimeout: cfg.UpstreamResponseTimeout,
|
|
SSEIdleTimeout: cfg.UpstreamStreamIdleTimeout,
|
|
UpstreamTLSInsecureSkipVerify: cfg.UpstreamTLSInsecureSkipVerify,
|
|
TrustedProxies: cfg.TrustedProxies,
|
|
MaxRequestBytes: cfg.MaxRequestBytes,
|
|
})
|
|
|
|
srv := &http.Server{
|
|
Addr: cfg.ListenAddr,
|
|
Handler: h,
|
|
ReadHeaderTimeout: 30 * time.Second,
|
|
ReadTimeout: cfg.ReadTimeout,
|
|
WriteTimeout: cfg.WriteTimeout,
|
|
IdleTimeout: cfg.IdleTimeout,
|
|
}
|
|
|
|
serverErr := make(chan error, 1)
|
|
go func() {
|
|
log.Printf("[main] listening on %s, upstream=%s", cfg.ListenAddr, cfg.UpstreamURL)
|
|
serverErr <- srv.ListenAndServe()
|
|
}()
|
|
|
|
runErr := waitForShutdown(rootCtx.Done(), serverErr)
|
|
|
|
// HTTP、升级连接、日志排空和数据库共享同一个关闭总预算。
|
|
shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second)
|
|
defer shutdownCancel()
|
|
shutdownErr := shutdownAll(shutdownCtx, cancel,
|
|
shutdownStep{name: "http", run: srv.Shutdown},
|
|
shutdownStep{name: "upgraded connection", run: h.Shutdown},
|
|
shutdownStep{name: "queue", run: func(ctx context.Context) error { return queue.Stop(ctx) }},
|
|
shutdownStep{name: "db", run: pool.Close},
|
|
)
|
|
log.Printf("[main] bye")
|
|
return errors.Join(runErr, shutdownErr)
|
|
}
|