proxy/server/server.go

package server

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"log/slog"
	"net/http"
	"os"
	"os/signal"
	"proxy-server/pkg/utils"
	"proxy-server/server/debug"
	"proxy-server/server/fwd"
	"proxy-server/server/pkg/env"
	"proxy-server/server/pkg/log"
	"proxy-server/server/pkg/orm"
	"proxy-server/server/web"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/google/uuid"

	"github.com/joho/godotenv"

	_ "net/http/pprof"
)

const (
	Version      = 1
	RestoreMagic = 0x72
)

type server struct {
	name string
}

func New() *server {
	return &server{}
}

func (s *server) Run() (err error) {

	// 初始化
	err = s.init()
	if err != nil {
		return err
	}

	// 恢复服务状态
	err = s.restore()
	if err != nil {
		return err
	}

	// 准备子服务
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	wg := sync.WaitGroup{}

	// 转发服务
	wg.Add(1)
	fwdQuit := make(chan struct{}, 1)
	go func() {
		defer wg.Done()
		defer close(fwdQuit)
		err := startFwdServer(ctx)
		if err != nil {
			slog.Error("转发服务发生错误", "err", err)
		}
		fwdQuit <- struct{}{}
	}()

	// 接口服务
	wg.Add(1)
	apiQuit := make(chan struct{}, 1)
	go func() {
		defer wg.Done()
		err := startWebServer(ctx)
		if err != nil {
			slog.Error("接口服务发生错误", "err", err)
		}
		apiQuit <- struct{}{}
	}()

	// debug
	go func() {
		debug.Start(ctx)
	}()

	// 性能监控
	// go func() {
	// 	runtime.SetBlockProfileRate(1)
	// 	err := http.ListenAndServe(":6060", nil)
	// 	if err != nil {
	// 		slog.Error("性能监控服务发生错误", "err", err)
	// 	}
	// }()

	// 报告上线
	slog.Debug("报告服务上线")
	go reportOnline(ctx, s.name)

	// 等待退出信号
	osQuit := make(chan os.Signal, 1)
	signal.Notify(osQuit, os.Interrupt, syscall.SIGTERM)

	select {
	case <-osQuit:
		slog.Info("服务主动退出")
	case <-fwdQuit:
		slog.Warn("fwd 服务异常退出")
	case <-apiQuit:
		slog.Warn("web 服务异常退出")
	}

	// 报告下线
	slog.Debug("报告服务下线")
	go reportOffline(ctx, s.name)

	// 退出其他服务
	cancel()

	timeout, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	select {
	case <-utils.ChanWgWait(timeout, &wg):
		slog.Info("服务已退出")
	case <-timeout.Done():
		slog.Warn("退出超时，强制退出")
	}

	return nil
}

func (s *server) restore() (err error) {
	var file = "proxy.lock"

	bytes, err := os.ReadFile(file)
	if err != nil {
		return err
	}

	if len(bytes) == 17 && bytes[0] == RestoreMagic {
		s.name = uuid.UUID(bytes[1:]).String()
		slog.Info("恢复服务名称", "name", s.name)
	} else {
		var u = uuid.New()
		s.name = u.String()

		bytes = make([]byte, 17)
		bytes[0] = RestoreMagic
		copy(bytes[1:], u[:])
		err := os.WriteFile(file, bytes, 0644)
		if err != nil {
			return err
		}

		slog.Info("生成服务名称", "name", s.name)
	}

	return nil
}

func (s *server) init() error {

	err := godotenv.Load()
	if err != nil {
		println("没有本地环境变量文件")
	}

	log.Init()
	env.Init()
	orm.Init()

	return nil
}

func startFwdServer(ctx context.Context) error {
	server := fwd.New(nil)
	go func() {
		<-ctx.Done()
		server.Close()
	}()
	server.Run()
	return nil
}

func startWebServer(ctx context.Context) error {
	return web.Start(ctx)
}

func reportOnline(ctx context.Context, name string) {
	reportRepeat(ctx, env.EndpointOnline, map[string]any{
		"name":    name,
		"version": Version,
	})
}

func reportOffline(ctx context.Context, name string) {
	reportRepeat(ctx, env.EndpointOffline, map[string]any{
		"name":    name,
		"version": Version,
	})
}

func reportRepeat(ctx context.Context, endpoint string, body any) {
	bodyStr, err := json.Marshal(body)
	if err != nil {
		panic(err)
	}

	for {
		req, err := http.NewRequest("POST", endpoint, strings.NewReader(string(bodyStr)))
		if err != nil {
			panic(err)
		}

		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Authorization", "Basic "+base64.RawURLEncoding.EncodeToString([]byte("proxy:proxy")))

		resp, err := http.DefaultClient.Do(req)
		if resp != nil && resp.StatusCode == http.StatusOK {
			return
		}
		select {
		case <-ctx.Done():
			return
		default:
		}

		slog.Warn("服务注册失败，五秒后重试", "err", err)
		time.Sleep(5 * time.Second)
	}
}