internal/home/cluster_balance_regression_test.go
Rama main · réplica de lectura del repositorio original.
package home
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"net"
"strings"
"sync"
"testing"
"time"
"github.com/router-for-me/CLIProxyAPI/v8/internal/config"
log "github.com/sirupsen/logrus"
)
func TestClusterBalanceResponseErrors(t *testing.T) {
for _, operation := range []string{"ENABLE", "RESULT", "CLAIM"} {
for _, response := range []struct {
name string
raw string
fail bool
}{
{name: "malformed", raw: `{`, fail: true},
{name: "wrong type", raw: `{"ok":"true"}`, fail: true},
{name: "rejected", raw: `{"ok":false}`},
} {
t.Run(operation+"/"+response.name, func(t *testing.T) {
client, _ := newRedisCommandTestClient(t, func(args []string) string {
if strings.EqualFold(args[0], "CLUSTER") {
return fmt.Sprintf("$%d\r\n%s\r\n", len(response.raw), response.raw)
}
return balanceTestRedisReply(args)
})
enableBalanceTestClient(client)
var errResponse error
switch operation {
case "ENABLE":
client.balance.registered = false
errResponse = client.syncClusterBalance(context.Background())
case "RESULT":
client.balance.request.ID, client.balance.result = "pending", "failed"
errResponse = client.syncClusterBalance(context.Background())
if client.balance.request.ID != "pending" || client.balance.result != "failed" {
t.Fatal("unacknowledged result was discarded")
}
case "CLAIM":
errResponse = client.executeClusterRebalance(context.Background(), []byte(`{"id":"plan","term":1,"target":{"ip":"other","port":8327}}`))
}
if (errResponse != nil) != response.fail {
t.Fatalf("response error = %v, want failure %v", errResponse, response.fail)
}
if response.fail && !strings.Contains(errResponse.Error(), operation+" response") {
t.Fatalf("error lacks protocol context: %v", errResponse)
}
if client.balance.unsupported || client.balance.moving || (operation == "ENABLE" && client.balance.registered) {
t.Fatal("invalid/rejected response advanced balance state or disabled capability")
}
})
}
}
}
func TestClusterBalanceRegistrationErrorsDoNotDisableCapability(t *testing.T) {
client, _ := newRedisCommandTestClient(t, func(args []string) string {
if strings.EqualFold(args[0], "CLUSTER") {
return "-ERR membership_required\r\n"
}
return balanceTestRedisReply(args)
})
enableBalanceTestClient(client)
client.balance.registered = false
if errSync := client.syncClusterBalance(context.Background()); errSync == nil {
t.Fatal("membership failure was swallowed")
}
if client.balance.unsupported || client.balance.registered {
t.Fatal("membership failure disabled capability or registered client")
}
}
func TestClusterBalanceBlockedRPCDoesNotDelayHeartbeatExpiry(t *testing.T) {
for _, blockedOperation := range []string{"ENABLE", "CLAIM", "PING", "RESULT"} {
t.Run(blockedOperation, func(t *testing.T) {
blocked := make(chan struct{})
release := make(chan struct{})
var blockOnce sync.Once
block := func() {
blockOnce.Do(func() { close(blocked) })
<-release
}
t.Cleanup(func() { close(release) })
target, _ := newRedisCommandTestClient(t, func(args []string) string {
if blockedOperation == "PING" && strings.EqualFold(args[0], "PING") {
block()
}
return balanceTestRedisReply(args)
})
instruction, errMarshal := json.Marshal(rebalanceInstruction{
ID: "blocked-plan", Term: 1, Target: clusterNode{IP: target.homeCfg.Host, Port: target.homeCfg.Port},
})
if errMarshal != nil {
t.Fatal(errMarshal)
}
client, commands := newRedisCommandTestClient(t, func(args []string) string {
switch strings.ToUpper(args[0]) {
case "CLUSTER":
if args[1] == "BALANCE" {
if args[2] == blockedOperation {
block()
}
return "$11\r\n{\"ok\":true}\r\n"
}
return "$22\r\n{\"ok\":true,\"nodes\":[]}\r\n"
case "GET":
return "$2\r\n{}\r\n"
case "SUBSCRIBE":
reply := "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:1\r\n"
if blockedOperation == "CLAIM" || blockedOperation == "PING" {
reply += fmt.Sprintf("*3\r\n$7\r\nmessage\r\n$9\r\nrebalance\r\n$%d\r\n%s\r\n", len(instruction), instruction)
}
return reply
default:
return balanceTestRedisReply(args)
}
})
enableBalanceTestClient(client)
client.balance.registered = false
client.testOperationTimeout = 5 * time.Second
if blockedOperation == "RESULT" {
client.balance.request.ID, client.balance.result = "previous-plan", "success"
}
ctx, cancel := context.WithCancel(context.Background())
t.Cleanup(cancel)
done := make(chan error, 1)
go func() {
done <- client.RunConfigSubscriberLifetime(ctx, func([]byte) error {
return client.SetLifecycleConfig(config.CredentialConcurrencyConfig{
LifecycleConfigRevision: 1, CPAHeartbeatTimeout: 200 * time.Millisecond,
})
}, nil)
}()
select {
case <-blocked:
case errRun := <-done:
t.Fatalf("subscriber ended before blocked RPC: %v", errRun)
case <-time.After(3 * time.Second):
t.Fatal("RPC did not start")
}
select {
case errRun := <-done:
if !isTimeoutError(errRun) {
t.Fatalf("subscriber error = %v, want heartbeat timeout", errRun)
}
case <-time.After(2 * time.Second):
t.Fatal("blocked RPC delayed heartbeat expiry or subscriber shutdown")
}
if client.HeartbeatOK() || !client.dispatchFenced.Load() {
t.Fatal("heartbeat expiry did not fence dispatch admission")
}
if _, errDispatch := client.RPopAuth(context.Background(), "test-model", "", nil, 1); !errors.Is(errDispatch, ErrDispatchFenced) {
t.Fatalf("late dispatch = %v, want fenced", errDispatch)
}
if len(findRedisCommand(commands.All(), "RPOP")) != 0 || client.balance.moving {
t.Fatal("expired subscriber issued dispatch or migrated")
}
client.mu.Lock()
commandPool := client.cmd
client.mu.Unlock()
if commandPool == nil {
t.Fatal("migration cancellation detached the source command pool")
}
if errPing := commandPool.Ping(context.Background()).Err(); errPing != nil {
t.Fatalf("migration cancellation interrupted the source command pool: %v", errPing)
}
})
}
}
func TestClusterBalanceMigrationLogIncludesPlanAndAddresses(t *testing.T) {
var output bytes.Buffer
logger := log.StandardLogger()
previousOutput, previousLevel := logger.Out, logger.GetLevel()
logger.SetOutput(&output)
logger.SetLevel(log.InfoLevel)
t.Cleanup(func() {
logger.SetOutput(previousOutput)
logger.SetLevel(previousLevel)
})
target, _ := newRedisCommandTestClient(t, balanceTestRedisReply)
client, _ := newRedisCommandTestClient(t, func(args []string) string {
if strings.EqualFold(args[0], "CLUSTER") {
return "$11\r\n{\"ok\":true}\r\n"
}
return balanceTestRedisReply(args)
})
enableBalanceTestClient(client)
origin := clusterNodeAddress(clusterNode{IP: client.homeCfg.Host, Port: client.homeCfg.Port})
request := rebalanceInstruction{ID: "logged-plan", Term: 1, Target: clusterNode{IP: target.homeCfg.Host, Port: target.homeCfg.Port}}
payload, errMarshal := json.Marshal(request)
if errMarshal != nil {
t.Fatal(errMarshal)
}
if errExecute := client.executeClusterRebalance(context.Background(), payload); !errors.Is(errExecute, ErrRebalanceRequested) {
t.Fatal(errExecute)
}
// The message itself must carry these values, independently of formatter field allowlists.
line := strings.ReplaceAll(output.String(), `\"`, `"`)
for key, value := range map[string]string{"plan": request.ID, "from": origin, "to": clusterNodeAddress(request.Target)} {
if !strings.Contains(line, fmt.Sprintf("%s=%q", key, value)) {
t.Fatalf("migration log missing %s: %q", key, line)
}
}
}
func TestClusterBalanceMalformedRegistrationEndsSubscriberBeforeReady(t *testing.T) {
client, _ := newRedisCommandTestClient(t, func(args []string) string {
switch strings.ToUpper(args[0]) {
case "CLUSTER":
if args[1] == "BALANCE" {
return "$1\r\n{\r\n"
}
return "$22\r\n{\"ok\":true,\"nodes\":[]}\r\n"
case "GET":
return "$2\r\n{}\r\n"
case "SUBSCRIBE":
return "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:1\r\n"
default:
return balanceTestRedisReply(args)
}
})
enableBalanceTestClient(client)
client.balance.registered = false
ready := false
errRun := client.RunConfigSubscriberLifetime(context.Background(), func([]byte) error { return nil }, func() { ready = true })
var syntaxError *json.SyntaxError
if !errors.As(errRun, &syntaxError) || !strings.Contains(errRun.Error(), "ENABLE response") {
t.Fatalf("subscriber did not propagate decoding failure: %v", errRun)
}
if ready || client.HeartbeatOK() {
t.Fatal("malformed registration published a healthy subscriber")
}
}
type blockingOuterCloseConn struct {
net.Conn
enteredClose chan struct{}
enterOnce sync.Once
underlyingClosed chan struct{}
}
func (c *blockingOuterCloseConn) NetConn() net.Conn {
return c.Conn
}
func (c *blockingOuterCloseConn) Close() error {
if c.enteredClose != nil {
c.enterOnce.Do(func() { close(c.enteredClose) })
}
<-c.underlyingClosed
return c.Conn.Close()
}
type trackingUnderlyingConn struct {
net.Conn
closed chan struct{}
once sync.Once
}
func (c *trackingUnderlyingConn) Close() error {
c.once.Do(func() { close(c.closed) })
return c.Conn.Close()
}
func TestCancelableConnClosesUnderlyingTransportWhenCloseRacesCancel(t *testing.T) {
for _, mode := range []string{"pre-canceled", "after-wrap", "during-close"} {
t.Run(mode, func(t *testing.T) {
for range 32 {
local, peer := net.Pipe()
t.Cleanup(func() {
_ = local.Close()
_ = peer.Close()
})
underlying := &trackingUnderlyingConn{Conn: local, closed: make(chan struct{})}
enteredClose := make(chan struct{})
outer := &blockingOuterCloseConn{
Conn: underlying,
enteredClose: enteredClose,
underlyingClosed: underlying.closed,
}
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
if mode == "pre-canceled" {
cancel()
}
wrapped := newPluginSyncCancelableConn(ctx, outer)
if mode == "after-wrap" {
cancel()
}
done := make(chan error, 1)
go func() {
done <- wrapped.Close()
}()
if mode == "during-close" {
<-enteredClose
cancel()
}
select {
case errClose := <-done:
if errClose != nil && !errors.Is(errClose, net.ErrClosed) {
t.Fatalf("close error = %v", errClose)
}
case <-time.After(500 * time.Millisecond):
_ = underlying.Close()
t.Fatal("Close blocked because underlying transport was not closed before outer Close")
}
}
})
}
}