proxy/gateway/fwd/fwd.go

package fwd

import (
	"context"
	"log/slog"
	"proxy-server/gateway/app"
	"proxy-server/gateway/env"
	"sync"
)

type Service struct {
	ctx    context.Context
	cancel context.CancelFunc
}

func New() *Service {
	ctx, cancel := context.WithCancel(context.Background())
	return &Service{
		ctx:    ctx,
		cancel: cancel,
	}
}

func (s *Service) Run() error {
	slog.Info("启动转发服务", "控制通道", env.AppCtrlPort, "数据通道", env.AppDataPort)

	errQuit := make(chan struct{}, 2)
	defer close(errQuit)

	wg := sync.WaitGroup{}

	// 控制通道
	wg.Add(1)
	go func() {
		defer wg.Done()
		err := ListenCtrl(s.ctx)
		if err != nil {
			slog.Error("fwd 控制通道监听发生错误", "err", err)
			errQuit <- struct{}{}
			return
		}
	}()

	// 数据通道
	wg.Add(1)
	go func() {
		defer wg.Done()
		err := ListenData(s.ctx)
		if err != nil {
			slog.Error("fwd 数据通道监听发生错误", "err", err)
			errQuit <- struct{}{}
			return
		}
	}()

	// 等待退出
	select {
	case <-s.ctx.Done():
	case <-errQuit:
		slog.Warn("fwd 服务异常退出")
		s.Stop()
	}

	wg.Wait()
	app.FwdLesWg.Wait()
	app.CtrlConnWg.Wait()
	app.UserConnWg.Wait()
	app.DataConnWg.Wait()

	return nil
}

func (s *Service) Stop() {
	s.cancel()
}