proxy/server/fwd/fwd.go

package fwd

import (
	"context"
	"log/slog"
	"proxy-server/pkg/utils"
	"proxy-server/server/core"
	"sync"
)

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

	userConnMap core.SyncMap[string, *core.Conn]

	fwdLesWg   utils.CountWaitGroup
	ctrlConnWg utils.CountWaitGroup
	dataConnWg utils.CountWaitGroup
	userConnWg utils.CountWaitGroup
}

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

func (s *Service) Run() error {
	slog.Debug("启动转发服务")

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

	wg := sync.WaitGroup{}

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

	// 数据通道
	wg.Add(1)
	go func() {
		defer wg.Done()
		err := s.listenData()
		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()
	s.fwdLesWg.Wait()
	s.ctrlConnWg.Wait()
	s.userConnWg.Wait()
	s.dataConnWg.Wait()

	return nil
}

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