base-golang/pkg/proxy/service.go

package proxy

import (
	"fmt"
	"go.uber.org/zap"
	"net"
	"runtime/debug"
	"sync"
)

var (
	servicesMap = new(sync.Map)
	logger      *zap.Logger
)

type Service struct {
	TCPConn TCP
	Name    string
}

func (s *Service) Stop() {
	servicesMap.Delete(s.Name)
	s.TCPConn.Close()
}

func Run(name string, args TCPArgs, zapLogger *zap.Logger) *Service {
	logger = zapLogger
	service := &Service{
		TCPConn: &tcp{cfg: args},
		Name:    name,
	}
	store, loaded := servicesMap.LoadOrStore(name, service)
	if loaded {
		service = store.(*Service)
	}

	go func() {
		defer func() {
			recoverErr := recover()
			if recoverErr != nil {
				logger.Sugar().Errorf("%s servcie crashed, ERR: %s\ntrace:%s", name, recoverErr, string(debug.Stack()))
			}
		}()
		startErr := service.TCPConn.Start()
		if startErr != nil {
			logger.Sugar().Errorf("%s servcie fail, ERR: %s", name, startErr)
		}
	}()

	return service
}

///////////////////////////////////////////////////////////////////////////

type Listener struct {
	ip               string
	port             int
	Listener         net.Listener
	errAcceptHandler func(err error)
}

func NewListener(ip string, port int) Listener {
	return Listener{
		ip:   ip,
		port: port,
		errAcceptHandler: func(err error) {
			logger.Sugar().Errorf("accept error , ERR:%s", err)
		},
	}
}

func (sc *Listener) ListenTCP(fn func(conn net.Conn)) (err error) {
	sc.Listener, err = net.Listen("tcp", fmt.Sprintf("%s:%d", sc.ip, sc.port))
	if err == nil {
		go func() {
			defer func() {
				if e := recover(); e != nil {
					logger.Sugar().Infof("ListenTCP crashed , err : %s , \ntrace:%s", e, string(debug.Stack()))
				}
			}()
			for {
				var conn net.Conn
				conn, err = sc.Listener.Accept()
				if err == nil {
					go func() {
						defer func() {
							if e := recover(); e != nil {
								logger.Sugar().Infof("connection handler crashed , err : %s , \ntrace:%s", e, string(debug.Stack()))
							}
						}()
						fn(conn)
					}()
				} else {
					sc.errAcceptHandler(err)
					break
				}
			}
		}()
	}
	return
}

func (sc *Listener) CloseListen() error {
	return sc.Listener.Close()
}