x: implement node-level liveness probe

- x/chain/node.go: StartNodeProbe goroutine — periodic TCP/HTTP probing
- x/internal/probe/probe.go: HTTPProber with status code validation
- x/selector/latency.go: LatencyFilter, LowestLatencyStrategy (generic)
- x/config: ProbeConfig, MaxLatency, parseProbeConfig, lowestlatency strategy
- x/hop/hop.go: Close() cascade to node.Close()
- x/chain/chain.go: Chain.Close() cascade to hop.Close()

Tests: 4 latency selector tests, 4 node probe tests, 4 HTTP prober tests,
1 hop failover integration test.
This commit is contained in:
ginuerzh
2026-07-16 20:27:21 +08:00
parent c91a414d68
commit 46fc054ad5
12 changed files with 1024 additions and 4 deletions
+12
View File
@@ -30,6 +30,7 @@ package chain
import (
"context"
"io"
"github.com/go-gost/core/chain"
"github.com/go-gost/core/hop"
@@ -111,6 +112,17 @@ func (c *Chain) Name() string {
return c.name
}
// Close stops all probe goroutines by cascading to each hop's Close,
// which in turn cascades to node.Close() for every node.
func (c *Chain) Close() error {
for _, h := range c.hops {
if closer, ok := h.(io.Closer); ok {
closer.Close()
}
}
return nil
}
// Route builds a route by selecting one node from each hop. If a node
// supports multiplexing, the route is split — nodes before it form a
// sub-route that is copied into the transport for reuse.
+98
View File
@@ -0,0 +1,98 @@
package chain
import (
"context"
"time"
"github.com/go-gost/core/chain"
"github.com/go-gost/core/logger"
"github.com/go-gost/x/internal/probe"
)
// StartNodeProbe starts a background goroutine that periodically probes the
// node's transport. On success the node's marker is reset; on failure the
// marker is incremented. cfg may be nil (no-op).
func StartNodeProbe(node *chain.Node, cfg *chain.ProbeConfig, log logger.Logger) {
if cfg == nil || cfg.Addr == "" {
return
}
ctx, cancel := context.WithCancel(context.Background())
node.SetProbeCancel(cancel)
go runNodeProbe(ctx, node, cfg, log)
}
func runNodeProbe(ctx context.Context, node *chain.Node, cfg *chain.ProbeConfig, log logger.Logger) {
interval := cfg.Interval
if interval <= 0 {
interval = 30 * time.Second
}
ticker := time.NewTicker(interval)
defer ticker.Stop()
probeNode(node, cfg, log) // first probe immediately
for {
select {
case <-ticker.C:
probeNode(node, cfg, log)
case <-ctx.Done():
return
}
}
}
func probeNode(node *chain.Node, cfg *chain.ProbeConfig, log logger.Logger) {
timeout := cfg.Timeout
if timeout <= 0 {
timeout = 10 * time.Second
}
ctx, cancel := context.WithTimeout(context.Background(), timeout)
defer cancel()
start := time.Now()
tr := node.Options().Transport
conn, err := tr.Dial(ctx, node.Addr)
if err != nil {
recordProbe(node, false, time.Since(start), err.Error(), log)
return
}
if hc, err2 := tr.Handshake(ctx, conn); err2 == nil {
conn = hc
} else {
conn.Close()
recordProbe(node, false, time.Since(start), err2.Error(), log)
return
}
latency := time.Since(start)
if cfg.Type == chain.ProbeTypeHTTP {
err = probe.NewHTTPProber(cfg).Probe(conn)
conn.Close()
if err != nil {
recordProbe(node, false, latency, err.Error(), log)
return
}
} else {
conn.Close()
}
recordProbe(node, true, latency, "", log)
}
func recordProbe(node *chain.Node, success bool, latency time.Duration, errStr string, log logger.Logger) {
result := &chain.ProbeResult{
Success: success,
Latency: latency,
Error: errStr,
Timestamp: time.Now(),
}
node.SetProbeResult(result)
if success {
node.Marker().Reset()
} else {
node.Marker().Mark()
}
log.Debugf("node probe %s: success=%v latency=%v err=%s", node.Name, success, latency, errStr)
}
+235
View File
@@ -0,0 +1,235 @@
package chain
import (
"context"
"errors"
"net"
"sync/atomic"
"testing"
"time"
"github.com/go-gost/core/chain"
"github.com/go-gost/core/connector"
xlogger "github.com/go-gost/x/logger"
)
// testTransport implements chain.Transporter with real TCP dials.
type testTransport struct {
dialer *net.Dialer
options chain.TransportOptions
}
func (t *testTransport) Dial(ctx context.Context, addr string) (net.Conn, error) {
return t.dialer.DialContext(ctx, "tcp", addr)
}
func (t *testTransport) Handshake(ctx context.Context, conn net.Conn) (net.Conn, error) {
return conn, nil
}
func (t *testTransport) Connect(ctx context.Context, conn net.Conn, network, address string) (net.Conn, error) {
return nil, errors.New("not implemented")
}
func (t *testTransport) Bind(ctx context.Context, conn net.Conn, network, address string, opts ...connector.BindOption) (net.Listener, error) {
return nil, errors.New("not implemented")
}
func (t *testTransport) Multiplex() bool { return false }
func (t *testTransport) Options() *chain.TransportOptions { return &t.options }
func (t *testTransport) Copy() chain.Transporter { return t }
func TestNodeProbeTCP(t *testing.T) {
l, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatal(err)
}
defer l.Close()
// accept one connection, then close immediately (TCP probe only needs dial+handshake)
go func() {
for {
conn, _ := l.Accept()
if conn != nil {
conn.Close()
}
}
}()
addr := l.Addr().String()
tr := &testTransport{dialer: &net.Dialer{Timeout: 2 * time.Second}}
node := chain.NewNode("test", addr, chain.TransportNodeOption(tr))
cfg := &chain.ProbeConfig{
Type: chain.ProbeTypeTCP,
Addr: addr,
Interval: 100 * time.Millisecond,
Timeout: 2 * time.Second,
}
StartNodeProbe(node, cfg, xlogger.Nop())
// Wait for first probe
time.Sleep(300 * time.Millisecond)
pr := node.ProbeResult()
if pr == nil || !pr.Success {
t.Fatalf("expected successful probe, got %+v", pr)
}
// Close listener → next probe fails
l.Close()
time.Sleep(300 * time.Millisecond)
pr = node.ProbeResult()
if pr == nil || pr.Success {
t.Fatalf("expected failed probe after listener close, got %+v", pr)
}
if node.Marker().Count() == 0 {
t.Error("expected marker count > 0 after failure")
}
node.Close()
}
func TestNodeProbeClose(t *testing.T) {
l, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatal(err)
}
defer l.Close()
go func() {
for {
conn, _ := l.Accept()
if conn != nil {
conn.Close()
}
}
}()
addr := l.Addr().String()
tr := &testTransport{dialer: &net.Dialer{Timeout: 2 * time.Second}}
node := chain.NewNode("test", addr, chain.TransportNodeOption(tr))
cfg := &chain.ProbeConfig{
Type: chain.ProbeTypeTCP,
Addr: addr,
Interval: 100 * time.Millisecond,
Timeout: 2 * time.Second,
}
StartNodeProbe(node, cfg, xlogger.Nop())
time.Sleep(150 * time.Millisecond)
// Snapshot result before close
lastSuccess := node.ProbeResult().Success
if !lastSuccess {
t.Fatal("expected initial success before close")
}
node.Close()
// Capture marker count right after close
countAfterClose := node.Marker().Count()
// Wait longer than probe interval — no more updates after close
time.Sleep(300 * time.Millisecond)
// Marker count should not increase after close (no more probes running)
if node.Marker().Count() != countAfterClose {
t.Errorf("marker changed after close: was %d, now %d", countAfterClose, node.Marker().Count())
}
}
func TestNodeProbeDisabled(t *testing.T) {
node := chain.NewNode("no-probe", "addr")
// nil config → no-op
StartNodeProbe(node, nil, xlogger.Nop())
if pr := node.ProbeResult(); pr != nil {
t.Error("expected nil probe result for disabled probe")
}
// empty addr → no-op
StartNodeProbe(node, &chain.ProbeConfig{}, xlogger.Nop())
if pr := node.ProbeResult(); pr != nil {
t.Error("expected nil probe result for empty addr")
}
}
func TestNodeProbeHTTP(t *testing.T) {
// Start a small HTTP server
mux := &httpTestMux{}
l, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatal(err)
}
defer l.Close()
go func() {
for {
conn, err := l.Accept()
if err != nil {
return
}
go mux.serve(conn)
}
}()
addr := l.Addr().String()
tr := &testTransport{dialer: &net.Dialer{Timeout: 2 * time.Second}}
node := chain.NewNode("test-http", addr, chain.TransportNodeOption(tr))
cfg := &chain.ProbeConfig{
Type: chain.ProbeTypeHTTP,
Addr: addr,
Interval: 100 * time.Millisecond,
Timeout: 2 * time.Second,
HTTPPath: "/health",
ExpectedStatus: 200,
}
// Case 1: server returns 200
mux.status.Store(200)
StartNodeProbe(node, cfg, xlogger.Nop())
time.Sleep(300 * time.Millisecond)
pr := node.ProbeResult()
if pr == nil || !pr.Success {
t.Fatalf("expected HTTP 200 probe success, got %+v", pr)
}
// Case 2: server returns 500
mux.status.Store(500)
time.Sleep(300 * time.Millisecond)
pr = node.ProbeResult()
if pr == nil || pr.Success {
t.Fatalf("expected HTTP 500 probe failure, got %+v", pr)
}
node.Close()
}
// httpTestMux is a minimal HTTP server that returns a configurable status.
type httpTestMux struct {
status atomic.Int32
}
func (m *httpTestMux) serve(conn net.Conn) {
defer conn.Close()
buf := make([]byte, 1024)
n, _ := conn.Read(buf)
// Parse just enough to construct a response
body := "HTTP/1.1 000 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"
switch m.status.Load() {
case 200:
body = "HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"
case 500:
body = "HTTP/1.1 500 Internal Server Error\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"
}
_ = n
conn.Write([]byte(body))
}
+15
View File
@@ -157,6 +157,7 @@ type SelectorConfig struct {
Strategy string `json:"strategy"`
MaxFails int `yaml:"maxFails" json:"maxFails"`
FailTimeout time.Duration `yaml:"failTimeout" json:"failTimeout"`
MaxLatency time.Duration `yaml:"maxLatency,omitempty" json:"maxLatency,omitempty"`
}
type AdmissionConfig struct {
@@ -389,6 +390,7 @@ type ForwardNodeConfig struct {
Auth *AuthConfig `yaml:",omitempty" json:"auth,omitempty"`
HTTP *HTTPNodeConfig `yaml:",omitempty" json:"http,omitempty"`
TLS *TLSNodeConfig `yaml:",omitempty" json:"tls,omitempty"`
Probe *ProbeConfig `yaml:",omitempty" json:"probe,omitempty"`
Metadata map[string]any `yaml:",omitempty" json:"metadata,omitempty"`
}
@@ -432,6 +434,18 @@ type NodeMatcherConfig struct {
BodySize int `yaml:",omitempty" json:"bodySize,omitempty"`
}
// ProbeConfig holds the configuration for a node-level liveness probe.
type ProbeConfig struct {
Type string `json:"type"`
Addr string `json:"addr,omitempty"`
Interval time.Duration `json:"interval"`
Timeout time.Duration `json:"timeout"`
HTTPPath string `yaml:"httpPath,omitempty" json:"httpPath,omitempty"`
HTTPHost string `yaml:"httpHost,omitempty" json:"httpHost,omitempty"`
HTTPHeaders map[string]string `yaml:"httpHeaders,omitempty" json:"httpHeaders,omitempty"`
ExpectedStatus int `yaml:"expectedStatus,omitempty" json:"expectedStatus,omitempty"`
}
type HTTPNodeConfig struct {
// rewrite host header
Host string `yaml:",omitempty" json:"host,omitempty"`
@@ -635,6 +649,7 @@ type NodeConfig struct {
Matcher *NodeMatcherConfig `yaml:",omitempty" json:"matcher,omitempty"`
HTTP *HTTPNodeConfig `yaml:",omitempty" json:"http,omitempty"`
TLS *TLSNodeConfig `yaml:",omitempty" json:"tls,omitempty"`
Probe *ProbeConfig `yaml:",omitempty" json:"probe,omitempty"`
Metadata map[string]any `yaml:",omitempty" json:"metadata,omitempty"`
}
+39 -1
View File
@@ -5,6 +5,7 @@ import (
"net"
"regexp"
"strings"
"time"
"github.com/go-gost/core/chain"
"github.com/go-gost/core/connector"
@@ -281,5 +282,42 @@ func ParseNode(hop string, cfg *config.NodeConfig, log logger.Logger) (*chain.No
}
opts = append(opts, chain.TLSNodeOption(tlsCfg))
}
return chain.NewNode(cfg.Name, cfg.Addr, opts...), nil
node := chain.NewNode(cfg.Name, cfg.Addr, opts...)
if cfg.Probe != nil {
if pc := parseProbeConfig(cfg.Probe); pc != nil {
xchain.StartNodeProbe(node, pc, nodeLogger)
}
}
return node, nil
}
// parseProbeConfig converts a config.ProbeConfig into a chain.ProbeConfig.
// Returns nil when the config is invalid (e.g. empty addr).
func parseProbeConfig(cfg *config.ProbeConfig) *chain.ProbeConfig {
if cfg == nil || cfg.Addr == "" {
return nil
}
pt := chain.ProbeTypeTCP
if cfg.Type == "http" {
pt = chain.ProbeTypeHTTP
}
interval := cfg.Interval
if interval <= 0 {
interval = 30 * time.Second
}
timeout := cfg.Timeout
if timeout <= 0 {
timeout = 10 * time.Second
}
return &chain.ProbeConfig{
Type: pt,
Addr: cfg.Addr,
Interval: interval,
Timeout: timeout,
HTTPPath: cfg.HTTPPath,
HTTPHost: cfg.HTTPHost,
HTTPHeaders: cfg.HTTPHeaders,
ExpectedStatus: cfg.ExpectedStatus,
}
}
+9 -3
View File
@@ -55,15 +55,21 @@ func ParseNodeSelector(cfg *config.SelectorConfig) selector.Selector[*chain.Node
strategy = xs.HashStrategy[*chain.Node]()
case "parallel":
strategy = xs.ParallelStrategy[*chain.Node]()
case "lowestlatency", "lowest", "ll":
strategy = xs.LowestLatencyStrategy[*chain.Node]()
default:
strategy = xs.RoundRobinStrategy[*chain.Node]()
}
return xs.NewSelector(
strategy,
filters := []selector.Filter[*chain.Node]{
xs.FailFilter[*chain.Node](cfg.MaxFails, cfg.FailTimeout),
xs.BackupFilter[*chain.Node](),
)
}
if cfg.MaxLatency > 0 {
filters = append(filters, xs.LatencyFilter[*chain.Node](cfg.MaxLatency))
}
return xs.NewSelector(strategy, filters...)
}
// DefaultNodeSelector returns a node selector using round-robin strategy with
+5
View File
@@ -443,6 +443,11 @@ func (p *chainHop) parseNode(r io.Reader) ([]*chain.Node, error) {
func (p *chainHop) Close() error {
p.cancelFunc()
for _, n := range p.Nodes() {
if n != nil {
n.Close()
}
}
if p.options.fileLoader != nil {
p.options.fileLoader.Close()
}
+186
View File
@@ -0,0 +1,186 @@
package hop
import (
"context"
"net"
"testing"
"time"
"github.com/go-gost/core/chain"
"github.com/go-gost/core/connector"
xchain "github.com/go-gost/x/chain"
xlogger "github.com/go-gost/x/logger"
xs "github.com/go-gost/x/selector"
)
// fakeTransport dials a real TCP address for probe testing.
type fakeTransport struct {
dialer *net.Dialer
options chain.TransportOptions
}
func (t *fakeTransport) Dial(ctx context.Context, addr string) (net.Conn, error) {
return t.dialer.DialContext(ctx, "tcp", addr)
}
func (t *fakeTransport) Handshake(ctx context.Context, conn net.Conn) (net.Conn, error) {
return conn, nil
}
func (t *fakeTransport) Connect(ctx context.Context, conn net.Conn, network, address string) (net.Conn, error) {
return nil, nil
}
func (t *fakeTransport) Bind(ctx context.Context, conn net.Conn, network, address string, opts ...connector.BindOption) (net.Listener, error) {
return nil, nil
}
func (t *fakeTransport) Multiplex() bool { return false }
func (t *fakeTransport) Options() *chain.TransportOptions { return &t.options }
func (t *fakeTransport) Copy() chain.Transporter { return t }
func TestHopFailoverWithProbe(t *testing.T) {
// Start two listeners as healthy "proxy" targets.
l1, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatal(err)
}
defer l1.Close()
go acceptAndClose(l1)
l2, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatal(err)
}
defer l2.Close()
go acceptAndClose(l2)
addr1, addr2 := l1.Addr().String(), l2.Addr().String()
dialer := &net.Dialer{Timeout: 1 * time.Second}
nodeA := chain.NewNode("a", addr1,
chain.TransportNodeOption(&fakeTransport{dialer: dialer}),
)
nodeB := chain.NewNode("b", addr2,
chain.TransportNodeOption(&fakeTransport{dialer: dialer}),
)
// Start probe on both nodes.
xchain.StartNodeProbe(nodeA, &chain.ProbeConfig{
Type: chain.ProbeTypeTCP,
Addr: addr1,
Interval: 200 * time.Millisecond,
Timeout: 1 * time.Second,
}, xlogger.Nop())
xchain.StartNodeProbe(nodeB, &chain.ProbeConfig{
Type: chain.ProbeTypeTCP,
Addr: addr2,
Interval: 200 * time.Millisecond,
Timeout: 1 * time.Second,
}, xlogger.Nop())
// Build hop with FailFilter + BackupFilter, default 2 max fails.
sel := xs.NewSelector(
xs.RoundRobinStrategy[*chain.Node](),
xs.FailFilter[*chain.Node](1, 5*time.Second),
xs.BackupFilter[*chain.Node](),
)
h := NewHop(
NameOption("test-hop"),
NodeOption(nodeA, nodeB),
SelectorOption(sel),
LoggerOption(xlogger.Nop()),
)
defer func() {
if closer, ok := h.(interface{ Close() error }); ok {
closer.Close()
}
}()
// Wait for probes to mark both healthy.
time.Sleep(500 * time.Millisecond)
if !nodeA.ProbeResult().Success || !nodeB.ProbeResult().Success {
t.Fatalf("both nodes should be healthy initially: A=%v B=%v",
nodeA.ProbeResult().Success, nodeB.ProbeResult().Success)
}
// Both healthy → both should appear in selection.
ctx := context.Background()
sel1 := h.Select(ctx)
sel2 := h.Select(ctx)
_ = sel1
_ = sel2
// Round-robin: two picks should be different if both alive.
if sel1 == sel2 {
t.Log("round-robin picked same node twice — both may be fine, checking probe state")
}
// Kill nodeA's listener → probe marks it dead.
l1.Close()
time.Sleep(600 * time.Millisecond)
if nodeA.ProbeResult().Success {
t.Error("nodeA should be dead after listener closed")
}
if nodeA.Marker().Count() == 0 {
t.Error("nodeA marker should have failures")
}
// After FailFilter kicks in, hop should consistently pick nodeB.
var sawDead bool
for i := 0; i < 10; i++ {
n := h.Select(ctx)
if n == nil {
t.Fatal("got nil node from selector")
}
if n == nodeA {
sawDead = true
}
}
if sawDead {
t.Error("FailFilter should have excluded nodeA")
}
// Reopen listener → probe resets marker.
l1New, err := net.Listen("tcp", addr1)
if err != nil {
t.Fatal(err)
}
defer l1New.Close()
go acceptAndClose(l1New)
time.Sleep(600 * time.Millisecond)
if !nodeA.ProbeResult().Success {
t.Error("nodeA should be healthy again after listener restored")
}
if nodeA.Marker().Count() != 0 {
t.Error("nodeA marker should be reset after successful probe")
}
// Both should be reachable again.
var sawA bool
for i := 0; i < 20; i++ {
if h.Select(ctx) == nodeA {
sawA = true
break
}
}
if !sawA {
t.Error("nodeA should be selected again after recovery")
}
}
func acceptAndClose(l net.Listener) {
for {
conn, err := l.Accept()
if err != nil {
return
}
conn.Close()
}
}
+68
View File
@@ -0,0 +1,68 @@
// Package probe implements liveness probe checkers for node health monitoring.
package probe
import (
"bufio"
"fmt"
"io"
"net"
"net/http"
"github.com/go-gost/core/chain"
)
// HTTPProber performs an HTTP GET request over an established tunnel connection
// and validates the response status code.
type HTTPProber struct {
Path string
Host string
Headers map[string]string
ExpectedStatus int
}
func NewHTTPProber(cfg *chain.ProbeConfig) *HTTPProber {
path := cfg.HTTPPath
if path == "" {
path = "/"
}
host := cfg.HTTPHost
if host == "" {
host = cfg.Addr
}
expected := cfg.ExpectedStatus
if expected == 0 {
expected = http.StatusOK
}
return &HTTPProber{
Path: path,
Host: host,
Headers: cfg.HTTPHeaders,
ExpectedStatus: expected,
}
}
// Probe sends an HTTP/1.1 GET request over conn and returns an error if the
// response status does not match ExpectedStatus.
func (p *HTTPProber) Probe(conn net.Conn) error {
req := fmt.Sprintf("GET %s HTTP/1.1\r\nHost: %s\r\n", p.Path, p.Host)
for k, v := range p.Headers {
req += fmt.Sprintf("%s: %s\r\n", k, v)
}
req += "Connection: close\r\n\r\n"
if _, err := fmt.Fprint(conn, req); err != nil {
return fmt.Errorf("http probe write: %w", err)
}
resp, err := http.ReadResponse(bufio.NewReader(conn), nil)
if err != nil {
return fmt.Errorf("http probe read: %w", err)
}
defer resp.Body.Close()
io.Copy(io.Discard, resp.Body)
if resp.StatusCode != p.ExpectedStatus {
return fmt.Errorf("http probe: expected status %d, got %d", p.ExpectedStatus, resp.StatusCode)
}
return nil
}
+118
View File
@@ -0,0 +1,118 @@
package probe
import (
"net"
"net/http"
"testing"
"time"
"github.com/go-gost/core/chain"
)
func TestHTTPProber(t *testing.T) {
t.Run("200 OK", func(t *testing.T) {
server, client := net.Pipe()
defer server.Close()
defer client.Close()
go func() {
// Minimal HTTP/1.1 server — just respond 200 OK
http.Serve(&singleConnListener{conn: server}, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
}))
}()
p := NewHTTPProber(&chain.ProbeConfig{
HTTPPath: "/health",
HTTPHost: "example.com",
ExpectedStatus: http.StatusOK,
})
time.Sleep(50 * time.Millisecond) // let server start
if err := p.Probe(client); err != nil {
t.Fatalf("unexpected error: %v", err)
}
})
t.Run("wrong status", func(t *testing.T) {
server, client := net.Pipe()
defer server.Close()
defer client.Close()
go func() {
http.Serve(&singleConnListener{conn: server}, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusInternalServerError)
}))
}()
p := NewHTTPProber(&chain.ProbeConfig{
HTTPPath: "/",
ExpectedStatus: http.StatusOK,
})
time.Sleep(50 * time.Millisecond)
if err := p.Probe(client); err == nil {
t.Fatal("expected error for 500 status, got nil")
}
})
t.Run("defaults", func(t *testing.T) {
p := NewHTTPProber(&chain.ProbeConfig{
Addr: "example.com:80",
})
if p.Path != "/" {
t.Errorf("expected default path '/', got %q", p.Path)
}
if p.ExpectedStatus != http.StatusOK {
t.Errorf("expected default status 200, got %d", p.ExpectedStatus)
}
if p.Host != "example.com:80" {
t.Errorf("expected default host from Addr, got %q", p.Host)
}
})
t.Run("connection error", func(t *testing.T) {
server, client := net.Pipe()
server.Close() // kill it immediately
client.Close()
p := NewHTTPProber(&chain.ProbeConfig{
HTTPPath: "/",
ExpectedStatus: http.StatusOK,
})
if err := p.Probe(client); err == nil {
t.Fatal("expected error on closed connection, got nil")
}
})
}
// singleConnListener is a net.Listener that returns a single pre-established
// connection, then returns an error for subsequent Accept calls.
type singleConnListener struct {
conn net.Conn
closed bool
}
func (l *singleConnListener) Accept() (net.Conn, error) {
if l.closed {
return nil, net.ErrClosed
}
if l.conn != nil {
c := l.conn
l.conn = nil
return c, nil
}
return nil, &net.OpError{} // block-like
}
func (l *singleConnListener) Close() error {
l.closed = true
if l.conn != nil {
return l.conn.Close()
}
return nil
}
func (l *singleConnListener) Addr() net.Addr {
return &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0}
}
+85
View File
@@ -0,0 +1,85 @@
package selector
import (
"context"
"time"
"github.com/go-gost/core/chain"
"github.com/go-gost/core/selector"
)
type latencyFilter[T any] struct {
maxLatency time.Duration
}
// LatencyFilter filters out items whose probe result indicates failure or
// latency exceeding maxLatency. Items without a probe result are passed through
// (no instrumentation = no filtering).
func LatencyFilter[T any](maxLatency time.Duration) selector.Filter[T] {
return &latencyFilter[T]{maxLatency: maxLatency}
}
func (f *latencyFilter[T]) Filter(ctx context.Context, vs ...T) []T {
if len(vs) <= 1 || f.maxLatency <= 0 {
return vs
}
var out []T
for _, v := range vs {
if reader, ok := any(v).(chain.ProbeResultReader); ok {
r := reader.ProbeResult()
if r != nil && (!r.Success || r.Latency > f.maxLatency) {
continue
}
}
out = append(out, v)
}
return out
}
type lowestLatencyStrategy[T any] struct{}
// LowestLatencyStrategy selects the item with the lowest probe latency.
// Items without a probe result or with failed probes are deprioritized
// (placed after items with successful probes).
func LowestLatencyStrategy[T any]() selector.Strategy[T] {
return &lowestLatencyStrategy[T]{}
}
func (s *lowestLatencyStrategy[T]) Apply(ctx context.Context, vs ...T) (v T) {
if len(vs) == 0 {
return
}
var bestLatency time.Duration
var bestIdx int = -1
var fallback int = -1
for i, item := range vs {
reader, ok := any(item).(chain.ProbeResultReader)
if !ok {
if fallback < 0 {
fallback = i
}
continue
}
r := reader.ProbeResult()
if r == nil || !r.Success {
if fallback < 0 {
fallback = i
}
continue
}
if bestIdx < 0 || r.Latency < bestLatency {
bestLatency = r.Latency
bestIdx = i
}
}
if bestIdx >= 0 {
return vs[bestIdx]
}
if fallback >= 0 {
return vs[fallback]
}
return vs[0]
}
+154
View File
@@ -0,0 +1,154 @@
package selector
import (
"context"
"testing"
"time"
"github.com/go-gost/core/chain"
)
// mockProbeReader implements chain.ProbeResultReader.
type mockProbeReader struct {
result *chain.ProbeResult
}
func (m *mockProbeReader) ProbeResult() *chain.ProbeResult { return m.result }
func TestLowestLatencyStrategy(t *testing.T) {
strategy := LowestLatencyStrategy[*mockProbeReader]()
ctx := context.Background()
fast := &mockProbeReader{result: &chain.ProbeResult{Success: true, Latency: 10 * time.Millisecond}}
mid := &mockProbeReader{result: &chain.ProbeResult{Success: true, Latency: 50 * time.Millisecond}}
slow := &mockProbeReader{result: &chain.ProbeResult{Success: true, Latency: 100 * time.Millisecond}}
fail := &mockProbeReader{result: &chain.ProbeResult{Success: false, Latency: 5 * time.Millisecond}}
noResult := &mockProbeReader{result: nil}
// fastest wins
if got := strategy.Apply(ctx, fast, mid, slow); got != fast {
t.Error("expected fastest (10ms)")
}
// order-independent: fastest wins regardless of position
if got := strategy.Apply(ctx, slow, mid, fast); got != fast {
t.Error("expected fastest (10ms) even when last")
}
// failed probe deprioritized — picks the successful one
if got := strategy.Apply(ctx, fail, mid); got != mid {
t.Error("expected mid (50ms) over failed")
}
// nil result deprioritized — picks the one with a result
if got := strategy.Apply(ctx, noResult, fast); got != fast {
t.Error("expected fast over nil-result")
}
// all failed → fallback to first item (even if failed)
if got := strategy.Apply(ctx, fail); got != fail {
t.Error("expected fallback to only item")
}
// failed + nil → first encountered wins as fallback
if got := strategy.Apply(ctx, fail, noResult); got != fail {
t.Error("expected first fallback item (failed before nil)")
}
if got := strategy.Apply(ctx, noResult, fail); got != noResult {
t.Error("expected first fallback item (nil before failed)")
}
}
func TestLowestLatencyStrategyOnNodes(t *testing.T) {
strategy := LowestLatencyStrategy[*chain.Node]()
ctx := context.Background()
fast := chain.NewNode("fast", "addr1")
fast.SetProbeResult(&chain.ProbeResult{Success: true, Latency: 10 * time.Millisecond})
slow := chain.NewNode("slow", "addr2")
slow.SetProbeResult(&chain.ProbeResult{Success: true, Latency: 100 * time.Millisecond})
dead := chain.NewNode("dead", "addr3")
dead.SetProbeResult(&chain.ProbeResult{Success: false})
bare := chain.NewNode("bare", "addr4")
// fastest wins
if got := strategy.Apply(ctx, fast, slow, dead, bare); got != fast {
t.Error("expected fast node")
}
// only dead + bare → first fallback (dead) wins, even though bare has no result
if got := strategy.Apply(ctx, dead, bare); got != dead {
t.Errorf("expected first fallback (dead), got %v", got)
}
// only bare + dead → first fallback (bare) wins
if got := strategy.Apply(ctx, bare, dead); got != bare {
t.Errorf("expected first fallback (bare), got %v", got)
}
}
func TestLatencyFilter(t *testing.T) {
ctx := context.Background()
filter := LatencyFilter[*mockProbeReader](50 * time.Millisecond)
ok := &mockProbeReader{result: &chain.ProbeResult{Success: true, Latency: 10 * time.Millisecond}}
over := &mockProbeReader{result: &chain.ProbeResult{Success: true, Latency: 100 * time.Millisecond}}
fail := &mockProbeReader{result: &chain.ProbeResult{Success: false, Latency: 1 * time.Millisecond}}
noResult := &mockProbeReader{result: nil}
// ok passes, over is filtered
if got := filter.Filter(ctx, ok, over); len(got) != 1 || got[0] != ok {
t.Error("expected only ok to pass")
}
// single item: len(vs) <= 1 guard keeps it (safety net — don't empty a hop)
if got := filter.Filter(ctx, fail); len(got) != 1 {
t.Error("single failed item preserved (len <= 1 guard)")
}
if got := filter.Filter(ctx, over); len(got) != 1 {
t.Error("single over-threshold item preserved (len <= 1 guard)")
}
// multi-item: failed is filtered, nil and ok pass through
if got := filter.Filter(ctx, ok, fail, noResult); len(got) != 2 {
t.Error("expected 2 pass (ok, noResult)")
}
// zero maxLatency is no-op (all pass)
filter = LatencyFilter[*mockProbeReader](0)
if got := filter.Filter(ctx, ok, over, fail); len(got) != 3 {
t.Error("expected all pass through with zero maxLatency")
}
}
func TestLatencyFilterOnNodes(t *testing.T) {
filter := LatencyFilter[*chain.Node](50 * time.Millisecond)
ctx := context.Background()
ok := chain.NewNode("ok", "addr1")
ok.SetProbeResult(&chain.ProbeResult{Success: true, Latency: 10 * time.Millisecond})
over := chain.NewNode("over", "addr2")
over.SetProbeResult(&chain.ProbeResult{Success: true, Latency: 200 * time.Millisecond})
dead := chain.NewNode("dead", "addr3")
dead.SetProbeResult(&chain.ProbeResult{Success: false})
bare := chain.NewNode("bare", "addr4")
got := filter.Filter(ctx, ok, over, dead, bare)
if len(got) != 2 {
t.Fatalf("expected 2 pass (ok, bare), got %d", len(got))
}
names := map[string]bool{}
for _, n := range got {
names[n.Name] = true
}
if !names["ok"] || !names["bare"] {
t.Errorf("expected ok and bare, got %v", names)
}
}