chore: pool and some

Author: dsx137

pkg/arch/dispatchers/dispatchers.go

@@ -20,7 +20,7 @@ type Dispatcher struct {
 	ctx              context.Context
 	ctxCancel        context.CancelFunc
 	forwarders       *structure.IndexMap[*arch.ForwarderWithValues]
-	totalWeights     map[uint32]uint32 // priority -> totalWeight
+	totalWeights     map[uint64]uint64 // priority -> totalWeight
 	currentIndex     int
 	senderPacket     *channel.SafeSender[packet.IPacket]
 	connsMap         *concurrent.ConcurrentMap[string, string]                     // conn -> forwarder
@@ -33,7 +33,7 @@ func NewDispatcher(parentCtx context.Context) *Dispatcher {
 		ctx:                     ctx,
 		ctxCancel:               cancel,
 		forwarders:              structure.NewIndexMap[*arch.ForwarderWithValues](),
-		totalWeights:            make(map[uint32]uint32),
+		totalWeights:            make(map[uint64]uint64),
 		currentIndex:            0,
 		MetaConcurrentStructure: *concurrent.NewMetaSyncStructure[Dispatcher](),
 		senderPacket:            channel.NewSafeSenderWithParentCtxAndSize[packet.IPacket](ctx, 16),
@@ -97,7 +97,7 @@ func (dispatcher *Dispatcher) Next() (uuid string, forwarder arch.IForwarder, ok
 		return
 	}
 
-	totalWeight, _ok := hof.NewStreamWithMap(dispatcher.totalWeights).Max(func(bigger container.Entry[uint32, uint32], smaller container.Entry[uint32, uint32]) bool {
+	totalWeight, _ok := hof.NewStreamWithMap(dispatcher.totalWeights).Max(func(bigger container.Entry[uint64, uint64], smaller container.Entry[uint64, uint64]) bool {
 		return bigger.GetKey() > smaller.GetKey()
 	})
 

pkg/arch/forwarders/forwarders.go

@@ -40,17 +40,6 @@ func (forwarder *Forwarder) HandlePacket(p packet.IPacket) bool {
 	return true
 }
 
-func (forwarder *Forwarder) receivePacket() packet.IPacket {
-	r, err := comm.ReceivePacket(forwarder.conn)
-	if err != nil {
-		if forwarder.GetCtx().Err() != nil {
-			return nil
-		}
-		r = &packet.PacketUnknown{Err: err}
-	}
-	return r
-}
-
 func (forwarder *Forwarder) GetCtx() context.Context {
 	return forwarder.conn.GetCtx()
 }
@@ -69,6 +58,17 @@ func (forwarder *Forwarder) GetChanSendPacket() <-chan packet.IPacket {
 
 // ---------------------------------------------------------------------
 
+func (forwarder *Forwarder) receivePacket() packet.IPacket {
+	r, err := comm.ReceivePacket(forwarder.conn)
+	if err != nil {
+		if forwarder.GetCtx().Err() != nil {
+			return nil
+		}
+		r = &packet.PacketUnknown{Err: err}
+	}
+	return r
+}
+
 func (forwarder *Forwarder) routineRead() {
 	pattern.NewConfigSelectContextAndChannel[packet.IPacket]().
 		WithCtx(forwarder.GetCtx()).

pkg/base/network/udp_listener.go

@@ -16,7 +16,7 @@ import (
 
 var logger = logging.GetLogger()
 
-var udpListenerBufSize uint32 = 10 * 1024
+var udpListenerBufSize uint64 = 10 * 1024
 var udpListenerBufPool = concurrent.NewPool(func() *[]byte {
 	buf := make([]byte, udpListenerBufSize)
 	return &buf
@@ -218,7 +218,7 @@ func (ul *UDPListener) dispatch() {
 			continue
 		}
 
-		data := make([]byte, n)
+		data := make([]byte, n) //TODO 这里可以用pool
 		copy(data, buf[:n])
 		ok := c.receiver.TryPush(data)
 		if !ok {

pkg/comm/comm.go

@@ -10,8 +10,8 @@ import (
 
 // var logger = logging.GetLogger()
 
-var commDataSize uint32 = 128 * 1024
-var commBufSize uint32 = 4 + commDataSize
+var commDataSize uint64 = 64 * 1024
+var commBufSize uint64 = 8 + 1024 + commDataSize
 
 var commBufPool = concurrent.NewPool(func() *[]byte {
 	buf := make([]byte, commBufSize)
@@ -28,37 +28,37 @@ func SendPacket(dst io.Writer, p packet.IPacket) (int, error) {
 		return 0, err
 	}
 
-	bufPtr := commBufPool.Get()
-	defer commBufPool.Put(bufPtr)
-	buf := *bufPtr
-
-	length := uint32(len(data))
+	length := uint64(len(data))
 	if length > commBufSize {
 		return 0, io.ErrShortBuffer
 	}
 
-	buf = buf[:4+length]
-	binary.BigEndian.PutUint32(buf[:4], length)
-	copy(buf[4:], data)
-	return dst.Write(buf)
-}
-
-func ReceivePacket(src io.Reader) (packet.IPacket, error) {
 	bufPtr := commBufPool.Get()
 	defer commBufPool.Put(bufPtr)
 	buf := *bufPtr
 
-	buf = buf[:4]
-	_, err := io.ReadFull(src, buf)
+	buf = buf[:8+length]
+	binary.BigEndian.PutUint64(buf[:8], length)
+	copy(buf[8:], data)
+	return dst.Write(buf)
+}
+
+func ReceivePacket(src io.Reader) (packet.IPacket, error) {
+	lenBuf := make([]byte, 8)
+	_, err := io.ReadFull(src, lenBuf)
 	if err != nil {
 		return nil, err
 	}
 
-	length := binary.BigEndian.Uint32(buf)
+	length := binary.BigEndian.Uint64(lenBuf)
 	if length > commBufSize {
 		return nil, io.ErrShortBuffer
 	}
 
+	bufPtr := commBufPool.Get()
+	defer commBufPool.Put(bufPtr)
+	buf := *bufPtr
+
 	buf = buf[:length]
 	n, err := io.ReadFull(src, buf)
 	if err != nil {

pkg/config/client.go

@@ -12,8 +12,8 @@ type ConfigItemProxy struct {
 	Remotes  []string `toml:"remotes"`
 	Backend  string   `toml:"backend"`
 	Frontend string   `toml:"frontend"`
-	Priority uint32   `toml:"priority"`
-	Weight   uint32   `toml:"weight"`
+	Priority uint64   `toml:"priority"`
+	Weight   uint64   `toml:"weight"`
 }
 
 type ConfigClient struct {

pkg/packet/packets.go

@@ -11,8 +11,8 @@ type PacketProxyNegotiationRequest struct {
 	Name         string
 	Proto        string
 	FrontendAddr string
-	Priority     uint32
-	Weight       uint32
+	Priority     uint64
+	Weight       uint64
 	Token        string
 }