From befcba239a36abaef2a1f5421e3370de3b577db0 Mon Sep 17 00:00:00 2001 From: hkfires <10558748+hkfires@users.noreply.github.com> Date: Wed, 21 Oct 2026 15:14:26 +0800 Subject: [PATCH] feat(home): add planned cluster rebalance handoff - Register for Home-directed balancing and execute only claimed, preflighted rebalance plans via a replacement takeover lifetime, reporting success or rolling back to the origin on failure - Select least-loaded nodes for bootstrap and failover instead of always switching to the first discovered node - Split dispatch and release fencing so the acknowledged source keeps final concurrency releases until close, and fence admission and settle pending dispatches before planned migration teardown - Isolate rebalance and plugin-status transports and receive subscription events concurrently so blocked RPCs or plugin cancellation cannot interrupt issued dispatches or releases --- internal/home/client.go | 230 +++++++--- internal/home/client_cancellation_test.go | 121 ++++++ internal/home/client_fence_test.go | 119 ++++++ .../home/client_release_lifecycle_test.go | 182 ++++++++ internal/home/cluster_balance.go | 249 +++++++++++ .../home/cluster_balance_recovery_test.go | 386 +++++++++++++++++ .../home/cluster_balance_regression_test.go | 243 +++++++++++ internal/home/cluster_balance_test.go | 263 ++++++++++++ internal/logging/home_rebalance_log_test.go | 48 +++ .../service_executionregistry_test.go | 4 +- sdk/cliproxy/service_home.go | 43 +- .../service_home_cancellation_test.go | 139 ++++++ ...ervice_home_rebalance_finalization_test.go | 159 +++++++ .../service_home_rebalance_release_test.go | 221 ++++++++++ sdk/cliproxy/service_home_rebalance_test.go | 402 ++++++++++++++++++ 15 files changed, 2732 insertions(+), 77 deletions(-) create mode 100644 internal/home/client_cancellation_test.go create mode 100644 internal/home/client_fence_test.go create mode 100644 internal/home/client_release_lifecycle_test.go create mode 100644 internal/home/cluster_balance.go create mode 100644 internal/home/cluster_balance_recovery_test.go create mode 100644 internal/home/cluster_balance_regression_test.go create mode 100644 internal/home/cluster_balance_test.go create mode 100644 internal/logging/home_rebalance_log_test.go create mode 100644 sdk/cliproxy/service_home_cancellation_test.go create mode 100644 sdk/cliproxy/service_home_rebalance_finalization_test.go create mode 100644 sdk/cliproxy/service_home_rebalance_release_test.go create mode 100644 sdk/cliproxy/service_home_rebalance_test.go diff --git a/internal/home/client.go b/internal/home/client.go index 5f7777dc7e..fe0eec05b2 100644 --- a/internal/home/client.go +++ b/internal/home/client.go @@ -7,6 +7,7 @@ import ( "encoding/json" "errors" "fmt" + "math/rand/v2" "net" "net/http" "net/url" @@ -190,9 +191,12 @@ type Client struct { lifecycle config.CredentialConcurrencyConfig limiter atomic.Pointer[config.CredentialConcurrencyConfig] managed bool + // releaseReady belongs to this acknowledged membership and is never copied to a replacement. + releaseReady bool heartbeatOK atomic.Bool dispatchFenced atomic.Bool + releaseFenced atomic.Bool ambiguousDispatch atomic.Bool // casUnsupported latches when Home does not implement the CAS command. // It is deliberately NOT carried across NewLifetime: CAS support is a @@ -207,6 +211,7 @@ type Client struct { legacyMembership bool clusterNodes []clusterNode reconnectFailures int + balance clusterBalanceState } func New(homeCfg config.HomeConfig) *Client { @@ -234,8 +239,18 @@ func (c *Client) NewLifetime() *Client { testOperationTimeout: c.testOperationTimeout, instanceID: c.instanceID, legacyMembership: c.legacyMembership, + balance: c.balance, } + next.balance.registered = false + next.balance.unsupported = false next.recoveryState.Store(c.recoveryState.Load()) + if next.balance.moving { + next.homeCfg.Host = next.balance.target.IP + next.homeCfg.Port = next.balance.target.Port + if !next.recoveryState.CompareAndSwap(uint32(recoveryStateTakeoverEligible), uint32(recoveryStateSwitchingTakeover)) { + next.recoveryState.CompareAndSwap(uint32(recoveryStateStable), uint32(recoveryStateSwitching)) + } + } return next } @@ -289,12 +304,20 @@ func (c *Client) HeartbeatOK() bool { return c.heartbeatOK.Load() } -// Close permanently ends this client's dispatch lifetime. +// FenceDispatch stops new credential acquisitions while preserving transports and final concurrency releases. +func (c *Client) FenceDispatch() { + if c != nil { + c.dispatchFenced.Store(true) + } +} + +// Close permanently ends this client's dispatch and release lifetime. func (c *Client) Close() { if c == nil { return } c.dispatchFenced.Store(true) + c.releaseFenced.Store(true) c.heartbeatOK.Store(false) c.mu.Lock() commandClient, subscriptionClient, connections := c.detachClientsLocked() @@ -318,6 +341,7 @@ func (c *Client) closeBootstrapPools() { } c.heartbeatOK.Store(false) c.mu.Lock() + c.releaseReady = false commandClient, subscriptionClient, connections := c.detachClientsLocked() c.mu.Unlock() closeDetachedClients(commandClient, subscriptionClient, connections) @@ -330,6 +354,7 @@ func (c *Client) AbortAmbiguousDispatch() { } c.ambiguousDispatch.Store(true) c.dispatchFenced.Store(true) + c.releaseFenced.Store(true) c.heartbeatOK.Store(false) c.mu.Lock() commandClient, subscriptionClient, connections := c.detachClientsLocked() @@ -785,6 +810,7 @@ func (c *Client) refreshClusterNodes(ctx context.Context) (bool, error) { if !c.clusterDiscoveryEnabled() { return false, nil } + if ctx == nil { ctx = context.Background() } @@ -805,7 +831,6 @@ func (c *Client) refreshClusterNodes(ctx context.Context) (bool, error) { if errText != nil { return false, errText } - nodes, errParse := parseClusterNodesPayload([]byte(raw)) if errParse != nil { return false, errParse @@ -813,12 +838,16 @@ func (c *Client) refreshClusterNodes(ctx context.Context) (bool, error) { if len(nodes) == 0 { return false, nil } - c.mu.Lock() defer c.mu.Unlock() c.clusterNodes = nodes - c.reconnectFailures = 0 - return c.switchToNodeLocked(nodes[0]), nil + // Discovery alone does not establish recovery or reset reconnect failures. + // Keep planned handoffs pinned here; repeated failures may still trigger failover. + if c.balance.moving || c.balance.returning || c.heartbeatOK.Load() { + return false, nil + } + node, ok := c.chooseClusterNodeLocked(nodes, false, rand.Float64()) + return ok && c.switchToNodeLocked(node), nil } func parseClusterNodesPayload(raw []byte) ([]clusterNode, error) { @@ -845,11 +874,17 @@ func (c *Client) updateClusterNodesFromPayload(raw []byte) error { func normalizeClusterNodes(nodes []clusterNode) []clusterNode { out := make([]clusterNode, 0, len(nodes)) + seen := make(map[string]bool, len(nodes)) for _, node := range nodes { node.IP = strings.TrimSpace(node.IP) - if node.IP == "" || node.Port <= 0 { + if node.IP == "" || node.Port <= 0 || node.Port > 65535 { continue } + address := clusterNodeAddress(node) + if seen[address] { + continue + } + seen[address] = true if node.ClientCount < 0 { node.ClientCount = 0 } @@ -869,6 +904,7 @@ func (c *Client) switchToNodeLocked(node clusterNode) bool { if strings.TrimSpace(c.homeCfg.Host) == host && c.homeCfg.Port == node.Port { return false } + c.releaseReady = false c.homeCfg.Host = host c.homeCfg.Port = node.Port if !c.recoveryState.CompareAndSwap(uint32(recoveryStateStable), uint32(recoveryStateSwitching)) { @@ -923,22 +959,13 @@ func (c *Client) failoverAfterSubscriptionTimeout() (bool, string) { func (c *Client) switchToNextNodeLocked() (bool, string) { currentHost := strings.TrimSpace(c.homeCfg.Host) currentPort := c.homeCfg.Port - candidates := append([]clusterNode(nil), c.clusterNodes...) - if strings.TrimSpace(c.seedHost) != "" && c.seedPort > 0 { - candidates = append(candidates, clusterNode{IP: c.seedHost, Port: c.seedPort}) + node, ok := c.chooseClusterNodeLocked(c.clusterNodes, true, rand.Float64()) + if !ok && strings.TrimSpace(c.seedHost) != "" && c.seedPort > 0 && (c.seedHost != currentHost || c.seedPort != currentPort) { + node, ok = clusterNode{IP: c.seedHost, Port: c.seedPort}, true } - for _, node := range candidates { - host := strings.TrimSpace(node.IP) - if host == "" || node.Port <= 0 { - continue - } - if host == currentHost && node.Port == currentPort { - continue - } - if c.switchToNodeLocked(clusterNode{IP: host, Port: node.Port}) { - addr, _ := c.addrLocked() - return true, addr - } + if ok && c.switchToNodeLocked(node) { + addr, _ := c.addrLocked() + return true, addr } return false, "" } @@ -1505,24 +1532,23 @@ func (c *Client) PushConcurrencyRelease(ctx context.Context, frame ConcurrencyRe } func (c *Client) concurrencyReleaseClient() (*redis.Client, error) { - if c == nil || c.dispatchFenced.Load() { + if c == nil || c.releaseFenced.Load() { return nil, ErrDispatchFenced } - state := recoveryState(c.recoveryState.Load()) - if state == recoveryStateTakeoverEligible || state == recoveryStateSwitching || state == recoveryStateSwitchingTakeover { - return nil, ErrNotConnected - } if !c.Enabled() { return nil, ErrDisabled } c.mu.Lock() defer c.mu.Unlock() - if c.dispatchFenced.Load() { + if c.releaseFenced.Load() { return nil, ErrDispatchFenced } - state = recoveryState(c.recoveryState.Load()) - if state == recoveryStateTakeoverEligible || state == recoveryStateSwitching || state == recoveryStateSwitchingTakeover { + state := recoveryState(c.recoveryState.Load()) + // Preparing takeover does not move membership. The acknowledged source owns final releases + // until Close; an unacknowledged replacement or a switched address does not inherit that right. + ownsSource := state == recoveryStateTakeoverEligible && c.releaseReady + if state != recoveryStateStable && !ownsSource { return nil, ErrNotConnected } if c.release != nil { @@ -1564,14 +1590,38 @@ func (c *Client) RPushAppLog(ctx context.Context, payload []byte) error { } func (c *Client) RPushPluginStatus(ctx context.Context, payload []byte) error { - cmd, errClient := c.commandClient() - if errClient != nil { - return errClient + options, errOptions := c.pluginSyncCommandOptions() + if errOptions != nil { + return errOptions } if len(payload) == 0 { return nil } - return cmd.RPush(ctx, redisKeyPluginStatus, payload).Err() + if ctx == nil { + ctx = context.Background() + } + // Plugin cancellation owns only this command's transport, never dispatch or release connections. + baseDialer := options.Dialer + if baseDialer == nil { + baseDialer = pluginSyncDialer(options) + } + options.Dialer = func(dialCtx context.Context, network, address string) (net.Conn, error) { + conn, errDial := baseDialer(dialCtx, network, address) + if errDial != nil { + return nil, errDial + } + return newPluginSyncCancelableConn(ctx, conn), nil + } + statusClient := redis.NewClient(options) + errPush := statusClient.RPush(ctx, redisKeyPluginStatus, payload).Err() + errClose := statusClient.Close() + if errContext := ctx.Err(); errContext != nil { + return errContext + } + if errClose != nil { + return errors.Join(errPush, fmt.Errorf("close plugin status client: %w", errClose)) + } + return errPush } func (c *Client) GetPluginTasks(ctx context.Context) ([]PluginTask, error) { @@ -1873,7 +1923,10 @@ func (c *Client) rebuildCommandPoolAndProbe(ctx context.Context) error { if errPing := c.Ping(ctx); errPing != nil { return errPing } + c.mu.Lock() + c.releaseReady = true c.recoveryState.Store(uint32(recoveryStateStable)) + c.mu.Unlock() return nil } @@ -1906,7 +1959,12 @@ func (c *Client) handleSubscriptionPayload(ctx context.Context, channel string, } return onConfig([]byte(payload)) case redisChannelCluster: - return c.updateClusterNodesFromPayload([]byte(payload)) + if errUpdate := c.updateClusterNodesFromPayload([]byte(payload)); errUpdate != nil { + return errUpdate + } + return c.syncClusterBalance(ctx) + case redisChannelRebalance: + return c.executeClusterRebalance(ctx, []byte(payload)) default: return nil } @@ -1979,37 +2037,99 @@ func (c *Client) RunConfigSubscriberLifetime(ctx context.Context, onConfig func( } c.resetReconnectFailures() c.heartbeatOK.Store(true) + // Reception must keep monitoring silence while the consumer performs migration RPCs. + receiveCtx, cancelReceive := context.WithCancelCause(ctx) + events := make(chan any, 16) + receiveDone := make(chan struct{}) + stopReceive := context.AfterFunc(receiveCtx, func() { _ = pubsub.Close() }) + go func() { + defer close(receiveDone) + for { + _, timeout := c.subscriptionParameters() + event, errReceive := pubsub.ReceiveTimeout(receiveCtx, timeout) + if errReceive == nil { + select { + case events <- event: + continue + case <-receiveCtx.Done(): + errReceive = context.Cause(receiveCtx) + default: + // Bounded backpressure fails the lifetime rather than dropping control messages. + errReceive = fmt.Errorf("Home subscription consumer backlog exceeded") + } + } + if !errors.Is(context.Cause(receiveCtx), ErrRebalanceRequested) { + if ctx.Err() == nil { + if c.heartbeatOK.Load() { + c.markMembershipTakeoverEligible() + } + c.releaseFenced.Store(true) + } + // Graceful cancellation stops new work but keeps final releases available until Close. + c.dispatchFenced.Store(true) + } + c.heartbeatOK.Store(false) + cancelReceive(errReceive) + return + } + }() + defer func() { + cancelReceive(context.Canceled) + stopReceive() + _ = pubsub.Close() + <-receiveDone + }() + endReceiveFailure := func(err error) error { + if cause := context.Cause(receiveCtx); cause != nil { + err = cause + } + if ctx.Err() == nil { + // Processing failures must preserve the committed membership before cleanup clears heartbeatOK. + if len(args) > 1 { + c.markMembershipTakeoverEligible() + } + if isTimeoutError(err) { + c.markSubscriptionTimeout() + } else { + c.markReconnectFailure("subscription") + } + } + return c.endConfigSubscriberLifetimeWithSubscription(err, pubsub, "subscription failure") + } + if errFinish := c.finishClusterRebalance(receiveCtx); errFinish != nil || receiveCtx.Err() != nil { + return endReceiveFailure(errFinish) + } if onReady != nil { onReady() } for { - _, receiveTimeout = c.subscriptionParameters() - event, errReceive := pubsub.ReceiveTimeout(ctx, receiveTimeout) - if errReceive != nil { - if ctx.Err() == nil { - if c.heartbeatOK.Load() { - c.markMembershipTakeoverEligible() - } - if isTimeoutError(errReceive) { - c.markSubscriptionTimeout() - } else { - c.markReconnectFailure("subscription") - } - } - return c.endConfigSubscriberLifetimeWithSubscription(errReceive, pubsub, "heartbeat loss") + var event any + select { + case <-receiveCtx.Done(): + return endReceiveFailure(context.Cause(receiveCtx)) + case event = <-events: + } + if receiveCtx.Err() != nil { + return endReceiveFailure(context.Cause(receiveCtx)) } switch msg := event.(type) { case *redis.Message: if msg == nil { continue } - if errApply := c.handleSubscriptionPayload(ctx, msg.Channel, msg.Payload, onConfig); errApply != nil { - if strings.EqualFold(strings.TrimSpace(msg.Channel), redisChannelCluster) { - log.Warn("failed to apply cluster update from home control center, ignoring") - } else { - log.Warn("failed to apply config update from home control center, ignoring") + if errApply := c.handleSubscriptionPayload(receiveCtx, msg.Channel, msg.Payload, onConfig); errApply != nil { + if receiveCtx.Err() != nil { + return endReceiveFailure(errApply) } + if errors.Is(errApply, ErrRebalanceRequested) { + cancelReceive(ErrRebalanceRequested) + return c.endConfigSubscriberLifetimeWithSubscription(errApply, pubsub, "planned rebalance") + } + if strings.EqualFold(strings.TrimSpace(msg.Channel), redisChannelCluster) || strings.EqualFold(strings.TrimSpace(msg.Channel), redisChannelRebalance) { + return endReceiveFailure(errApply) + } + log.Warn("failed to apply config update from home control center, ignoring") } case *redis.Pong: c.resetReconnectFailures() @@ -2040,6 +2160,7 @@ func receiveSubscriptionACKs(ctx context.Context, pubsub *redis.PubSub, receiveT func (c *Client) endConfigSubscriberLifetime(err error) error { c.heartbeatOK.Store(false) + c.failClusterRebalance() if !c.managedLifetime() { c.Close() } @@ -2048,6 +2169,9 @@ func (c *Client) endConfigSubscriberLifetime(err error) error { func (c *Client) endConfigSubscriberLifetimeWithSubscription(err error, subscription subscriptionCloser, reason string) error { c.heartbeatOK.Store(false) + if !errors.Is(err, ErrRebalanceRequested) { + c.failClusterRebalance() + } if subscription != nil { if errClose := subscription.Close(); errClose != nil { log.WithError(errClose).Debugf("Home subscription close after %s", reason) diff --git a/internal/home/client_cancellation_test.go b/internal/home/client_cancellation_test.go new file mode 100644 index 0000000000..a146013959 --- /dev/null +++ b/internal/home/client_cancellation_test.go @@ -0,0 +1,121 @@ +package home + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strings" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" +) + +func TestManagedSubscriberCancellationPreservesFinalRelease(t *testing.T) { + for _, existingReleaseClient := range []bool{false, true} { + t.Run(fmt.Sprintf("existingReleaseClient=%v", existingReleaseClient), func(t *testing.T) { + client, commands := newRedisCommandTestClient(t, func(args []string) string { + switch strings.ToUpper(args[0]) { + 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" + case "LPUSH": + return ":1\r\n" + default: + return balanceTestRedisReply(args) + } + }) + client.SetManagedLifetime(true) + client.testOperationTimeout = time.Second + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + frame := ConcurrencyReleaseFrame{CredentialID: "credential-a", Model: "test-model", ReleaseSeq: 1} + var errInitial error + ready := false + errRun := client.RunConfigSubscriberLifetime(ctx, func([]byte) error { + return client.SetLifecycleConfig(config.CredentialConcurrencyConfig{LifecycleConfigRevision: 1}) + }, func() { + ready = client.HeartbeatOK() + if existingReleaseClient { + errInitial = client.PushConcurrencyRelease(ctx, frame) + } + cancel() + }) + if !ready || !errors.Is(errRun, context.Canceled) || errInitial != nil { + t.Fatalf("subscriber cancellation: ready=%v error=%v initial release=%v", ready, errRun, errInitial) + } + if client.HeartbeatOK() || client.AmbiguousDispatch() { + t.Fatal("graceful cancellation retained a heartbeat or introduced ambiguity") + } + if _, errDispatch := client.RPopAuth(context.Background(), "late-model", "", nil, 1); !errors.Is(errDispatch, ErrDispatchFenced) { + t.Fatalf("graceful cancellation reopened dispatch: %v", errDispatch) + } + frame.ReleaseSeq = 2 + if errRelease := client.PushConcurrencyRelease(context.WithoutCancel(ctx), frame); errRelease != nil { + t.Fatalf("final release after normal cancellation: %v", errRelease) + } + var releases []ConcurrencyReleaseFrame + for _, args := range commands.All() { + if len(args) == 3 && strings.EqualFold(args[0], "LPUSH") && args[1] == redisKeyConcurrencyRelease { + var sent ConcurrencyReleaseFrame + if errJSON := json.Unmarshal([]byte(args[2]), &sent); errJSON != nil { + t.Fatal(errJSON) + } + releases = append(releases, sent) + } + } + wantCount := 1 + if existingReleaseClient { + wantCount++ + } + if len(releases) != wantCount || releases[len(releases)-1] != frame { + t.Fatalf("release frames = %#v", releases) + } + client.Close() + if errRelease := client.PushConcurrencyRelease(context.Background(), frame); !errors.Is(errRelease, ErrDispatchFenced) { + t.Fatalf("release after Close = %v, want fenced", errRelease) + } + }) + } +} + +func TestManagedSubscriberReceiveFailureKeepsReleaseFenced(t *testing.T) { + client, commands := newRedisCommandTestClient(t, func(args []string) string { + switch strings.ToUpper(args[0]) { + case "GET": + return "$2\r\n{}\r\n" + case "SUBSCRIBE": + // A malformed RESP frame fails reception after the membership ACK. + return "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:1\r\n?invalid\r\n" + default: + return balanceTestRedisReply(args) + } + }) + client.SetManagedLifetime(true) + client.testOperationTimeout = time.Second + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + errRun := client.RunConfigSubscriberLifetime(ctx, func([]byte) error { + return client.SetLifecycleConfig(config.CredentialConcurrencyConfig{LifecycleConfigRevision: 1}) + }, nil) + if errRun == nil || errors.Is(errRun, context.Canceled) { + t.Fatalf("receive failure = %v", errRun) + } + // Later graceful cancellation must not undo a fence installed by a receive failure. + cancel() + if client.HeartbeatOK() || !client.dispatchFenced.Load() { + t.Fatal("abnormal reception did not fence the lifetime") + } + frame := ConcurrencyReleaseFrame{CredentialID: "credential-a", Model: "test-model", ReleaseSeq: 1} + if errRelease := client.PushConcurrencyRelease(context.WithoutCancel(ctx), frame); !errors.Is(errRelease, ErrDispatchFenced) { + t.Fatalf("release after receive failure = %v, want fenced", errRelease) + } + if _, errDispatch := client.RPopAuth(context.Background(), "test-model", "", nil, 1); !errors.Is(errDispatch, ErrDispatchFenced) { + t.Fatalf("dispatch after receive failure = %v, want fenced", errDispatch) + } + if commands.CountCommandKey("LPUSH", redisKeyConcurrencyRelease) != 0 { + t.Fatal("fenced lifetime sent a release") + } +} diff --git a/internal/home/client_fence_test.go b/internal/home/client_fence_test.go new file mode 100644 index 0000000000..65e48e7fd7 --- /dev/null +++ b/internal/home/client_fence_test.go @@ -0,0 +1,119 @@ +package home + +import ( + "context" + "errors" + "strings" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executionregistry" +) + +func TestFenceDispatchRejectsRequestPausedBeforeRegistration(t *testing.T) { + client, commands := newRedisCommandTestClient(t, balanceTestRedisReply) + client.heartbeatOK.Store(true) + registry := executionregistry.New() + checkedHeartbeat := make(chan struct{}) + resume := make(chan struct{}) + result := make(chan error, 1) + go func() { + if !client.HeartbeatOK() { + result <- errors.New("initial heartbeat was unavailable") + return + } + close(checkedHeartbeat) + <-resume + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + result <- errBegin + return + } + defer pending.End() + _, errDispatch := client.RPopAuth(context.Background(), "test-model", "", nil, 1) + result <- errDispatch + }() + select { + case <-checkedHeartbeat: + case <-time.After(time.Second): + t.Fatal("request did not reach the heartbeat barrier") + } + client.FenceDispatch() + if errPending := registry.WaitPending(context.Background()); errPending != nil { + t.Fatal(errPending) + } + if !client.HeartbeatOK() || client.AmbiguousDispatch() { + t.Fatal("admission fencing changed heartbeat or ambiguity") + } + close(resume) + select { + case errDispatch := <-result: + if !errors.Is(errDispatch, ErrDispatchFenced) { + t.Fatalf("late dispatch = %v, want fenced", errDispatch) + } + case <-time.After(time.Second): + t.Fatal("late dispatch did not finish") + } + if len(commands.All()) != 0 { + t.Fatalf("fenced dispatch sent commands: %#v", commands.All()) + } +} + +func TestFenceDispatchPreservesIssuedRPopTransport(t *testing.T) { + issued := make(chan struct{}) + reply := make(chan struct{}, 1) + t.Cleanup(func() { close(reply) }) + client, commands := newRedisCommandTestClient(t, func(args []string) string { + if strings.EqualFold(args[0], "RPOP") { + close(issued) + <-reply + return "$2\r\n{}\r\n" + } + return balanceTestRedisReply(args) + }) + client.testOperationTimeout = time.Second + client.closeBootstrapPools() + if errEnsure := client.ensureClients(); errEnsure != nil { + t.Fatal(errEnsure) + } + client.heartbeatOK.Store(true) + commandPool, subscriptionPool := client.cmd, client.sub + result := make(chan error, 1) + go func() { + raw, errDispatch := client.RPopAuth(context.Background(), "test-model", "", nil, 1) + if errDispatch == nil && string(raw) != "{}" { + errDispatch = errors.New("issued response was lost") + } + result <- errDispatch + }() + select { + case <-issued: + case <-time.After(time.Second): + t.Fatal("RPOP did not reach Home") + } + client.FenceDispatch() + client.FenceDispatch() + if client.cmd != commandPool || client.sub != subscriptionPool || !client.HeartbeatOK() || client.AmbiguousDispatch() { + t.Fatal("fencing changed transport, heartbeat, or ambiguity") + } + select { + case errDispatch := <-result: + t.Fatalf("fencing interrupted the issued RPOP: %v", errDispatch) + default: + } + if _, errLate := client.RPopAuth(context.Background(), "late-model", "", nil, 1); !errors.Is(errLate, ErrDispatchFenced) { + t.Fatalf("late RPOP = %v, want fenced", errLate) + } + reply <- struct{}{} + select { + case errDispatch := <-result: + if errDispatch != nil { + t.Fatal(errDispatch) + } + case <-time.After(time.Second): + t.Fatal("issued RPOP did not finish after its response arrived") + } + if len(findRedisCommand(commands.All(), "RPOP")) == 0 { + t.Fatal("issued RPOP was not recorded") + } +} diff --git a/internal/home/client_release_lifecycle_test.go b/internal/home/client_release_lifecycle_test.go new file mode 100644 index 0000000000..de93c84de2 --- /dev/null +++ b/internal/home/client_release_lifecycle_test.go @@ -0,0 +1,182 @@ +package home + +import ( + "context" + "encoding/json" + "errors" + "strings" + "testing" + "time" +) + +func TestDispatchFencePreservesConcurrencyReleaseUntilClose(t *testing.T) { + client, commands := newRedisCommandTestClient(t, func(args []string) string { + if strings.EqualFold(args[0], "LPUSH") { + return ":1\r\n" + } + return balanceTestRedisReply(args) + }) + client.FenceDispatch() + frame := ConcurrencyReleaseFrame{CredentialID: "credential-a", Model: "test-model", ReleaseSeq: 1} + if errRelease := client.PushConcurrencyRelease(context.Background(), frame); errRelease != nil { + t.Fatalf("admission fence blocked final release: %v", errRelease) + } + if _, errDispatch := client.RPopAuth(context.Background(), "test-model", "", nil, 1); !errors.Is(errDispatch, ErrDispatchFenced) { + t.Fatalf("fenced dispatch = %v", errDispatch) + } + client.Close() + if errRelease := client.PushConcurrencyRelease(context.Background(), frame); !errors.Is(errRelease, ErrDispatchFenced) { + t.Fatalf("closed lifetime accepted release: %v", errRelease) + } + if commands.CountCommandKey("LPUSH", redisKeyConcurrencyRelease) != 1 { + t.Fatal("closed lifetime sent another release") + } +} + +func TestPreparedMigrationKeepsReleaseAtSourceUntilTakeover(t *testing.T) { + target, targetCommands := newRedisCommandTestClient(t, balanceTestRedisReply) + source, sourceCommands := newRedisCommandTestClient(t, func(args []string) string { + switch strings.ToUpper(args[0]) { + case "CLUSTER": + return "$11\r\n{\"ok\":true}\r\n" + case "LPUSH": + return ":1\r\n" + default: + return balanceTestRedisReply(args) + } + }) + enableBalanceTestClient(source) + if errProbe := source.rebuildCommandPoolAndProbe(context.Background()); errProbe != nil { + t.Fatal(errProbe) + } + payload, errJSON := json.Marshal(rebalanceInstruction{ID: "release-plan", Term: 1, Target: clusterNode{IP: target.homeCfg.Host, Port: target.homeCfg.Port}}) + if errJSON != nil { + t.Fatal(errJSON) + } + if errExecute := source.executeClusterRebalance(context.Background(), payload); !errors.Is(errExecute, ErrRebalanceRequested) { + t.Fatal(errExecute) + } + source.FenceDispatch() + source.heartbeatOK.Store(false) + next := source.NewLifetime() + t.Cleanup(next.Close) + frame := ConcurrencyReleaseFrame{CredentialID: "credential-a", Model: "test-model", ReleaseSeq: 1} + if errRelease := source.PushConcurrencyRelease(context.Background(), frame); errRelease != nil { + t.Fatalf("prepared source lost release ownership: %v", errRelease) + } + if errRelease := next.PushConcurrencyRelease(context.Background(), frame); !errors.Is(errRelease, ErrNotConnected) { + t.Fatalf("unacknowledged target accepted release: %v", errRelease) + } + // Cancellation before the subscriber reports the plan also rolls back balance state. + source.failClusterRebalance() + frame.ReleaseSeq++ + if errRelease := source.PushConcurrencyRelease(context.Background(), frame); errRelease != nil { + t.Fatalf("canceling the prepared plan lost source release ownership: %v", errRelease) + } + if sourceCommands.CountCommandKey("LPUSH", redisKeyConcurrencyRelease) != 2 || targetCommands.CountCommandKey("LPUSH", redisKeyConcurrencyRelease) != 0 { + t.Fatal("release was not restricted to the source membership") + } + source.AbortAmbiguousDispatch() + if errRelease := source.PushConcurrencyRelease(context.Background(), frame); !errors.Is(errRelease, ErrDispatchFenced) { + t.Fatalf("ambiguous source accepted release: %v", errRelease) + } +} + +func TestSourceReleaseOwnershipDoesNotSurviveAddressSwitch(t *testing.T) { + target, commands := newRedisCommandTestClient(t, balanceTestRedisReply) + source, _ := newRedisCommandTestClient(t, balanceTestRedisReply) + if errProbe := source.rebuildCommandPoolAndProbe(context.Background()); errProbe != nil { + t.Fatal(errProbe) + } + source.markMembershipTakeoverEligible() + source.mu.Lock() + switched := source.switchToNodeLocked(clusterNode{IP: target.homeCfg.Host, Port: target.homeCfg.Port}) + source.mu.Unlock() + if !switched { + t.Fatal("source address was not switched") + } + frame := ConcurrencyReleaseFrame{CredentialID: "credential-a", Model: "test-model", ReleaseSeq: 1} + if errRelease := source.PushConcurrencyRelease(context.Background(), frame); !errors.Is(errRelease, ErrNotConnected) { + t.Fatalf("switched address retained source release ownership: %v", errRelease) + } + if len(commands.All()) != 0 { + t.Fatal("unacknowledged target received release traffic") + } +} + +func TestPluginStatusCancellationPreservesDispatchAndReleaseTransports(t *testing.T) { + statusStarted := make(chan struct{}) + dispatchStarted := make(chan struct{}) + statusReply := make(chan struct{}) + dispatchReply := make(chan struct{}, 1) + t.Cleanup(func() { close(statusReply); close(dispatchReply) }) + client, _ := newRedisCommandTestClient(t, func(args []string) string { + switch strings.ToUpper(args[0]) { + case "RPUSH": + close(statusStarted) + <-statusReply + return ":1\r\n" + case "RPOP": + close(dispatchStarted) + <-dispatchReply + return "$2\r\n{}\r\n" + case "LPUSH": + return ":1\r\n" + default: + return balanceTestRedisReply(args) + } + }) + client.testOperationTimeout = 5 * time.Second + client.closeBootstrapPools() + if errEnsure := client.ensureClients(); errEnsure != nil { + t.Fatal(errEnsure) + } + dispatchResult := make(chan error, 1) + go func() { + _, errDispatch := client.RPopAuth(context.Background(), "test-model", "", nil, 1) + dispatchResult <- errDispatch + }() + select { + case <-dispatchStarted: + case <-time.After(time.Second): + t.Fatal("dispatch did not start") + } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + statusResult := make(chan error, 1) + go func() { statusResult <- client.RPushPluginStatus(ctx, []byte(`{"ok":true}`)) }() + select { + case <-statusStarted: + case <-time.After(time.Second): + t.Fatal("plugin status did not start") + } + cancel() + select { + case errStatus := <-statusResult: + if !errors.Is(errStatus, context.Canceled) { + t.Fatalf("plugin cancellation error = %v", errStatus) + } + case <-time.After(time.Second): + t.Fatal("plugin status transport was not canceled") + } + if client.dispatchFenced.Load() || client.releaseFenced.Load() || client.AmbiguousDispatch() { + t.Fatal("plugin cancellation ended the Home lifetime") + } + select { + case errDispatch := <-dispatchResult: + t.Fatalf("plugin cancellation interrupted the dispatch: %v", errDispatch) + default: + } + if errRelease := client.PushConcurrencyRelease(context.Background(), ConcurrencyReleaseFrame{CredentialID: "credential-a", Model: "test-model", ReleaseSeq: 1}); errRelease != nil { + t.Fatalf("plugin cancellation blocked release: %v", errRelease) + } + dispatchReply <- struct{}{} + select { + case errDispatch := <-dispatchResult: + if errDispatch != nil { + t.Fatal(errDispatch) + } + case <-time.After(time.Second): + t.Fatal("issued dispatch did not finish") + } +} diff --git a/internal/home/cluster_balance.go b/internal/home/cluster_balance.go new file mode 100644 index 0000000000..deb1f09b09 --- /dev/null +++ b/internal/home/cluster_balance.go @@ -0,0 +1,249 @@ +package home + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net" + "strconv" + "strings" + + "github.com/redis/go-redis/v9" + log "github.com/sirupsen/logrus" +) + +const redisChannelRebalance = "rebalance" + +// ErrRebalanceRequested asks the supervisor to reconnect without draining safe active executions. +var ErrRebalanceRequested = errors.New("home cluster rebalance requested") + +type rebalanceInstruction struct { + ID string `json:"id"` + Term int64 `json:"term"` + Target clusterNode `json:"target"` +} + +type clusterBalanceState struct { + moving bool + returning bool + target clusterNode + origin clusterNode + request rebalanceInstruction + result string + registered bool + unsupported bool +} + +func (c *Client) canExecuteRebalanceLocked() bool { + return c.managed && c.clusterDiscoveryEnabledLocked() && !c.legacyMembership && + c.lifecycle.LifecycleConfigRevision > 0 && c.heartbeatOK.Load() && + !c.dispatchFenced.Load() && !c.ambiguousDispatch.Load() && + recoveryState(c.recoveryState.Load()) == recoveryStateStable && !c.balance.moving +} + +// rebalanceCommand owns a cancellable transport independent of issued dispatches. +func (c *Client) rebalanceCommand(ctx context.Context, address string, args ...any) (string, error) { + if errContext := ctx.Err(); errContext != nil { + return "", errContext + } + c.mu.Lock() + if address == "" { + address, _ = c.addrLocked() + } + options, errOptions := c.redisOptionsLocked(address) + c.mu.Unlock() + if errOptions != nil { + return "", errOptions + } + options.ReadTimeout = -1 + options.WriteTimeout = -1 + options.Dialer = redis.NewDialer(options) + command := redis.NewClient(options) + stopCancel := context.AfterFunc(ctx, func() { _ = command.Close() }) + defer func() { + stopCancel() + if errClose := command.Close(); errClose != nil && ctx.Err() == nil { + log.Debugf("Home rebalance transport close failed: %v", errClose) + } + }() + raw, errCommand := command.Do(ctx, args...).Text() + if errCommand != nil { + return "", fmt.Errorf("Home rebalance %v: %w", args[0], errCommand) + } + return raw, nil +} + +func (c *Client) syncClusterBalance(ctx context.Context) error { + c.mu.Lock() + eligible := c.canExecuteRebalanceLocked() && !c.balance.unsupported + registered, request, result := c.balance.registered, c.balance.request, c.balance.result + c.mu.Unlock() + if !eligible { + return nil + } + if !registered { + raw, errEnable := c.rebalanceCommand(ctx, "", "CLUSTER", "BALANCE", "ENABLE") + if errEnable != nil { + message := strings.ToLower(errEnable.Error()) + if isHomeCommandUnsupported(errEnable) || strings.Contains(message, "unsupported subcommand") || strings.Contains(message, "rebalance unsupported") { + // Legacy Home has no capability code; only an explicit unsupported reply disables balancing. + c.mu.Lock() + c.balance.unsupported = true + c.mu.Unlock() + return nil + } + return fmt.Errorf("register Home rebalance: %w", errEnable) + } + var response struct { + OK bool `json:"ok"` + } + if errJSON := json.Unmarshal([]byte(raw), &response); errJSON != nil { + return fmt.Errorf("decode Home rebalance ENABLE response: %w", errJSON) + } + if !response.OK { + return nil + } + c.mu.Lock() + c.balance.registered = true + c.mu.Unlock() + } + if result == "" { + return nil + } + + raw, errResult := c.rebalanceCommand(ctx, "", "CLUSTER", "BALANCE", "RESULT", request.ID, result) + if errResult != nil { + return fmt.Errorf("report Home rebalance result: %w", errResult) + } + var response struct { + OK bool `json:"ok"` + } + if errJSON := json.Unmarshal([]byte(raw), &response); errJSON != nil { + return fmt.Errorf("decode Home rebalance RESULT response: %w", errJSON) + } + if !response.OK { + return nil + } + c.mu.Lock() + if c.balance.request.ID == request.ID && c.balance.result == result { + c.balance.result = "" + c.balance.request = rebalanceInstruction{} + } + c.mu.Unlock() + return nil +} + +// executeClusterRebalance claims an exact persisted plan before touching any transport. +// This lifetime keeps its original address; only the supervisor's replacement uses the target. +func (c *Client) executeClusterRebalance(ctx context.Context, payload []byte) error { + var request rebalanceInstruction + if errJSON := json.Unmarshal(payload, &request); errJSON != nil { + return fmt.Errorf("decode Home rebalance instruction: %w", errJSON) + } + request.Target.IP = strings.TrimSpace(request.Target.IP) + if request.ID == "" || request.Term <= 0 || request.Target.IP == "" || request.Target.Port <= 0 || request.Target.Port > 65535 { + return nil + } + c.mu.Lock() + eligible := c.canExecuteRebalanceLocked() && c.balance.registered && c.balance.result == "" && + (request.Target.IP != strings.TrimSpace(c.homeCfg.Host) || request.Target.Port != c.homeCfg.Port) + c.mu.Unlock() + if !eligible || ctx.Err() != nil { + return nil + } + raw, errClaim := c.rebalanceCommand(ctx, "", "CLUSTER", "BALANCE", "CLAIM", request.ID, request.Term, request.Target.IP, request.Target.Port) + if errClaim != nil { + return errClaim + } + var response struct { + OK bool `json:"ok"` + } + if errJSON := json.Unmarshal([]byte(raw), &response); errJSON != nil { + return fmt.Errorf("decode Home rebalance CLAIM response: %w", errJSON) + } + if !response.OK { + return nil + } + c.mu.Lock() + c.balance.request = request + c.balance.result = "failed" + c.mu.Unlock() + if _, errPing := c.rebalanceCommand(ctx, clusterNodeAddress(request.Target), "PING"); errPing != nil { + if errReport := c.syncClusterBalance(ctx); errReport != nil { + return errors.Join(errPing, errReport) + } + c.mu.Lock() + recovered := ctx.Err() == nil && c.canExecuteRebalanceLocked() && c.balance.result == "" + c.mu.Unlock() + if !recovered { + return errPing + } + // No switch occurred. An acknowledged failure and a healthy source safely end this plan, not the subscription. + log.Warnf("Home rebalance preflight failed; failure reported, keeping source: plan=%q target=%q error=%v", request.ID, clusterNodeAddress(request.Target), errPing) + return nil + } + c.mu.Lock() + if ctx.Err() != nil || !c.canExecuteRebalanceLocked() { + c.mu.Unlock() + return c.syncClusterBalance(ctx) + } + c.balance.origin = clusterNode{IP: strings.TrimSpace(c.homeCfg.Host), Port: c.homeCfg.Port} + c.balance.target = request.Target + c.balance.result = "" + c.balance.moving = true + c.markMembershipTakeoverEligible() + origin := c.balance.origin + c.mu.Unlock() + log.Infof("executing Home rebalance plan: plan=%q from=%q to=%q", request.ID, clusterNodeAddress(origin), clusterNodeAddress(request.Target)) + return ErrRebalanceRequested +} + +func (c *Client) finishClusterRebalance(ctx context.Context) error { + c.mu.Lock() + if c.balance.moving { + c.balance.moving = false + c.balance.result = "success" + } + c.balance.returning = false + c.mu.Unlock() + return c.syncClusterBalance(ctx) +} + +func (c *Client) failClusterRebalance() { + c.mu.Lock() + defer c.mu.Unlock() + if !c.balance.moving { + return + } + c.balance.moving = false + c.balance.returning = true + c.balance.result = "failed" + c.switchToNodeLocked(c.balance.origin) +} + +// chooseClusterNodeLocked only chooses bootstrap/fault-recovery targets, never active migrations. +func (c *Client) chooseClusterNodeLocked(nodes []clusterNode, excludeCurrent bool, draw float64) (clusterNode, bool) { + var choices []clusterNode + for _, node := range nodes { + if node.IP == "" || node.Port <= 0 || node.Port > 65535 { + continue + } + if excludeCurrent && node.IP == strings.TrimSpace(c.homeCfg.Host) && node.Port == c.homeCfg.Port { + continue + } + if len(choices) == 0 || node.ClientCount < choices[0].ClientCount { + choices = []clusterNode{node} + } else if node.ClientCount == choices[0].ClientCount { + choices = append(choices, node) + } + } + if len(choices) == 0 { + return clusterNode{}, false + } + return choices[int(draw*float64(len(choices)))], true +} + +func clusterNodeAddress(node clusterNode) string { + return net.JoinHostPort(node.IP, strconv.Itoa(node.Port)) +} diff --git a/internal/home/cluster_balance_recovery_test.go b/internal/home/cluster_balance_recovery_test.go new file mode 100644 index 0000000000..b34ae6fc3e --- /dev/null +++ b/internal/home/cluster_balance_recovery_test.go @@ -0,0 +1,386 @@ +package home + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "reflect" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/redis/go-redis/v9" + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" +) + +func TestClusterBalanceTargetPreflightFailureRecovery(t *testing.T) { + for _, report := range []struct { + name string + reply string + }{ + {name: "acknowledged", reply: "$11\r\n{\"ok\":true}\r\n"}, + {name: "RPC failure", reply: "-TEST_REPORT_FAILURE\r\n"}, + {name: "malformed", reply: "$1\r\n{\r\n"}, + {name: "rejected", reply: "$12\r\n{\"ok\":false}\r\n"}, + } { + t.Run(report.name, func(t *testing.T) { + target, _ := newRedisCommandTestClient(t, func(args []string) string { + if strings.EqualFold(args[0], "PING") { + return "-TEST_TARGET_UNAVAILABLE\r\n" + } + return balanceTestRedisReply(args) + }) + instruction, errMarshal := json.Marshal(rebalanceInstruction{ + ID: "failed-plan", Term: 1, Target: clusterNode{IP: target.homeCfg.Host, Port: target.homeCfg.Port}, + }) + if errMarshal != nil { + t.Fatal(errMarshal) + } + const marker = `{"after_rebalance":true}` + client, commands := newRedisCommandTestClient(t, func(args []string) string { + switch strings.ToUpper(args[0]) { + case "CLUSTER": + if args[1] == "BALANCE" { + if args[2] == "RESULT" { + return report.reply + } + return "$11\r\n{\"ok\":true}\r\n" + } + return "$22\r\n{\"ok\":true,\"nodes\":[]}\r\n" + case "GET", "RPOP": + 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" + + fmt.Sprintf("*3\r\n$7\r\nmessage\r\n$9\r\nrebalance\r\n$%d\r\n%s\r\n", len(instruction), instruction) + + fmt.Sprintf("*3\r\n$7\r\nmessage\r\n$6\r\nconfig\r\n$%d\r\n%s\r\n", len(marker), marker) + default: + return balanceTestRedisReply(args) + } + }) + enableBalanceTestClient(client) + client.testOperationTimeout = time.Second + client.balance.registered = false + origin, _ := client.addr() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + var sourceCommand, sourceSubscription *redis.Client + var dispatchErr error + var dispatched []byte + var continued, healthy, samePools bool + errRun := client.RunConfigSubscriberLifetime(ctx, func(payload []byte) error { + if string(payload) == marker { + continued = true + healthy = client.HeartbeatOK() && !client.dispatchFenced.Load() + client.mu.Lock() + samePools = client.cmd == sourceCommand && client.sub == sourceSubscription + client.mu.Unlock() + dispatched, dispatchErr = client.RPopAuth(ctx, "test-model", "", nil, 1) + cancel() + } + return nil + }, func() { + client.mu.Lock() + sourceCommand, sourceSubscription = client.cmd, client.sub + client.mu.Unlock() + }) + var result []string + for _, args := range commands.All() { + if len(args) >= 3 && args[2] == "RESULT" { + result = args + } + } + if !reflect.DeepEqual(result, []string{"CLUSTER", "BALANCE", "RESULT", "failed-plan", "failed"}) { + t.Fatalf("failure report = %#v", result) + } + if address, _ := client.addr(); address != origin || client.balance.moving { + t.Fatal("failed preflight switched away from the source") + } + if report.name == "acknowledged" { + if !continued || !healthy || !samePools || !errors.Is(errRun, context.Canceled) { + t.Fatalf("healthy source was not retained: continued=%v healthy=%v samePools=%v error=%v", continued, healthy, samePools, errRun) + } + if dispatchErr != nil || string(dispatched) != "{}" || len(findRedisCommand(commands.All(), "RPOP")) == 0 { + t.Fatalf("dispatch after failed preflight = %q, %v", dispatched, dispatchErr) + } + if client.balance.result != "" || client.balance.request.ID != "" { + t.Fatal("acknowledged failure report remained pending") + } + return + } + var targetError redis.Error + if continued || !errors.As(errRun, &targetError) || targetError.Error() != "TEST_TARGET_UNAVAILABLE" { + t.Fatalf("unacknowledged failure was swallowed: continued=%v error=%v", continued, errRun) + } + if client.balance.result != "failed" || client.balance.request.ID != "failed-plan" { + t.Fatal("unacknowledged failure report was discarded") + } + if report.name == "malformed" { + var syntaxError *json.SyntaxError + if !errors.As(errRun, &syntaxError) { + t.Fatalf("report decoding failure was lost: %v", errRun) + } + } + if report.name == "RPC failure" { + joined, ok := errRun.(interface{ Unwrap() []error }) + if !ok || len(joined.Unwrap()) != 2 { + t.Fatalf("report RPC failure was lost: %v", errRun) + } + var reportError redis.Error + if !errors.As(joined.Unwrap()[1], &reportError) || reportError.Error() != "TEST_REPORT_FAILURE" { + t.Fatalf("report RPC failure code was lost: %v", errRun) + } + } + }) + } +} + +func TestClusterBalanceRollbackRecoveryAcrossLifetimes(t *testing.T) { + for _, failure := range []string{"GET", "SUBSCRIBE", "PING"} { + for _, sourceRecovers := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/source_recovers=%v", failure, sourceRecovers), func(t *testing.T) { + healthy, healthyCommands := newRedisCommandTestClient(t, func(args []string) string { + switch strings.ToUpper(args[0]) { + case "CLUSTER": + if args[1] == "BALANCE" { + 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": + return "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:1\r\n" + default: + return balanceTestRedisReply(args) + } + }) + target, _ := newRedisCommandTestClient(t, func(args []string) string { + switch strings.ToUpper(args[0]) { + case "CLUSTER": + return "$22\r\n{\"ok\":true,\"nodes\":[]}\r\n" + case "GET": + return "-TEST_TARGET_CONFIG_UNAVAILABLE\r\n" + default: + return balanceTestRedisReply(args) + } + }) + nodes, errNodes := json.Marshal(clusterNodesEnvelope{OK: true, Nodes: []clusterNode{ + {IP: healthy.homeCfg.Host, Port: healthy.homeCfg.Port}, + }}) + if errNodes != nil { + t.Fatal(errNodes) + } + var sourceRecovered atomic.Bool + source, sourceCommands := newRedisCommandTestClient(t, func(args []string) string { + if strings.EqualFold(args[0], failure) && !sourceRecovered.Load() { + return "-TEST_SOURCE_UNAVAILABLE\r\n" + } + switch strings.ToUpper(args[0]) { + case "CLUSTER": + if args[1] == "BALANCE" { + return "$11\r\n{\"ok\":true}\r\n" + } + return fmt.Sprintf("$%d\r\n%s\r\n", len(nodes), nodes) + 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(source) + source.testOperationTimeout = time.Second + instruction, errInstruction := json.Marshal(rebalanceInstruction{ + ID: "failed-plan", Term: 1, Target: clusterNode{IP: target.homeCfg.Host, Port: target.homeCfg.Port}, + }) + if errInstruction != nil { + t.Fatal(errInstruction) + } + if errExecute := source.executeClusterRebalance(context.Background(), instruction); !errors.Is(errExecute, ErrRebalanceRequested) { + t.Fatalf("prepare migration = %v", errExecute) + } + source.Close() + client := source.NewLifetime() + defer func() { client.Close() }() + client.SetManagedLifetime(true) + apply := func([]byte) error { + return client.SetLifecycleConfig(config.CredentialConcurrencyConfig{LifecycleConfigRevision: 1}) + } + errTarget := client.RunConfigSubscriberLifetime(context.Background(), apply, nil) + var targetError redis.Error + if !errors.As(errTarget, &targetError) || targetError.Error() != "TEST_TARGET_CONFIG_UNAVAILABLE" { + t.Fatalf("target failure = %v", errTarget) + } + if client.homeCfg.Port != source.homeCfg.Port || !client.balance.returning || client.reconnectFailures != 1 { + t.Fatalf("failed target did not enter rollback: port=%d returning=%v failures=%d", client.homeCfg.Port, client.balance.returning, client.reconnectFailures) + } + + attempts := homeReconnectFailoverThreshold - 1 + if sourceRecovers { + attempts-- + } + for attempt := 1; attempt <= attempts; attempt++ { + client.Close() + client = client.NewLifetime() + client.SetManagedLifetime(true) + errRun := client.RunConfigSubscriberLifetime(context.Background(), apply, nil) + var sourceError redis.Error + if !errors.As(errRun, &sourceError) || sourceError.Error() != "TEST_SOURCE_UNAVAILABLE" { + t.Fatalf("source attempt %d lost original failure: %v", attempt, errRun) + } + wantPort, wantFailures := source.homeCfg.Port, attempt+1 + if wantFailures == homeReconnectFailoverThreshold { + wantPort, wantFailures = healthy.homeCfg.Port, 0 + } + if client.homeCfg.Port != wantPort || client.reconnectFailures != wantFailures { + t.Fatalf("source attempt %d: port=%d failures=%d, want port=%d failures=%d", attempt, client.homeCfg.Port, client.reconnectFailures, wantPort, wantFailures) + } + if client.balance.moving || !client.balance.returning || client.balance.result != "failed" || client.balance.request.ID != "failed-plan" { + t.Fatal("rollback lost its pending failed result") + } + } + if count := sourceCommands.CountCommandKey("CLUSTER", "NODES"); count != attempts { + t.Fatalf("source discovery count = %d, want %d", count, attempts) + } + + sourceRecovered.Store(sourceRecovers) + client.Close() + client = client.NewLifetime() + client.SetManagedLifetime(true) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ready := false + errRecovered := client.RunConfigSubscriberLifetime(ctx, apply, func() { + ready = client.HeartbeatOK() && client.reconnectFailures == 0 && + recoveryState(client.recoveryState.Load()) == recoveryStateStable + cancel() + }) + if !ready || !errors.Is(errRecovered, context.Canceled) { + t.Fatalf("recovery did not complete: ready=%v error=%v", ready, errRecovered) + } + if client.balance.returning || client.balance.result != "" || client.balance.request.ID != "" { + t.Fatal("recovery did not finish the failed migration") + } + recoveredCommands, wantPort := healthyCommands, healthy.homeCfg.Port + if sourceRecovers { + recoveredCommands, wantPort = sourceCommands, source.homeCfg.Port + } + if client.homeCfg.Port != wantPort { + t.Fatalf("recovered port = %d, want %d", client.homeCfg.Port, wantPort) + } + var subscription, result []string + for _, args := range recoveredCommands.All() { + if strings.EqualFold(args[0], "SUBSCRIBE") { + subscription = args + } + if len(args) >= 3 && args[2] == "RESULT" { + result = args + } + } + wantSubscription := []string{"subscribe", redisChannelConfig, "1", "takeover", source.MembershipInstanceID()} + if !reflect.DeepEqual(subscription, wantSubscription) { + t.Fatalf("recovery subscription = %#v, want %#v", subscription, wantSubscription) + } + if !reflect.DeepEqual(result, []string{"CLUSTER", "BALANCE", "RESULT", "failed-plan", "failed"}) { + t.Fatalf("failure report = %#v", result) + } + }) + } + } +} + +func TestClusterBalancePostACKFailurePreservesTakeover(t *testing.T) { + for _, failure := range []struct { + operation string + malformed bool + }{ + {operation: "ENABLE"}, + {operation: "ENABLE", malformed: true}, + {operation: "RESULT"}, + {operation: "RESULT", malformed: true}, + {operation: redisChannelCluster, malformed: true}, + } { + t.Run(fmt.Sprintf("%s/malformed=%v", failure.operation, failure.malformed), func(t *testing.T) { + var recovered atomic.Bool + client, commands := newRedisCommandTestClient(t, func(args []string) string { + switch strings.ToUpper(args[0]) { + case "CLUSTER": + if args[1] == "BALANCE" { + if args[2] == failure.operation && !recovered.Load() { + if failure.malformed { + return "$1\r\n{\r\n" + } + return "-TEST_BALANCE_FAILURE\r\n" + } + 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 failure.operation == redisChannelCluster && !recovered.Load() { + reply += fmt.Sprintf("*3\r\n$7\r\nmessage\r\n$%d\r\n%s\r\n$1\r\n{\r\n", len(redisChannelCluster), redisChannelCluster) + } + return reply + default: + return balanceTestRedisReply(args) + } + }) + enableBalanceTestClient(client) + client.testOperationTimeout = time.Second + client.balance.registered = false + if failure.operation == "RESULT" { + client.balance.request.ID, client.balance.result = "previous-plan", "success" + } + errRun := client.RunConfigSubscriberLifetime(context.Background(), func([]byte) error { return nil }, nil) + if failure.malformed { + var syntaxError *json.SyntaxError + if !errors.As(errRun, &syntaxError) { + t.Fatalf("original JSON failure was lost: %v", errRun) + } + } else { + var redisError redis.Error + if !errors.As(errRun, &redisError) || redisError.Error() != "TEST_BALANCE_FAILURE" { + t.Fatalf("original RPC failure was lost: %v", errRun) + } + } + if client.HeartbeatOK() { + t.Fatal("failed subscriber retained healthy heartbeat") + } + + next := client.NewLifetime() + next.SetManagedLifetime(true) + t.Cleanup(next.Close) + if errSet := next.SetLifecycleConfig(config.CredentialConcurrencyConfig{LifecycleConfigRevision: 1}); errSet != nil { + t.Fatal(errSet) + } + want := []string{redisChannelConfig, "1", "takeover", client.MembershipInstanceID()} + if args, _ := next.subscriptionParameters(); !reflect.DeepEqual(args, want) { + t.Fatalf("replacement subscription = %#v, want %#v", args, want) + } + recovered.Store(true) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ready := false + errNext := next.RunConfigSubscriberLifetime(ctx, func([]byte) error { return nil }, func() { + ready = next.HeartbeatOK() + cancel() + }) + if !ready || !errors.Is(errNext, context.Canceled) { + t.Fatalf("replacement did not recover: ready=%v error=%v", ready, errNext) + } + var subscriptions [][]string + for _, args := range commands.All() { + if strings.EqualFold(args[0], "SUBSCRIBE") { + subscriptions = append(subscriptions, args) + } + } + if len(subscriptions) != 2 || !reflect.DeepEqual(subscriptions[1][1:], want) { + t.Fatalf("replacement wire subscriptions = %#v", subscriptions) + } + }) + } +} diff --git a/internal/home/cluster_balance_regression_test.go b/internal/home/cluster_balance_regression_test.go new file mode 100644 index 0000000000..aed91b67ac --- /dev/null +++ b/internal/home/cluster_balance_regression_test.go @@ -0,0 +1,243 @@ +package home + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "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") + } +} diff --git a/internal/home/cluster_balance_test.go b/internal/home/cluster_balance_test.go new file mode 100644 index 0000000000..c42be8ae0c --- /dev/null +++ b/internal/home/cluster_balance_test.go @@ -0,0 +1,263 @@ +package home + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strings" + "sync/atomic" + "testing" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/config" +) + +func balanceTestRedisReply(args []string) string { + if strings.EqualFold(args[0], "HELLO") { + return "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n" + } + return "+PONG\r\n" +} + +func enableBalanceTestClient(client *Client) { + client.homeCfg.DisableClusterDiscovery = false + client.managed = true + client.lifecycle = config.CredentialConcurrencyConfig{LifecycleConfigRevision: 1}.WithDefaults() + client.heartbeatOK.Store(true) + client.balance.registered = true +} + +func TestClusterBalanceSnapshotsNeverInitiateMigration(t *testing.T) { + client, commands := newRedisCommandTestClient(t, balanceTestRedisReply) + enableBalanceTestClient(client) + payload, errMarshal := json.Marshal(clusterNodesEnvelope{OK: true, Nodes: []clusterNode{ + {IP: client.homeCfg.Host, Port: client.homeCfg.Port, ClientCount: 8}, {IP: "other", Port: 8327}, + }}) + if errMarshal != nil { + t.Fatal(errMarshal) + } + for range 10 { + if errUpdate := client.handleSubscriptionPayload(context.Background(), redisChannelCluster, string(payload), nil); errUpdate != nil { + t.Fatal(errUpdate) + } + } + if client.balance.moving || len(commands.All()) != 0 { + t.Fatal("load snapshots initiated autonomous migration") + } +} + +func TestClusterBalanceExecutesOnlyClaimedPlans(t *testing.T) { + for _, mode := range []string{"success", "rejected", "target failure", "claim failure"} { + t.Run(mode, func(t *testing.T) { + target, _ := newRedisCommandTestClient(t, func(args []string) string { + if mode == "target failure" && strings.EqualFold(args[0], "PING") { + return "-ERR unavailable\r\n" + } + return balanceTestRedisReply(args) + }) + client, commands := newRedisCommandTestClient(t, func(args []string) string { + if strings.EqualFold(args[0], "CLUSTER") { + if args[2] == "CLAIM" && mode == "rejected" { + return "$12\r\n{\"ok\":false}\r\n" + } + if args[2] == "CLAIM" && mode == "claim failure" { + return "-ERR unavailable\r\n" + } + return "$11\r\n{\"ok\":true}\r\n" + } + return balanceTestRedisReply(args) + }) + enableBalanceTestClient(client) + request := rebalanceInstruction{ID: "plan-1", Term: 7, Target: clusterNode{IP: target.homeCfg.Host, Port: target.homeCfg.Port}} + payload, errMarshal := json.Marshal(request) + if errMarshal != nil { + t.Fatal(errMarshal) + } + oldCommand, oldSubscription := client.cmd, client.sub + errExecute := client.executeClusterRebalance(context.Background(), payload) + if errors.Is(errExecute, ErrRebalanceRequested) != (mode == "success") { + t.Fatalf("execute = %v", errExecute) + } + if client.cmd != oldCommand || client.sub != oldSubscription { + t.Fatal("preparation closed source pools") + } + if mode == "success" { + next := client.NewLifetime() + t.Cleanup(next.Close) + if next.homeCfg.Port != target.homeCfg.Port || client.homeCfg.Port == target.homeCfg.Port || next.balance.registered { + t.Fatal("replacement target/state not isolated") + } + if recoveryState(next.recoveryState.Load()) != recoveryStateSwitchingTakeover { + t.Fatal("missing takeover") + } + next.failClusterRebalance() + if next.homeCfg.Port != client.homeCfg.Port || !next.balance.returning || next.balance.result != "failed" { + t.Fatal("failed handoff did not roll back") + } + client.SuppressTakeover() + if recoveryState(client.NewLifetime().recoveryState.Load()) == recoveryStateSwitchingTakeover { + t.Fatal("suppressed takeover restored") + } + } + if mode == "target failure" { + found := false + for _, args := range commands.All() { + if len(args) == 5 && args[2] == "RESULT" && args[4] == "failed" { + found = true + } + } + if !found || client.balance.moving { + t.Fatal("failed probe not reported while keeping source") + } + } + }) + } +} + +func TestClusterBalanceInstructionSafetyGates(t *testing.T) { + for name, disable := range map[string]func(*Client){ + "unmanaged": func(c *Client) { c.managed = false }, + "disabled discovery": func(c *Client) { c.homeCfg.DisableClusterDiscovery = true }, + "legacy": func(c *Client) { c.legacyMembership = true }, + "no revision": func(c *Client) { c.lifecycle.LifecycleConfigRevision = 0 }, + "fenced": func(c *Client) { c.dispatchFenced.Store(true) }, + "ambiguous": func(c *Client) { c.ambiguousDispatch.Store(true) }, + "no registration": func(c *Client) { c.balance.registered = false }, + "no heartbeat": func(c *Client) { c.heartbeatOK.Store(false) }, + "recovering": func(c *Client) { c.recoveryState.Store(uint32(recoveryStateTakeoverEligible)) }, + "unreported result": func(c *Client) { c.balance.result = "failed" }, + } { + t.Run(name, func(t *testing.T) { + client, commands := newRedisCommandTestClient(t, balanceTestRedisReply) + enableBalanceTestClient(client) + disable(client) + errExecute := client.executeClusterRebalance(context.Background(), []byte(`{"id":"plan","term":1,"target":{"ip":"other","port":8327}}`)) + if errExecute != nil || len(commands.All()) != 0 { + t.Fatalf("unsafe instruction executed: %v", errExecute) + } + }) + } +} + +func TestClusterBalanceOldHomeAndResultRetry(t *testing.T) { + for _, oldHome := range []bool{true, false} { + t.Run(fmt.Sprint(oldHome), func(t *testing.T) { + var results atomic.Int32 + client, commands := newRedisCommandTestClient(t, func(args []string) string { + if strings.EqualFold(args[0], "CLUSTER") { + if oldHome { + return "-ERR cluster resp: unsupported subcommand\r\n" + } + if args[2] == "RESULT" && results.Add(1) == 1 { + return "-ERR retry later\r\n" + } + return "$11\r\n{\"ok\":true}\r\n" + } + return balanceTestRedisReply(args) + }) + enableBalanceTestClient(client) + client.balance.registered = false + client.balance.request.ID, client.balance.result = "plan", "failed" + errFirst := client.syncClusterBalance(context.Background()) + if (errFirst != nil) != !oldHome { + t.Fatalf("first sync = %v, old Home = %v", errFirst, oldHome) + } + if errSecond := client.syncClusterBalance(context.Background()); errSecond != nil { + t.Fatalf("second sync = %v", errSecond) + } + if oldHome { + if !client.balance.unsupported || commands.CountCommandKey("CLUSTER", "BALANCE") != 1 { + t.Fatal("old Home was repeatedly probed") + } + } else if results.Load() != 2 || client.balance.result != "" { + t.Fatal("result was not retried") + } + }) + } +} + +func TestRunConfigSubscriberLifetimeExecutesHomePlanAndReportsSuccess(t *testing.T) { + var directive atomic.Value + target, targetCommands := newRedisCommandTestClient(t, func(args []string) string { + switch strings.ToUpper(args[0]) { + case "CLUSTER": + if args[1] == "BALANCE" { + return "$11\r\n{\"ok\":true}\r\n" + } + return "$22\r\n{\"ok\":true,\"nodes\":[]}\r\n" + case "GET": + return "$3\r\n{}\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) + } + }) + client, _ := newRedisCommandTestClient(t, func(args []string) string { + switch strings.ToUpper(args[0]) { + case "CLUSTER": + if args[1] == "BALANCE" { + return "$11\r\n{\"ok\":true}\r\n" + } + return "$22\r\n{\"ok\":true,\"nodes\":[]}\r\n" + case "GET": + return "$3\r\n{}\n\r\n" + case "SUBSCRIBE": + payload := directive.Load().(string) + return "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:1\r\n" + fmt.Sprintf("*3\r\n$7\r\nmessage\r\n$9\r\nrebalance\r\n$%d\r\n%s\r\n", len(payload), payload) + default: + return balanceTestRedisReply(args) + } + }) + payload, errMarshal := json.Marshal(rebalanceInstruction{ID: "plan", Term: 1, Target: clusterNode{IP: target.homeCfg.Host, Port: target.homeCfg.Port}}) + if errMarshal != nil { + t.Fatal(errMarshal) + } + directive.Store(string(payload)) + enableBalanceTestClient(client) + client.balance.registered = false + apply := func([]byte) error { + return client.SetLifecycleConfig(config.CredentialConcurrencyConfig{LifecycleConfigRevision: 1}) + } + errRun := client.RunConfigSubscriberLifetime(context.Background(), apply, nil) + if !errors.Is(errRun, ErrRebalanceRequested) || client.HeartbeatOK() || client.dispatchFenced.Load() { + t.Fatalf("planned end = %v", errRun) + } + next := client.NewLifetime() + next.SetManagedLifetime(true) + t.Cleanup(next.Close) + ready := false + _ = next.RunConfigSubscriberLifetime(context.Background(), func([]byte) error { + return next.SetLifecycleConfig(config.CredentialConcurrencyConfig{LifecycleConfigRevision: 1}) + }, func() { ready = next.HeartbeatOK(); next.Close() }) + if !ready || next.balance.moving || next.balance.result != "" { + t.Fatal("target ACK/probe/result did not complete") + } + args := findRedisCommand(targetCommands.All(), "SUBSCRIBE") + if len(args) < 4 || args[3] != "takeover" || next.homeCfg.Port != target.homeCfg.Port { + t.Fatalf("replacement subscription = %#v", args) + } + found := false + for _, command := range targetCommands.All() { + if len(command) == 5 && command[2] == "RESULT" && command[3] == "plan" && command[4] == "success" { + found = true + } + } + if !found { + t.Fatal("successful migration not reported") + } +} + +func TestClusterBalanceBootstrapTieAndNormalization(t *testing.T) { + client := New(config.HomeConfig{Host: "home-0", Port: 8327}) + nodes := normalizeClusterNodes([]clusterNode{{IP: "home-0", Port: 8327}, {IP: "home-1", Port: 8327}, {IP: " home-1 ", Port: 8327}, {IP: "bad", Port: 65536}}) + if len(nodes) != 2 { + t.Fatalf("nodes = %#v", nodes) + } + for index, draw := range []float64{0, 0.9} { + node, ok := client.chooseClusterNodeLocked(nodes, false, draw) + if !ok || node.IP != fmt.Sprintf("home-%d", index) { + t.Fatalf("choice = %#v", node) + } + } +} diff --git a/internal/logging/home_rebalance_log_test.go b/internal/logging/home_rebalance_log_test.go new file mode 100644 index 0000000000..12a2458aa4 --- /dev/null +++ b/internal/logging/home_rebalance_log_test.go @@ -0,0 +1,48 @@ +package logging + +import ( + "bytes" + "encoding/json" + "fmt" + "strings" + "testing" + + log "github.com/sirupsen/logrus" +) + +func TestHomeRebalanceLogPreservesPlanAndAddressesLocallyAndWhenForwarded(t *testing.T) { + var local bytes.Buffer + logger := log.New() + logger.SetOutput(&local) + logger.SetFormatter(&LogFormatter{}) + stub := &stubHomeAppLogClient{heartbeatOK: true} + forwarder := &HomeAppLogForwarder{ + formatter: &LogFormatter{}, + queue: make(chan homeAppLogPayload, 1), + } + forwarder.enabled.Store(true) + forwarder.bind(stub) + logger.AddHook(forwarder) + + logger.Infof("executing Home rebalance plan: plan=%q from=%q to=%q", "plan-1", "home-1:8327", "home-2:8327") + select { + case payload := <-forwarder.queue: + forwarder.forward(payload) + default: + t.Fatal("rebalance log was not queued for Home") + } + if stub.pushedCount() != 1 { + t.Fatal("rebalance log was not forwarded to Home") + } + var forwarded homeAppLogPayload + if errJSON := json.Unmarshal(stub.pushedAt(0), &forwarded); errJSON != nil { + t.Fatal(errJSON) + } + for destination, line := range map[string]string{"local": local.String(), "Home": forwarded.Line} { + for key, value := range map[string]string{"plan": "plan-1", "from": "home-1:8327", "to": "home-2:8327"} { + if !strings.Contains(line, fmt.Sprintf("%s=%q", key, value)) { + t.Fatalf("%s log missing %s: %q", destination, key, line) + } + } + } +} diff --git a/sdk/cliproxy/service_executionregistry_test.go b/sdk/cliproxy/service_executionregistry_test.go index 575847f81e..8be8724209 100644 --- a/sdk/cliproxy/service_executionregistry_test.go +++ b/sdk/cliproxy/service_executionregistry_test.go @@ -131,9 +131,11 @@ func TestServiceConcurrentReplacementWaitsForInFlightDrain(t *testing.T) { }() deadline := time.Now().Add(time.Second) for { - if _, errLate := registry.BeginDispatch(); errLate != nil { + late, errLate := registry.BeginDispatch() + if errLate != nil { break } + late.End() if time.Now().After(deadline) { t.Fatal("first replacement did not begin draining") } diff --git a/sdk/cliproxy/service_home.go b/sdk/cliproxy/service_home.go index cc751909a7..ca9a23a1ec 100644 --- a/sdk/cliproxy/service_home.go +++ b/sdk/cliproxy/service_home.go @@ -256,8 +256,6 @@ func (s *Service) homeLifetimeActive(homeCtx, lifetimeCtx context.Context, gener } func (s *Service) finalizeHomePluginWorkUntilDone(ctx, homeCtx context.Context, generation uint64, client *home.Client, work *homePluginFinalization, publish func() bool) error { - stopClose := closeHomeClientOnCancellation(ctx, client) - defer stopClose() for { if errContext := ctx.Err(); errContext != nil { return errContext @@ -289,21 +287,6 @@ func (s *Service) finalizeHomePluginWorkUntilDone(ctx, homeCtx context.Context, } } -func closeHomeClientOnCancellation(ctx context.Context, client *home.Client) func() { - if ctx == nil || client == nil { - return func() {} - } - stop := make(chan struct{}) - go func() { - select { - case <-ctx.Done(): - client.Close() - case <-stop: - } - }() - return func() { close(stop) } -} - func logHomeConfigChanges(oldCfg, newCfg *config.Config) { if oldCfg == nil || newCfg == nil || !newCfg.Home.Enabled || (!oldCfg.Debug && !newCfg.Debug) { return @@ -576,6 +559,16 @@ func (s *Service) runHomeSubscriber(homeCtx context.Context, parentCtx context.C }, func() { readyOnce.Do(func() { close(ready) }) }) + plannedRebalance := errors.Is(errRun, home.ErrRebalanceRequested) + settleBound := time.Duration(cancelBound.Load()) + var errPending error + if plannedRebalance { + // Settle issued dispatches before deciding between target takeover and source drain/flush. + client.FenceDispatch() + settleCtx, cancelSettle := context.WithTimeout(context.WithoutCancel(parentCtx), settleBound) + errPending = registry.WaitPending(settleCtx) + cancelSettle() + } lifetimeCancel() <-workerDone if publisherDone := supervisor.publisherCompletion(); publisherDone != nil { @@ -588,11 +581,11 @@ func (s *Service) runHomeSubscriber(homeCtx context.Context, parentCtx context.C releaseCancel() <-releaseDone client.Close() - - settleBound := time.Duration(cancelBound.Load()) - settleCtx, cancelSettle := context.WithTimeout(context.WithoutCancel(parentCtx), settleBound) - errPending := registry.WaitPending(settleCtx) - cancelSettle() + if !plannedRebalance { + settleCtx, cancelSettle := context.WithTimeout(context.WithoutCancel(parentCtx), settleBound) + errPending = registry.WaitPending(settleCtx) + cancelSettle() + } if errPending != nil { log.WithError(errPending).Error("failed to settle pending Home dispatches before subscriber replacement") s.cancelServiceRun() @@ -617,7 +610,11 @@ func (s *Service) runHomeSubscriber(homeCtx context.Context, parentCtx context.C releaseFlusher = home.NewReleaseFlusher(nil, nil) registry.SetReleaseSink(releaseFlusher.MarkDirty) } - log.WithError(errRun).Warn("home config subscription lifetime ended") + if errors.Is(errRun, home.ErrRebalanceRequested) { + log.Info("home config subscription reconnecting for load balance") + } else { + log.WithError(errRun).Warn("home config subscription lifetime ended") + } if !published.Load() && !waitForHomeSubscriberRetry(homeCtx, homeSubscriberPreAckRetryBackoff) { return } diff --git a/sdk/cliproxy/service_home_cancellation_test.go b/sdk/cliproxy/service_home_cancellation_test.go new file mode 100644 index 0000000000..8a2af03e60 --- /dev/null +++ b/sdk/cliproxy/service_home_cancellation_test.go @@ -0,0 +1,139 @@ +package cliproxy + +import ( + "context" + "encoding/json" + "errors" + "sync" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/home" + "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executionregistry" +) + +func TestServiceNormalReplacementDrainsAndFlushesAccountedExecution(t *testing.T) { + server := newRebalanceAdmissionHome(t, false) + releaseACK := make(chan struct{}) + var ackOnce sync.Once + server.mu.Lock() + server.releaseACK = releaseACK + server.mu.Unlock() + forwarder := &rebalanceAdmissionForwarder{bound: make(chan *home.Client, 4), detached: make(chan *home.Client, 4)} + previousStart := startHomeLogForwarder + startHomeLogForwarder = func(int) homeLogForwarder { return forwarder } + t.Cleanup(func() { startHomeLogForwarder = previousStart }) + service := newRegistryTestService(t, server.listener) + ctx, cancel := context.WithCancel(context.Background()) + service.runCancel = cancel + service.startHomeSubscriber(ctx) + var replacementDone chan struct{} + t.Cleanup(func() { + ackOnce.Do(func() { close(releaseACK) }) + cancel() + if replacementDone != nil { + awaitRebalanceEvent(t, replacementDone) + } + awaitRebalanceEvent(t, service.homeSupervisor.done) + home.ClearCurrent() + }) + client := awaitRebalanceEvent(t, forwarder.bound) + awaitRebalanceEvent(t, server.subscribed) + service.homeMu.Lock() + registry := service.homeRegistry + service.homeMu.Unlock() + oldSupervisor := service.homeSupervisor + if registry == nil || !client.HeartbeatOK() { + t.Fatal("initial subscriber is not ready") + } + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scope, errInstall := registry.Install(pending, executionregistry.ScopeSpec{ + RequestID: "accounted-request", CredentialID: "credential-a", Model: "test-model", Accounted: true, + }) + if errInstall != nil { + t.Fatal(errInstall) + } + t.Cleanup(func() { scope.End("test cleanup") }) + draining := make(chan struct{}) + tickets := make(chan *executionregistry.ReleaseTicket, 1) + if errBind := scope.Bind(func() error { + close(draining) + go func() { tickets <- scope.EndWithRelease("normal replacement") }() + return nil + }); errBind != nil { + t.Fatal(errBind) + } + + replacementDone = make(chan struct{}) + go func() { + service.startHomeSubscriber(ctx) + close(replacementDone) + }() + awaitRebalanceEvent(t, draining) + ticket := awaitRebalanceEvent(t, tickets) + if ticket == nil { + t.Fatal("accounted execution did not produce a release ticket") + } + var args []string + select { + case args = <-server.releases: + case <-oldSupervisor.done: + t.Fatalf("old supervisor ended without sending its final release: service error=%v", ctx.Err()) + case <-ctx.Done(): + t.Fatalf("normal replacement canceled the entire service: %v", ctx.Err()) + case <-time.After(7 * time.Second): + t.Fatal("normal replacement did not send its final release") + } + var frame home.ConcurrencyReleaseFrame + if errJSON := json.Unmarshal([]byte(args[2]), &frame); errJSON != nil { + t.Fatal(errJSON) + } + want := home.ConcurrencyReleaseFrame{CredentialID: "credential-a", Model: "test-model", ReleaseSeq: 1} + if frame != want { + t.Fatalf("release = %#v, want %#v", frame, want) + } + if client.HeartbeatOK() || client.AmbiguousDispatch() || ctx.Err() != nil { + t.Fatal("normal cancellation retained heartbeat, introduced ambiguity, or canceled the service") + } + select { + case <-oldSupervisor.done: + t.Fatal("old supervisor finished before release acknowledgement") + case <-replacementDone: + t.Fatal("replacement finished before release acknowledgement") + case <-server.subscribed: + t.Fatal("new subscription started before release acknowledgement") + default: + } + if late, errLate := registry.BeginDispatch(); !errors.Is(errLate, executionregistry.ErrRegistryNotAccepting) { + if late != nil { + late.End() + } + t.Fatalf("drained registry accepted new dispatch: %v", errLate) + } + + ackOnce.Do(func() { close(releaseACK) }) + awaitRebalanceEvent(t, oldSupervisor.done) + awaitRebalanceEvent(t, replacementDone) + next := awaitRebalanceEvent(t, forwarder.bound) + awaitRebalanceEvent(t, server.subscribed) + if ctx.Err() != nil || next == client || !next.HeartbeatOK() { + t.Fatalf("replacement did not keep service running: error=%v", ctx.Err()) + } + waitCtx, cancelWait := context.WithTimeout(context.Background(), time.Second) + defer cancelWait() + if errTicket := ticket.Wait(waitCtx); errTicket != nil { + t.Fatalf("final release ticket was not acknowledged: %v", errTicket) + } + if errRelease := client.PushConcurrencyRelease(context.Background(), want); !errors.Is(errRelease, home.ErrDispatchFenced) { + t.Fatalf("old client remained open after flush: %v", errRelease) + } + service.homeMu.Lock() + nextRegistry := service.homeRegistry + service.homeMu.Unlock() + if nextRegistry == nil || nextRegistry == registry { + t.Fatal("normal replacement did not publish a fresh registry") + } +} diff --git a/sdk/cliproxy/service_home_rebalance_finalization_test.go b/sdk/cliproxy/service_home_rebalance_finalization_test.go new file mode 100644 index 0000000000..9ec390cd9d --- /dev/null +++ b/sdk/cliproxy/service_home_rebalance_finalization_test.go @@ -0,0 +1,159 @@ +package cliproxy + +import ( + "context" + "errors" + "runtime" + "sync" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/home" + "github.com/router-for-me/CLIProxyAPI/v8/internal/homeplugins" + "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executionregistry" + "github.com/router-for-me/CLIProxyAPI/v8/sdk/config" +) + +func TestServicePlannedRebalanceSettlesDispatchBeforeCancelingPluginFinalization(t *testing.T) { + source := newRebalanceAdmissionHome(t, false) + target := newRebalanceAdmissionHome(t, false) + source.mu.Lock() + source.target = target.listener.Addr().String() + source.mu.Unlock() + forwarder := &rebalanceAdmissionForwarder{bound: make(chan *home.Client, 4), detached: make(chan *home.Client, 4)} + previousStart := startHomeLogForwarder + startHomeLogForwarder = func(int) homeLogForwarder { return forwarder } + t.Cleanup(func() { startHomeLogForwarder = previousStart }) + service := newRegistryTestService(t, source.listener) + service.cfg.Home.DisableClusterDiscovery = false + service.cfg.Home.NodeID = "source-node" + finalizing := make(chan context.Context, 1) + finalizationCanceled := make(chan struct{}) + allowFinalizationReturn := make(chan struct{}) + var returnOnce sync.Once + serviceCtx, cancelService := context.WithCancel(context.Background()) + service.homePluginDeleteTask = func(ctx context.Context, _ *config.Config, _ home.PluginTask) homeplugins.SyncReport { + finalizing <- ctx + <-ctx.Done() + close(finalizationCanceled) + select { + case <-allowFinalizationReturn: + case <-serviceCtx.Done(): + } + return homeplugins.SyncReport{Task: "delete", Error: ctx.Err().Error()} + } + service.runCancel = cancelService + service.startHomeSubscriber(serviceCtx) + t.Cleanup(func() { + returnOnce.Do(func() { close(allowFinalizationReturn) }) + cancelService() + awaitRebalanceEvent(t, service.homeSupervisor.done) + home.ClearCurrent() + }) + client := awaitRebalanceEvent(t, forwarder.bound) + t.Cleanup(client.Close) + service.homeMu.Lock() + registry := service.homeRegistry + service.homeMu.Unlock() + pendingActive, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + active, errInstall := registry.Install(pendingActive, executionregistry.ScopeSpec{RequestID: "running-model"}) + if errInstall != nil { + t.Fatal(errInstall) + } + t.Cleanup(func() { active.End("test cleanup") }) + drained := make(chan struct{}) + if errBind := active.Bind(func() error { + close(drained) + go active.End("drained") + return nil + }); errBind != nil { + t.Fatal(errBind) + } + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + allowProcessing := make(chan struct{}) + var processOnce sync.Once + t.Cleanup(func() { processOnce.Do(func() { close(allowProcessing) }) }) + response := make(chan error, 1) + processed := make(chan error, 1) + go func() { + defer pending.End() + _, errDispatch := client.RPopAuth(serviceCtx, "test-model", "", nil, 1) + response <- errDispatch + <-allowProcessing + if home.IsAmbiguousDispatchError(errDispatch) { + client.AbortAmbiguousDispatch() + } + if errDispatch == nil { + scope, errScope := registry.Install(pending, executionregistry.ScopeSpec{RequestID: "received-dispatch"}) + errDispatch = errScope + if scope != nil { + scope.End("processed") + } + } + processed <- errDispatch + }() + awaitRebalanceEvent(t, source.issued) + source.pluginTasks.Store(true) + source.configUpdates <- "credential-concurrency:\n lifecycle-config-revision: 2\n cpa-heartbeat-timeout: 10s\n cpa-cancel-bound: 5s\n" + finalizationCtx := awaitRebalanceEvent(t, finalizing) + close(source.migrate) + deadline := time.NewTimer(5 * time.Second) + defer deadline.Stop() + for { + _, errFence := client.RPopAuth(serviceCtx, "", "", nil, 1) + if errors.Is(errFence, home.ErrDispatchFenced) { + break + } + select { + case <-deadline.C: + t.Fatal("planned migration did not fence dispatch admission") + default: + runtime.Gosched() + } + } + if finalizationCtx.Err() != nil { + t.Fatal("plugin finalization was canceled while its close handler could interrupt the pending dispatch") + } + select { + case <-source.transportEnded: + t.Fatal("plugin finalization closed the pending RPOP connection") + case <-drained: + t.Fatal("planned migration drained a running model") + default: + } + close(source.reply) + if errResponse := awaitRebalanceEvent(t, response); errResponse != nil { + t.Fatalf("pending RPOP was interrupted: %v", errResponse) + } + if finalizationCtx.Err() != nil { + t.Fatal("plugin finalization was canceled before CPA processed the dispatch response") + } + processOnce.Do(func() { close(allowProcessing) }) + if errProcessed := awaitRebalanceEvent(t, processed); errProcessed != nil { + t.Fatalf("response processing failed: %v", errProcessed) + } + awaitRebalanceEvent(t, finalizationCanceled) + returnOnce.Do(func() { close(allowFinalizationReturn) }) + subscription := awaitRebalanceEvent(t, target.subscribed) + next := awaitRebalanceEvent(t, forwarder.bound) + service.homeMu.Lock() + nextRegistry := service.homeRegistry + service.homeMu.Unlock() + if client.AmbiguousDispatch() || nextRegistry != registry || next == client || !next.HeartbeatOK() || serviceCtx.Err() != nil { + t.Fatal("planned migration did not preserve the healthy registry and service") + } + if len(subscription) < 4 || subscription[3] != "takeover" { + t.Fatalf("replacement subscription lost takeover: %#v", subscription) + } + select { + case <-drained: + t.Fatal("plugin finalization cancellation drained the running model") + default: + } +} diff --git a/sdk/cliproxy/service_home_rebalance_release_test.go b/sdk/cliproxy/service_home_rebalance_release_test.go new file mode 100644 index 0000000000..f84d2d882f --- /dev/null +++ b/sdk/cliproxy/service_home_rebalance_release_test.go @@ -0,0 +1,221 @@ +package cliproxy + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "runtime" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/home" + "github.com/router-for-me/CLIProxyAPI/v8/internal/homeplugins" + "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executionregistry" + "github.com/router-for-me/CLIProxyAPI/v8/sdk/config" +) + +func TestServiceRebalanceReleaseOwnershipAcrossCancellation(t *testing.T) { + for _, mode := range []string{"success", "shutdown", "replacement"} { + for _, pluginFinalization := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/plugins=%v", mode, pluginFinalization), func(t *testing.T) { + source := newRebalanceAdmissionHome(t, false) + target := newRebalanceAdmissionHome(t, false) + ack := make(chan struct{}) + var ackOnce sync.Once + source.mu.Lock() + source.target = target.listener.Addr().String() + if mode != "success" { + source.releaseACK = ack + } + source.mu.Unlock() + if mode == "success" { + target.mu.Lock() + target.releaseACK = ack + target.mu.Unlock() + } + forwarder := &rebalanceAdmissionForwarder{bound: make(chan *home.Client, 4), detached: make(chan *home.Client, 4)} + previousStart := startHomeLogForwarder + startHomeLogForwarder = func(int) homeLogForwarder { return forwarder } + t.Cleanup(func() { startHomeLogForwarder = previousStart }) + service := newRegistryTestService(t, source.listener) + service.cfg.Home.DisableClusterDiscovery = false + service.cfg.Home.NodeID = "source-node" + finalizing := make(chan context.Context, 1) + service.homePluginDeleteTask = func(ctx context.Context, _ *config.Config, _ home.PluginTask) homeplugins.SyncReport { + finalizing <- ctx + <-ctx.Done() + return homeplugins.SyncReport{Task: "delete", Error: ctx.Err().Error()} + } + ctx, cancel := context.WithCancel(context.Background()) + var canceledBySupervisor atomic.Int32 + service.runCancel = func() { canceledBySupervisor.Add(1); cancel() } + service.startHomeSubscriber(ctx) + var replacementDone chan struct{} + t.Cleanup(func() { + ackOnce.Do(func() { close(ack) }) + cancel() + if replacementDone != nil { + awaitRebalanceEvent(t, replacementDone) + } + awaitRebalanceEvent(t, service.homeSupervisor.done) + home.ClearCurrent() + }) + client := awaitRebalanceEvent(t, forwarder.bound) + awaitRebalanceEvent(t, source.subscribed) + service.homeMu.Lock() + registry := service.homeRegistry + service.homeMu.Unlock() + oldSupervisor := service.homeSupervisor + activePending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + active, errInstall := registry.Install(activePending, executionregistry.ScopeSpec{ + RequestID: "accounted", CredentialID: "credential-a", Model: "test-model", Accounted: true, + }) + if errInstall != nil { + t.Fatal(errInstall) + } + t.Cleanup(func() { active.End("test cleanup") }) + drained := make(chan struct{}) + tickets := make(chan *executionregistry.ReleaseTicket, 1) + if errBind := active.Bind(func() error { + close(drained) + go func() { tickets <- active.EndWithRelease("drained") }() + return nil + }); errBind != nil { + t.Fatal(errBind) + } + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + result := make(chan error, 1) + go func() { + defer pending.End() + _, errDispatch := client.RPopAuth(context.Background(), "test-model", "", nil, 1) + if home.IsAmbiguousDispatchError(errDispatch) { + client.AbortAmbiguousDispatch() + } + result <- errDispatch + }() + t.Cleanup(client.Close) + awaitRebalanceEvent(t, source.issued) + if pluginFinalization { + source.pluginTasks.Store(true) + source.configUpdates <- "credential-concurrency:\n lifecycle-config-revision: 2\n cpa-heartbeat-timeout: 10s\n cpa-cancel-bound: 5s\n" + awaitRebalanceEvent(t, finalizing) + source.pluginTasks.Store(false) + } + close(source.migrate) + deadline := time.NewTimer(5 * time.Second) + defer deadline.Stop() + for { + _, errFence := client.RPopAuth(context.Background(), "", "", nil, 1) + if errors.Is(errFence, home.ErrDispatchFenced) { + break + } + select { + case <-deadline.C: + t.Fatal("migration did not fence admission") + default: + runtime.Gosched() + } + } + switch mode { + case "shutdown": + cancel() + case "replacement": + // Cancel synchronously at the pending wait, then run the normal replacement entrypoint. + oldSupervisor.cancel() + source.mu.Lock() + source.target = "" + source.mu.Unlock() + replacementDone = make(chan struct{}) + go func() { service.startHomeSubscriber(ctx); close(replacementDone) }() + } + close(source.reply) + if errDispatch := awaitRebalanceEvent(t, result); errDispatch != nil { + t.Fatalf("cancellation interrupted issued dispatch: %v", errDispatch) + } + owner := source + if mode == "success" { + awaitRebalanceEvent(t, target.subscribed) + next := awaitRebalanceEvent(t, forwarder.bound) + service.homeMu.Lock() + nextRegistry := service.homeRegistry + service.homeMu.Unlock() + if next == client || !next.HeartbeatOK() || nextRegistry != registry { + t.Fatal("successful migration did not preserve active execution ownership") + } + select { + case <-drained: + t.Fatal("successful migration drained the accounted execution") + default: + } + active.End("completed at destination") + owner = target + } + ticket := awaitRebalanceEvent(t, tickets) + if ticket == nil { + t.Fatal("accounted execution lost its release ticket") + } + var args []string + select { + case args = <-owner.releases: + case <-oldSupervisor.done: + t.Fatalf("supervisor exited without final release: canceled service=%d", canceledBySupervisor.Load()) + case <-time.After(7 * time.Second): + t.Fatal("release was not sent to the membership owner") + } + var frame home.ConcurrencyReleaseFrame + if errJSON := json.Unmarshal([]byte(args[2]), &frame); errJSON != nil { + t.Fatal(errJSON) + } + if frame != (home.ConcurrencyReleaseFrame{CredentialID: "credential-a", Model: "test-model", ReleaseSeq: 1}) { + t.Fatalf("unexpected release: %#v", frame) + } + if mode != "success" { + select { + case <-oldSupervisor.done: + t.Fatal("cleanup completed before the final release was acknowledged") + case <-target.subscribed: + t.Fatal("canceled migration contacted the target subscription") + default: + } + } else { + select { + case <-source.releases: + t.Fatal("completed migration released through its former source") + default: + } + } + ackOnce.Do(func() { close(ack) }) + waitCtx, cancelWait := context.WithTimeout(context.Background(), time.Second) + defer cancelWait() + if errTicket := ticket.Wait(waitCtx); errTicket != nil { + t.Fatal(errTicket) + } + if mode != "success" { + awaitRebalanceEvent(t, oldSupervisor.done) + } + if mode == "replacement" { + awaitRebalanceEvent(t, replacementDone) + next := awaitRebalanceEvent(t, forwarder.bound) + if next == client || !next.HeartbeatOK() { + t.Fatal("replacement was not published") + } + } + if canceledBySupervisor.Load() != 0 || (mode != "shutdown" && ctx.Err() != nil) { + t.Fatal("final release failed and canceled the parent service") + } + if _, errLate := client.RPopAuth(context.Background(), "late", "", nil, 1); !errors.Is(errLate, home.ErrDispatchFenced) { + t.Fatalf("old lifetime reopened dispatch: %v", errLate) + } + }) + } + } +} diff --git a/sdk/cliproxy/service_home_rebalance_test.go b/sdk/cliproxy/service_home_rebalance_test.go new file mode 100644 index 0000000000..c1a3c9e4cb --- /dev/null +++ b/sdk/cliproxy/service_home_rebalance_test.go @@ -0,0 +1,402 @@ +package cliproxy + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net" + "runtime" + "strconv" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/router-for-me/CLIProxyAPI/v8/internal/home" + "github.com/router-for-me/CLIProxyAPI/v8/sdk/cliproxy/executionregistry" +) + +// rebalanceAdmissionHome controls subscription migration and the delivery of an issued RPOP. +type rebalanceAdmissionHome struct { + listener net.Listener + target string + migrate chan struct{} + reply chan struct{} + stop chan struct{} + subscribed chan []string + configUpdates chan string + pluginTasks atomic.Bool + releases chan []string + releaseACK <-chan struct{} + issued chan struct{} + transportEnded chan struct{} + ambiguous bool + rpops atomic.Int32 + mu sync.Mutex + connections map[net.Conn]struct{} + workers sync.WaitGroup +} + +func newRebalanceAdmissionHome(t *testing.T, ambiguous bool) *rebalanceAdmissionHome { + t.Helper() + listener, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatal(errListen) + } + server := &rebalanceAdmissionHome{ + listener: listener, ambiguous: ambiguous, + migrate: make(chan struct{}), reply: make(chan struct{}), stop: make(chan struct{}), + subscribed: make(chan []string, 4), configUpdates: make(chan string, 1), releases: make(chan []string, 4), + issued: make(chan struct{}), transportEnded: make(chan struct{}), + connections: make(map[net.Conn]struct{}), + } + accepted := make(chan struct{}) + go func() { + defer close(accepted) + for { + conn, errAccept := listener.Accept() + if errAccept != nil { + return + } + server.mu.Lock() + server.connections[conn] = struct{}{} + server.mu.Unlock() + server.workers.Add(1) + go server.serve(conn) + } + }() + t.Cleanup(func() { + close(server.stop) + _ = listener.Close() + <-accepted + server.mu.Lock() + for conn := range server.connections { + _ = conn.Close() + } + server.mu.Unlock() + server.workers.Wait() + }) + return server +} + +func (s *rebalanceAdmissionHome) serve(conn net.Conn) { + defer s.workers.Done() + defer func() { + _ = conn.Close() + s.mu.Lock() + delete(s.connections, conn) + s.mu.Unlock() + }() + reader := bufio.NewReader(conn) + for { + args, errRead := readRegistryTestRedisCommand(reader) + if errRead != nil { + return + } + var reply string + switch strings.ToUpper(args[0]) { + case "HELLO": + reply = "%6\r\n$6\r\nserver\r\n$5\r\nredis\r\n$5\r\nproto\r\n:3\r\n$2\r\nid\r\n:1\r\n$4\r\nmode\r\n$10\r\nstandalone\r\n$4\r\nrole\r\n$6\r\nmaster\r\n$7\r\nmodules\r\n*0\r\n" + case "GET": + payload := "credential-concurrency:\n lifecycle-config-revision: 1\n cpa-heartbeat-timeout: 10s\n cpa-cancel-bound: 5s\n" + if args[1] != "config" { + payload = "[]" + if args[1] == "plugin-tasks" && s.pluginTasks.Load() { + payload = `[{"id":1,"operation":"delete","plugin_id":"test-plugin"}]` + } + } + reply = fmt.Sprintf("$%d\r\n%s\r\n", len(payload), payload) + case "CLUSTER": + payload := `{"ok":true,"nodes":[]}` + if args[1] == "BALANCE" { + payload = `{"ok":true}` + } + reply = fmt.Sprintf("$%d\r\n%s\r\n", len(payload), payload) + case "RPUSH": + reply = ":1\r\n" + case "LPUSH": + if len(args) == 3 && args[1] == "concurrency-release" { + select { + case s.releases <- append([]string(nil), args...): + case <-s.stop: + return + } + s.mu.Lock() + ack := s.releaseACK + s.mu.Unlock() + if ack != nil { + select { + case <-ack: + case <-s.stop: + return + } + } + } + reply = ":1\r\n" + case "SUBSCRIBE": + s.subscribed <- append([]string(nil), args...) + if _, errWrite := io.WriteString(conn, "*3\r\n$9\r\nsubscribe\r\n$6\r\nconfig\r\n:1\r\n"); errWrite != nil { + return + } + s.mu.Lock() + target := s.target + s.mu.Unlock() + if target != "" { + updates: + for { + select { + case payload := <-s.configUpdates: + writeRegistryTestMessage(conn, payload) + case <-s.migrate: + break updates + case <-s.stop: + return + } + } + host, portText, errSplit := net.SplitHostPort(target) + if errSplit != nil { + return + } + port, errPort := strconv.Atoi(portText) + if errPort != nil { + return + } + payload, errJSON := json.Marshal(map[string]any{"id": "admission-plan", "term": 1, "target": map[string]any{"ip": host, "port": port}}) + if errJSON != nil { + return + } + if _, errWrite := io.WriteString(conn, fmt.Sprintf("*3\r\n$7\r\nmessage\r\n$9\r\nrebalance\r\n$%d\r\n%s\r\n", len(payload), payload)); errWrite != nil { + return + } + } + <-s.stop + return + case "RPOP": + if s.rpops.Add(1) != 1 { + return + } + close(s.issued) + // No more commands are expected on this connection until the old client closes. + s.workers.Add(1) + go func() { + defer s.workers.Done() + _, _ = reader.ReadByte() + close(s.transportEnded) + }() + select { + case <-s.reply: + case <-s.stop: + return + } + if s.ambiguous { + return + } + if _, errWrite := io.WriteString(conn, "$2\r\n{}\r\n"); errWrite != nil { + return + } + select { + case <-s.transportEnded: + case <-s.stop: + } + return + default: + reply = "+OK\r\n" + } + if _, errWrite := io.WriteString(conn, reply); errWrite != nil { + return + } + } +} + +type rebalanceAdmissionForwarder struct { + bound chan *home.Client + detached chan *home.Client +} + +func (f *rebalanceAdmissionForwarder) Bind(client *home.Client) { f.bound <- client } +func (f *rebalanceAdmissionForwarder) Deactivate(client *home.Client) { + f.detached <- client +} +func (f *rebalanceAdmissionForwarder) Stop() {} + +func awaitRebalanceEvent[T any](t *testing.T, events <-chan T) T { + t.Helper() + select { + case event := <-events: + return event + case <-time.After(5 * time.Second): + t.Fatal("rebalance synchronization event did not arrive") + var zero T + return zero + } +} + +func TestServicePlannedRebalanceFencesAdmissionAndSettlesResponseProcessing(t *testing.T) { + for _, ambiguous := range []bool{false, true} { + t.Run(fmt.Sprintf("ambiguous=%v", ambiguous), func(t *testing.T) { + source := newRebalanceAdmissionHome(t, ambiguous) + target := newRebalanceAdmissionHome(t, false) + source.mu.Lock() + source.target = target.listener.Addr().String() + source.mu.Unlock() + forwarder := &rebalanceAdmissionForwarder{bound: make(chan *home.Client, 4), detached: make(chan *home.Client, 4)} + previousStart := startHomeLogForwarder + startHomeLogForwarder = func(int) homeLogForwarder { return forwarder } + t.Cleanup(func() { startHomeLogForwarder = previousStart }) + service := newRegistryTestService(t, source.listener) + service.cfg.Home.DisableClusterDiscovery = false + ctx, cancel := context.WithCancel(context.Background()) + service.runCancel = cancel + service.startHomeSubscriber(ctx) + t.Cleanup(func() { + cancel() + awaitRebalanceEvent(t, service.homeSupervisor.done) + home.ClearCurrent() + }) + client := awaitRebalanceEvent(t, forwarder.bound) + t.Cleanup(client.Close) + service.homeMu.Lock() + registry := service.homeRegistry + service.homeMu.Unlock() + if registry == nil || !client.HeartbeatOK() { + t.Fatal("source lifetime was not published healthy") + } + activePending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + active, errInstall := registry.Install(activePending, executionregistry.ScopeSpec{RequestID: "already-running"}) + if errInstall != nil { + t.Fatal(errInstall) + } + t.Cleanup(func() { active.End("test cleanup") }) + drained := make(chan struct{}) + var drainOnce sync.Once + if errBind := active.Bind(func() error { + drainOnce.Do(func() { close(drained) }) + go active.End("unsafe replacement") + return nil + }); errBind != nil { + t.Fatal(errBind) + } + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + processResponse := make(chan struct{}) + var processOnce sync.Once + t.Cleanup(func() { processOnce.Do(func() { close(processResponse) }) }) + response := make(chan error, 1) + processed := make(chan error, 1) + go func() { + defer pending.End() + raw, errDispatch := client.RPopAuth(ctx, "test-model", "", nil, 1) + response <- errDispatch + <-processResponse + if home.IsAmbiguousDispatchError(errDispatch) { + // Match conductor ordering: record ambiguity before releasing the pending token. + client.AbortAmbiguousDispatch() + } + if errDispatch != nil { + processed <- errDispatch + return + } + if string(raw) != "{}" { + processed <- fmt.Errorf("dispatch response = %q", raw) + return + } + scope, errScope := registry.Install(pending, executionregistry.ScopeSpec{RequestID: "newly-dispatched"}) + if errScope == nil { + // Keep this execution running across migration, without a model completion barrier. + errScope = scope.Bind(func() error { + drainOnce.Do(func() { close(drained) }) + go scope.End("unsafe replacement") + return nil + }) + } + processed <- errScope + }() + awaitRebalanceEvent(t, source.issued) + close(source.migrate) + // Observe the admission fence without issuing network traffic or adding sleeps. + fenceDeadline := time.NewTimer(5 * time.Second) + defer fenceDeadline.Stop() + for { + _, errFence := client.RPopAuth(ctx, "", "", nil, 1) + if errors.Is(errFence, home.ErrDispatchFenced) { + break + } + select { + case <-fenceDeadline.C: + t.Fatal("planned migration did not fence admission before settling pending dispatches") + default: + runtime.Gosched() + } + } + select { + case <-source.transportEnded: + t.Fatal("migration closed an issued dispatch before its response arrived") + case <-target.subscribed: + t.Fatal("migration completed before the pending dispatch settled") + case <-drained: + t.Fatal("planned migration canceled an active execution") + default: + } + close(source.reply) + errResponse := awaitRebalanceEvent(t, response) + if home.IsAmbiguousDispatchError(errResponse) != ambiguous || (!ambiguous && errResponse != nil) { + t.Fatalf("issued RPOP response = %v", errResponse) + } + select { + case <-target.subscribed: + t.Fatal("Home sending credentials was treated as completed CPA response processing") + default: + } + processOnce.Do(func() { close(processResponse) }) + errProcessed := awaitRebalanceEvent(t, processed) + if home.IsAmbiguousDispatchError(errProcessed) != ambiguous || (!ambiguous && errProcessed != nil) { + t.Fatalf("dispatch processing = %v", errProcessed) + } + if detached := awaitRebalanceEvent(t, forwarder.detached); detached != client { + t.Fatal("migration detached the wrong lifetime") + } + subscription := awaitRebalanceEvent(t, target.subscribed) + next := awaitRebalanceEvent(t, forwarder.bound) + service.homeMu.Lock() + nextRegistry := service.homeRegistry + service.homeMu.Unlock() + if next == client || !next.HeartbeatOK() { + t.Fatal("replacement lifetime was not published") + } + if ambiguous { + awaitRebalanceEvent(t, drained) + if !client.AmbiguousDispatch() || registry == nextRegistry || strings.Contains(strings.Join(subscription, " "), "takeover") { + t.Fatalf("unsafe replacement preserved registry or takeover: subscription=%#v", subscription) + } + } else { + if registry != nextRegistry || len(subscription) < 4 || subscription[3] != "takeover" { + t.Fatalf("safe migration did not preserve registry and takeover: %#v", subscription) + } + select { + case <-drained: + t.Fatal("safe migration drained an active model execution") + default: + } + // A stale request may register after migration's empty pending-set observation. + late, errLateBegin := registry.BeginDispatch() + if errLateBegin != nil { + t.Fatal(errLateBegin) + } + _, errLate := client.RPopAuth(ctx, "late-model", "", nil, 1) + late.End() + if !errors.Is(errLate, home.ErrDispatchFenced) || source.rpops.Load() != 1 { + t.Fatalf("stale request issued a new RPOP: error=%v count=%d", errLate, source.rpops.Load()) + } + } + }) + } +}