mirror of
https://github.com/go-gost/x.git
synced 2026-10-19 16:16:32 +00:00
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:
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
@@ -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"`
|
||||
}
|
||||
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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}
|
||||
}
|
||||
@@ -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]
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user