rts-sim-testing-service/third_party/udp/udp_server.go

package udp

import (
	"context"
	"log/slog"
	"net"
	"runtime/debug"
)

type UdpServer interface {
	Listen() error
	Close()
}

type UdpMsgHandler func(b []byte)

type server struct {
	addr     string
	conn     *net.UDPConn
	handler  UdpMsgHandler
	cancelFn context.CancelFunc
	done     chan struct{} // 服务协程退出信号
}

// NewServer creates a new instance of UdpServer.
func NewServer(addr string, handler UdpMsgHandler) UdpServer {
	return &server{addr: addr, handler: handler, done: make(chan struct{})}
}

func (s *server) Listen() error {
	udpAddr, err := net.ResolveUDPAddr("udp", s.addr)
	if err != nil {
		return err
	}
	conn, err := net.ListenUDP("udp", udpAddr)
	if err != nil {
		return err
	}
	s.conn = conn
	ctx, cfn := context.WithCancel(context.Background())
	// 启动监听处理
	go s.listenAndHandle(ctx)
	s.cancelFn = cfn
	return nil
}

func (s *server) Close() {
	err := s.conn.Close()
	if err != nil {
		slog.Error("udp server close error", "error", err)
	}
	s.cancelFn()
	<-s.done
}

func (s *server) listenAndHandle(ctx context.Context) {
	defer close(s.done)
	defer s.conn.Close()
mainLoop:
	for {
		select {
		case <-ctx.Done():
			// slog.Info("udp server listen 关闭", "addr", s.addr)
			break mainLoop
		default:
		}
		buf := make([]byte, 1024)
		n, _, err := s.conn.ReadFromUDP(buf)
		if err != nil {
			slog.Error("udp server read error", "addr", s.addr, "error", err)
			continue
		}
		// slog.Debug("udp server handle", "addr", ap.String())
		go s.handle(buf[:n])
	}
}

func (s *server) handle(b []byte) {
	defer func() {
		if err := recover(); err != nil {
			slog.Error("udp server handle error", "error", err, "stack", string(debug.Stack()))
			debug.PrintStack()
		}
	}()
	// slog.Info("udp server handle", "msg", string(b))
	s.handler(b)
}