This commit is contained in:
RMT
2026-03-15 18:37:50 +08:00
committed by ginuerzh
parent 802943bd28
commit 2a2643a6b7
2 changed files with 169 additions and 0 deletions
+2
View File
@@ -47,6 +47,8 @@ func ParseNodeSelector(cfg *config.SelectorConfig) selector.Selector[*chain.Node
strategy = xs.FIFOStrategy[*chain.Node]()
case "hash":
strategy = xs.HashStrategy[*chain.Node]()
case "parallel":
strategy = xs.ParallelStrategy[*chain.Node]()
default:
strategy = xs.RoundRobinStrategy[*chain.Node]()
}
+167
View File
@@ -0,0 +1,167 @@
package selector
import (
"context"
"errors"
"net"
"sync"
"github.com/go-gost/core/chain"
"github.com/go-gost/core/connector"
"github.com/go-gost/core/logger"
"github.com/go-gost/core/selector"
xnet "github.com/go-gost/x/internal/net"
)
type parallelStrategy[T any] struct{}
func ParallelStrategy[T any]() selector.Strategy[T] {
return &parallelStrategy[T]{}
}
func (s *parallelStrategy[T]) Apply(ctx context.Context, vs ...T) (v T) {
if len(vs) == 0 {
return
}
if len(vs) == 1 {
return vs[0]
}
nodes := make([]*chain.Node, 0, len(vs))
for _, node := range vs {
nodes = append(nodes, any(node).(*chain.Node))
}
vn := chain.NewNode("parallel", "", chain.TransportNodeOption(&parallelTransporter{nodes: nodes}))
return any(vn).(T)
}
type parallelConn struct {
net.Conn
node *chain.Node
}
type parallelTransporter struct {
nodes []*chain.Node
}
func (tr *parallelTransporter) Dial(ctx context.Context, addr string) (net.Conn, error) {
if len(tr.nodes) == 0 {
return nil, errors.New("no nodes available")
}
ctx, cancel := context.WithCancel(ctx)
type result struct {
conn net.Conn
node *chain.Node
err error
}
resCh := make(chan result, len(tr.nodes))
var wg sync.WaitGroup
for _, node := range tr.nodes {
wg.Add(1)
go func(n *chain.Node) {
defer wg.Done()
raddr, err := xnet.Resolve(ctx, "ip", n.Addr, n.Options().Resolver, n.Options().HostMapper, logger.Default())
if err != nil {
resCh <- result{err: err, node: n}
return
}
cc, err := n.Options().Transport.Dial(ctx, raddr)
resCh <- result{conn: cc, node: n, err: err}
}(node)
}
var best result
var errs []error
for i := 0; i < len(tr.nodes); i++ {
res := <-resCh
if res.err != nil {
errs = append(errs, res.err)
continue
}
best = res
break
}
cancel()
if best.conn == nil {
go func() {
wg.Wait()
close(resCh)
}()
return nil, errors.Join(errs...)
}
go func() {
// Drain the rest of the channel to clean up late successful connections
for i := len(errs) + 1; i < len(tr.nodes); i++ {
res := <-resCh
if res.err == nil && res.conn != nil {
res.conn.Close()
}
}
wg.Wait()
close(resCh)
}()
return &parallelConn{
Conn: best.conn,
node: best.node,
}, nil
}
func (tr *parallelTransporter) Handshake(ctx context.Context, conn net.Conn) (net.Conn, error) {
pc, ok := conn.(*parallelConn)
if !ok {
return conn, nil
}
cc, err := pc.node.Options().Transport.Handshake(ctx, pc.Conn)
if err != nil {
return nil, err
}
pc.Conn = cc
return pc, nil
}
func (tr *parallelTransporter) Connect(ctx context.Context, conn net.Conn, network, address string) (net.Conn, error) {
pc, ok := conn.(*parallelConn)
if !ok {
return nil, errors.New("invalid connection type")
}
return pc.node.Options().Transport.Connect(ctx, pc.Conn, network, address)
}
func (tr *parallelTransporter) Bind(ctx context.Context, conn net.Conn, network, address string, opts ...connector.BindOption) (net.Listener, error) {
pc, ok := conn.(*parallelConn)
if !ok {
return nil, errors.New("invalid connection type")
}
return pc.node.Options().Transport.Bind(ctx, pc.Conn, network, address, opts...)
}
func (tr *parallelTransporter) Multiplex() bool {
return false
}
func (tr *parallelTransporter) Options() *chain.TransportOptions {
return &chain.TransportOptions{}
}
func (tr *parallelTransporter) Copy() chain.Transporter {
return &parallelTransporter{
nodes: append([]*chain.Node(nil), tr.nodes...),
}
}