refactor: concurrent and chores

Author: dsx137

cmd/main.go

@@ -33,11 +33,6 @@ func loadConfig[T config.IConfig](path string) (T, error) {
 	return *cfg, nil
 }
 
-func start(prog program.IProgram) {
-	program.Program = prog
-	prog.Run()
-}
-
 func init() {
 	rootCmd.PersistentFlags().BoolVarP(&config.IsVerbose, "verbose", "v", false, "verbose output")
 
@@ -58,7 +53,8 @@ func init() {
 				logger.Error("Error creating server: ", err)
 				return
 			}
-			start(s)
+			program.Program = s
+			s.Run()
 		},
 	})
 
@@ -74,7 +70,9 @@ func init() {
 				logger.Error("Error loading config: ", err)
 				return
 			}
-			start(client.NewClient(ctx, cfg))
+			c := client.NewClient(ctx, cfg)
+			program.Program = c
+			c.Run()
 		},
 	})
 

pkg/arch/acceptors/acceptor_udp.go

@@ -2,13 +2,13 @@ package acceptors
 
 import (
 	"context"
-	"net"
 
+	"club.asynclab/asrp/pkg/base/network"
 	"club.asynclab/asrp/pkg/comm"
 )
 
 func NewAcceptorUDP(parentCtx context.Context, addr string) (*Acceptor[comm.UDP], error) {
-	listener, err := net.Listen("udp", addr) // TODO
+	listener, err := network.NewUDPListener(addr)
 	if err != nil {
 		return nil, err
 	}

pkg/arch/acceptors/acceptors.go

@@ -88,23 +88,27 @@ func (acceptor *Acceptor[T]) HandlePacket(pkt packet.IPacket) bool {
 	return true
 }
 
-func (acceptor *Acceptor[T]) init() {
-	go func() {
-		pattern.NewConfigSelectContextAndChannel[*comm.Conn]().
-			WithCtx(acceptor.GetCtx()).
-			WithGoroutine(func(ch chan *comm.Conn) {
-				for {
-					conn, err := acceptor.listener.Accept()
-					if err != nil {
-						if acceptor.GetCtx().Err() != nil {
-							return
-						}
-						continue
+// ----------------------------------------------
+
+func (acceptor *Acceptor[T]) routineRead() {
+	pattern.NewConfigSelectContextAndChannel[*comm.Conn]().
+		WithCtx(acceptor.GetCtx()).
+		WithGoroutine(func(ch chan *comm.Conn) {
+			for {
+				conn, err := acceptor.listener.Accept()
+				if err != nil {
+					if acceptor.GetCtx().Err() != nil {
+						return
 					}
-					ch <- conn
+					continue
 				}
-			}).
-			WithChannelHandler(func(conn *comm.Conn) { go acceptor.handleConnection(conn) }).
-			Run()
-	}()
+				ch <- conn
+			}
+		}).
+		WithChannelHandler(func(conn *comm.Conn) { go acceptor.handleConnection(conn) }).
+		Run()
+}
+
+func (acceptor *Acceptor[T]) init() {
+	go acceptor.routineRead()
 }

pkg/arch/arch.go

@@ -52,8 +52,5 @@ type IAcceptor interface {
 
 type ForwarderWithValues struct {
 	IForwarder
-	Name         string
-	FrontendAddr string
-	Priority     int
-	Weight       int
+	InitPacket *packet.PacketProxyNegotiationRequest
 }

pkg/arch/connectors/connectors.go

@@ -73,6 +73,7 @@ func (connector *Connector[T]) initConnection(conn *comm.Conn) error {
 
 	r := &packet.PacketProxyNegotiationRequest{
 		Name:         connector.proxyConfig.Name,
+		Proto:        connector.proxyConfig.Proto,
 		FrontendAddr: connector.proxyConfig.Frontend,
 		Token:        connector.remoteConfig.Token,
 		Priority:     connector.proxyConfig.Priority,

pkg/arch/dialers/dialer_udp.go

@@ -3,9 +3,7 @@ package dialers
 import (
 	"context"
 	"net"
-	"time"
 
-	"club.asynclab/asrp/pkg/base/network"
 	"club.asynclab/asrp/pkg/comm"
 )
 
@@ -23,5 +21,5 @@ func (impl *DialerImplUDP) Dial(ctx context.Context, addr string) (*comm.Conn, e
 	if err != nil {
 		return nil, err
 	}
-	return comm.NewConnWithParentCtx(ctx, network.NewConnWithTimeout(conn, 60*time.Second, 0)), nil
+	return comm.NewConnWithParentCtx(ctx, conn), nil
 }

pkg/arch/dialers/dialers.go

@@ -92,14 +92,14 @@ func (dialer *Dialer[T]) HandlePacket(pkt packet.IPacket) bool {
 		return false
 	}
 
-	switch pkt := pkt.(type) { // TODO dialer失败就发送给dispatcher重分配请求
+	switch pkt := pkt.(type) {
 	case *packet.PacketProxyData:
 		ok := true
 
 		dialer.conns.Compute(func(v *concurrent.ConcurrentIndexMap[*comm.Conn]) {
 			if _, ok := v.Load(pkt.Uuid); !ok {
 				conn, err := dialer.impl.Dial(dialer.ctx, dialer.addr)
-				if err != nil {
+				if err != nil { // TODO dialer失败就发送给dispatcher重分配请求
 					ok = false
 					logger.Error(fmt.Sprintf("Error dialing: %v", err))
 					return

pkg/arch/dispatchers/dispatchers.go

@@ -17,13 +17,14 @@ var logger = logging.GetLogger()
 
 type Dispatcher struct {
 	concurrent.MetaConcurrentStructure[Dispatcher]
-	ctx          context.Context
-	ctxCancel    context.CancelFunc
-	forwarders   *structure.IndexMap[*arch.ForwarderWithValues]
-	totalWeights map[int]int // priority -> totalWeight
-	currentIndex int
-	senderPacket *channel.SafeSender[packet.IPacket]
-	connsMap     *concurrent.ConcurrentMap[string, string]
+	ctx              context.Context
+	ctxCancel        context.CancelFunc
+	forwarders       *structure.IndexMap[*arch.ForwarderWithValues]
+	totalWeights     map[uint32]uint32 // priority -> totalWeight
+	currentIndex     int
+	senderPacket     *channel.SafeSender[packet.IPacket]
+	connsMap         *concurrent.ConcurrentMap[string, string]                     // conn -> forwarder
+	connsMapBackward *concurrent.ConcurrentMap[string, *structure.HashSet[string]] // forwarder -> conns
 }
 
 func NewDispatcher(parentCtx context.Context) *Dispatcher {
@@ -32,11 +33,12 @@ func NewDispatcher(parentCtx context.Context) *Dispatcher {
 		ctx:                     ctx,
 		ctxCancel:               cancel,
 		forwarders:              structure.NewIndexMap[*arch.ForwarderWithValues](),
-		totalWeights:            make(map[int]int),
+		totalWeights:            make(map[uint32]uint32),
 		currentIndex:            0,
 		MetaConcurrentStructure: *concurrent.NewMetaSyncStructure[Dispatcher](),
 		senderPacket:            channel.NewSafeSenderWithParentCtxAndSize[packet.IPacket](ctx, 16),
 		connsMap:                concurrent.NewSyncMap[string, string](),
+		connsMapBackward:        concurrent.NewSyncMap[string, *structure.HashSet[string]](),
 	}
 
 	return dispatcher
@@ -57,23 +59,12 @@ func (dispatcher *Dispatcher) Close() error {
 
 // ---------------------------------------------------------------------
 
-func (dispatcher *Dispatcher) packetHandlerMiddleware(forwarderUuid string, f func(packet.IPacket) bool) func(packet.IPacket) bool {
-	return func(pkt packet.IPacket) bool {
-		switch pkt := pkt.(type) {
-		case *packet.PacketProxyData:
-			dispatcher.connsMap.LoadOrStore(pkt.Uuid, forwarderUuid)
-		}
-
-		return f(pkt)
-	}
-}
-
 func (dispatcher *Dispatcher) AddForwarder(fwv *arch.ForwarderWithValues) (uuid string) {
 	dispatcher.Lock.Lock()
 	defer dispatcher.Lock.Unlock()
 
 	uuid = dispatcher.forwarders.Store(fwv)
-	dispatcher.totalWeights[fwv.Priority] += fwv.Weight
+	dispatcher.totalWeights[fwv.InitPacket.Priority] += fwv.InitPacket.Weight
 
 	go channel.ConsumeWithCtx(dispatcher.GetCtx(), fwv.GetChanSendPacket(), dispatcher.senderPacket.Push)
 
@@ -88,12 +79,14 @@ func (dispatcher *Dispatcher) RemoveForwarder(uuid string) {
 	if !ok {
 		return
 	}
-	dispatcher.totalWeights[conn.Priority] -= conn.Weight
-	if dispatcher.totalWeights[conn.Priority] == 0 {
-		delete(dispatcher.totalWeights, conn.Priority)
+	dispatcher.totalWeights[conn.InitPacket.Priority] -= conn.InitPacket.Weight
+	if dispatcher.totalWeights[conn.InitPacket.Priority] == 0 {
+		delete(dispatcher.totalWeights, conn.InitPacket.Priority)
 	}
 	dispatcher.forwarders.Delete(uuid)
-	dispatcher.connsMap.Delete(uuid)
+	if conns, ok := dispatcher.connsMapBackward.LoadAndDelete(uuid); ok {
+		conns.Stream().ForEach(func(t string) { dispatcher.connsMap.Delete(t) })
+	}
 }
 
 func (dispatcher *Dispatcher) Next() (uuid string, forwarder arch.IForwarder, ok bool) {
@@ -104,21 +97,21 @@ func (dispatcher *Dispatcher) Next() (uuid string, forwarder arch.IForwarder, ok
 		return
 	}
 
-	totalWeight, _ok := hof.NewStreamWithMap(dispatcher.totalWeights).Max(func(bigger container.Entry[int, int], smaller container.Entry[int, int]) bool {
+	totalWeight, _ok := hof.NewStreamWithMap(dispatcher.totalWeights).Max(func(bigger container.Entry[uint32, uint32], smaller container.Entry[uint32, uint32]) bool {
 		return bigger.GetKey() > smaller.GetKey()
 	})
 
 	if !_ok || totalWeight.GetValue() == 0 {
 		return
 	}
 
-	dispatcher.currentIndex = (dispatcher.currentIndex + 1) % totalWeight.GetValue()
+	dispatcher.currentIndex = (dispatcher.currentIndex + 1) % int(totalWeight.GetValue())
 	i := dispatcher.currentIndex
 
 	dispatcher.forwarders.Stream().Filter(func(t container.Entry[string, *arch.ForwarderWithValues]) bool {
-		return t.GetValue().Priority == totalWeight.GetKey()
+		return t.GetValue().InitPacket.Priority == totalWeight.GetKey()
 	}).Range(func(t container.Entry[string, *arch.ForwarderWithValues]) bool {
-		i -= t.GetValue().Weight
+		i -= int(t.GetValue().InitPacket.Weight)
 		if i < 0 {
 			uuid, forwarder, ok = t.GetKey(), t.GetValue(), true
 			return false
@@ -152,16 +145,18 @@ func (dispatcher *Dispatcher) HandlePacket(pkt packet.IPacket) bool {
 	switch pkt := pkt.(type) {
 	case packet.IPacketForConn:
 		if uuid, ok := dispatcher.connsMap.Load(pkt.GetUuid()); ok {
-			if v, ok := dispatcher.forwarders.Load(uuid); ok {
-				v.HandlePacket(pkt)
-				return true
+			if forwarder, ok := dispatcher.forwarders.Load(uuid); ok {
+				forwarder.HandlePacket(pkt)
+			}
+		} else {
+			if uuid, forwarder, ok := dispatcher.Next(); ok {
+				actual, _ := dispatcher.connsMap.LoadOrStore(pkt.GetUuid(), uuid)
+				actual2, _ := dispatcher.connsMapBackward.LoadOrStore(actual, structure.NewHashSet[string]())
+				actual2.Store(pkt.GetUuid())
+				forwarder.HandlePacket(pkt)
 			}
 		}
 	}
 
-	if uuid, forwarder, ok := dispatcher.Next(); ok {
-		dispatcher.packetHandlerMiddleware(uuid, forwarder.HandlePacket)(pkt)
-	}
-
 	return true
 }

pkg/arch/forwarders/forwarders.go

@@ -22,7 +22,7 @@ func NewForwarder(conn *comm.Conn) *Forwarder {
 	forwarder := &Forwarder{
 		Tip:          "[Common Forwarder]: ",
 		conn:         conn,
-		senderPacket: channel.NewSafeSenderWithSize[packet.IPacket](16),
+		senderPacket: channel.NewSafeSenderWithParentCtxAndSize[packet.IPacket](conn.GetCtx(), 16),
 	}
 
 	forwarder.init()
@@ -69,20 +69,22 @@ func (forwarder *Forwarder) GetChanSendPacket() <-chan packet.IPacket {
 
 // ---------------------------------------------------------------------
 
-func (forwarder *Forwarder) init() {
-	go func() {
-		pattern.NewConfigSelectContextAndChannel[packet.IPacket]().
-			WithCtx(forwarder.GetCtx()).
-			WithGoroutine(func(ch chan packet.IPacket) {
-				for {
-					pkt := forwarder.receivePacket()
-					if pkt == nil {
-						return
-					}
-					ch <- pkt
+func (forwarder *Forwarder) routineRead() {
+	pattern.NewConfigSelectContextAndChannel[packet.IPacket]().
+		WithCtx(forwarder.GetCtx()).
+		WithGoroutine(func(ch chan packet.IPacket) {
+			for {
+				pkt := forwarder.receivePacket()
+				if pkt == nil {
+					return
 				}
-			}).
-			WithChannelHandler(func(pkt packet.IPacket) { forwarder.senderPacket.Push(pkt) }).
-			Run()
-	}()
+				ch <- pkt
+			}
+		}).
+		WithChannelHandler(func(pkt packet.IPacket) { forwarder.senderPacket.Push(pkt) }).
+		Run()
+}
+
+func (forwarder *Forwarder) init() {
+	go forwarder.routineRead()
 }

pkg/arch/ushers/ushers.go

@@ -75,16 +75,15 @@ func (usher *Usher[T]) handleConnection(conn *comm.Conn) {
 	}
 
 	usher.senderForwarder.Push(&arch.ForwarderWithValues{
-		IForwarder:   f,
-		Name:         p.Name,
-		FrontendAddr: p.FrontendAddr,
-		Priority:     int(p.Priority),
-		Weight:       int(p.Weight),
+		IForwarder: f,
+		InitPacket: p,
 	})
 }
 
-func (usher *Usher[T]) init() {
-	go pattern.NewConfigSelectContextAndChannel[*comm.Conn]().
+// -------------------------------------------
+
+func (usher *Usher[T]) routineRead() {
+	pattern.NewConfigSelectContextAndChannel[*comm.Conn]().
 		WithCtx(usher.GetCtx()).
 		WithGoroutine(func(ch chan *comm.Conn) {
 			for {
@@ -102,3 +101,7 @@ func (usher *Usher[T]) init() {
 		WithChannelHandler(func(conn *comm.Conn) { go usher.handleConnection(conn) }).
 		Run()
 }
+
+func (usher *Usher[T]) init() {
+	go usher.routineRead()
+}