diff --git a/internal/home/client.go b/internal/home/client.go index a8c0a729a..3ac162ac1 100644 --- a/internal/home/client.go +++ b/internal/home/client.go @@ -93,6 +93,15 @@ var ( ErrDispatchFenced = errors.New("home auth dispatch is fenced") ) +// IsMembershipTakeoverUnavailableError reports whether Home cannot preserve the previous membership state. +func IsMembershipTakeoverUnavailableError(err error) bool { + if err == nil { + return false + } + message := strings.ToLower(err.Error()) + return strings.Contains(message, "membership_takeover_unavailable") || strings.Contains(message, "wrong number of arguments for 'subscribe' command") +} + type clusterNode struct { IP string `json:"ip"` Port int `json:"port"` @@ -127,6 +136,15 @@ type subscriptionCloser interface { Close() error } +type recoveryState uint32 + +const ( + recoveryStateStable recoveryState = iota + recoveryStateTakeoverEligible + recoveryStateSwitching + recoveryStateSwitchingTakeover +) + type Client struct { mu sync.Mutex @@ -139,12 +157,15 @@ type Client struct { sub *redis.Client release *redis.Client connections map[*homeDispatchConn]struct{} + closing chan struct{} lifecycle config.CredentialConcurrencyConfig limiter atomic.Pointer[config.CredentialConcurrencyConfig] managed bool heartbeatOK atomic.Bool dispatchFenced atomic.Bool + ambiguousDispatch atomic.Bool + recoveryState atomic.Uint32 clusterNodes []clusterNode reconnectFailures int } @@ -164,13 +185,15 @@ func (c *Client) NewLifetime() *Client { } c.mu.Lock() defer c.mu.Unlock() - return &Client{ + next := &Client{ homeCfg: c.homeCfg, seedHost: c.seedHost, seedPort: c.seedPort, clusterNodes: append([]clusterNode(nil), c.clusterNodes...), reconnectFailures: c.reconnectFailures, } + next.recoveryState.Store(c.recoveryState.Load()) + return next } func (c *Client) Enabled() bool { @@ -203,11 +226,15 @@ func (c *Client) Close() { commandClient, subscriptionClient, connections := c.detachClientsLocked() releaseClient := c.release c.release = nil + closing := c.closing c.mu.Unlock() closeDetachedClients(commandClient, subscriptionClient, connections) if releaseClient != nil { _ = releaseClient.Close() } + if closing != nil { + <-closing + } } // closeBootstrapPools replaces private bootstrap pools without ending the client lifetime. @@ -227,6 +254,7 @@ func (c *Client) AbortAmbiguousDispatch() { if c == nil { return } + c.ambiguousDispatch.Store(true) c.dispatchFenced.Store(true) c.heartbeatOK.Store(false) c.mu.Lock() @@ -254,6 +282,21 @@ func (c *Client) AbortAmbiguousDispatch() { } } +// AmbiguousDispatch reports whether this lifetime observed an issued dispatch with an unknown delivery result. +func (c *Client) AmbiguousDispatch() bool { + return c != nil && c.ambiguousDispatch.Load() +} + +// SuppressTakeover forces the next subscriber lifetime through normal membership recovery. +func (c *Client) SuppressTakeover() { + if c == nil { + return + } + if !c.recoveryState.CompareAndSwap(uint32(recoveryStateTakeoverEligible), uint32(recoveryStateStable)) { + c.recoveryState.CompareAndSwap(uint32(recoveryStateSwitchingTakeover), uint32(recoveryStateSwitching)) + } +} + func (c *Client) detachClientsLocked() (*redis.Client, *redis.Client, []*homeDispatchConn) { connections := make([]*homeDispatchConn, 0, len(c.connections)) for conn := range c.connections { @@ -284,7 +327,14 @@ func (c *Client) closeClientsLocked() { commandClient, subscriptionClient, connections := c.detachClientsLocked() releaseClient := c.release c.release = nil + previousClosing := c.closing + done := make(chan struct{}) + c.closing = done go func() { + defer close(done) + if previousClosing != nil { + <-previousClosing + } closeDetachedClients(commandClient, subscriptionClient, connections) if releaseClient != nil { _ = releaseClient.Close() @@ -292,6 +342,25 @@ func (c *Client) closeClientsLocked() { }() } +func (c *Client) waitForClientsClosed() { + for { + c.mu.Lock() + closing := c.closing + c.mu.Unlock() + if closing == nil { + return + } + <-closing + c.mu.Lock() + if c.closing == closing { + c.closing = nil + c.mu.Unlock() + return + } + c.mu.Unlock() + } +} + // SetManagedLifetime defers client shutdown to the Service lifetime owner. func (c *Client) SetManagedLifetime(managed bool) { if c == nil { @@ -341,6 +410,7 @@ func (c *Client) ensureClients() error { if !c.Enabled() { return ErrDisabled } + c.waitForClientsClosed() c.mu.Lock() defer c.mu.Unlock() if c.dispatchFenced.Load() { @@ -685,6 +755,9 @@ func (c *Client) switchToNodeLocked(node clusterNode) bool { } c.homeCfg.Host = host c.homeCfg.Port = node.Port + if !c.recoveryState.CompareAndSwap(uint32(recoveryStateStable), uint32(recoveryStateSwitching)) { + c.recoveryState.CompareAndSwap(uint32(recoveryStateTakeoverEligible), uint32(recoveryStateSwitchingTakeover)) + } c.closeClientsLocked() return true } @@ -1233,6 +1306,10 @@ func (c *Client) concurrencyReleaseClient() (*redis.Client, error) { if c == nil || c.dispatchFenced.Load() { return nil, ErrDispatchFenced } + state := recoveryState(c.recoveryState.Load()) + if state == recoveryStateSwitching || state == recoveryStateSwitchingTakeover { + return nil, ErrNotConnected + } if !c.Enabled() { return nil, ErrDisabled } @@ -1242,6 +1319,10 @@ func (c *Client) concurrencyReleaseClient() (*redis.Client, error) { if c.dispatchFenced.Load() { return nil, ErrDispatchFenced } + state = recoveryState(c.recoveryState.Load()) + if state == recoveryStateSwitching || state == recoveryStateSwitchingTakeover { + return nil, ErrNotConnected + } if c.release != nil { return c.release, nil } @@ -1525,13 +1606,21 @@ func (c *Client) subscriptionParameters() ([]string, time.Duration) { args := []string{redisChannelConfig} if cfg.LifecycleConfigRevision > 0 { args = append(args, strconv.FormatInt(cfg.LifecycleConfigRevision, 10)) + state := recoveryState(c.recoveryState.Load()) + if state == recoveryStateTakeoverEligible || state == recoveryStateSwitchingTakeover { + args = append(args, "takeover") + } } return args, cfg.CPAHeartbeatTimeout } func (c *Client) rebuildCommandPoolAndProbe(ctx context.Context) error { c.promoteSubscription() - return c.Ping(ctx) + if errPing := c.Ping(ctx); errPing != nil { + return errPing + } + c.recoveryState.Store(uint32(recoveryStateStable)) + return nil } func (c *Client) promoteSubscription() { @@ -1641,6 +1730,11 @@ func (c *Client) RunConfigSubscriberLifetime(ctx context.Context, onConfig func( event, errReceive := pubsub.ReceiveTimeout(ctx, receiveTimeout) if errReceive != nil { if ctx.Err() == nil { + if c.heartbeatOK.Load() { + if !c.recoveryState.CompareAndSwap(uint32(recoveryStateStable), uint32(recoveryStateTakeoverEligible)) { + c.recoveryState.CompareAndSwap(uint32(recoveryStateSwitching), uint32(recoveryStateSwitchingTakeover)) + } + } if isTimeoutError(errReceive) { c.markSubscriptionTimeout() } else { diff --git a/internal/home/client_test.go b/internal/home/client_test.go index 133a6ee61..b6ea29845 100644 --- a/internal/home/client_test.go +++ b/internal/home/client_test.go @@ -217,6 +217,75 @@ func TestNewLifetimePreservesClusterFailoverState(t *testing.T) { } } +func TestEnsureClientsWaitsForPreviousTargetClose(t *testing.T) { + client := New(config.HomeConfig{Enabled: true, Host: "next.example.com", Port: 8327}) + closing := make(chan struct{}) + client.closing = closing + done := make(chan error, 1) + go func() { + done <- client.ensureClients() + }() + + select { + case errEnsure := <-done: + t.Fatalf("ensureClients() returned before previous target closed: %v", errEnsure) + case <-time.After(20 * time.Millisecond): + } + close(closing) + select { + case errEnsure := <-done: + if errEnsure != nil { + t.Fatal(errEnsure) + } + case <-time.After(time.Second): + t.Fatal("ensureClients() did not continue after previous target closed") + } + client.Close() +} + +func TestConcurrencyReleaseDoesNotOpenSwitchingTarget(t *testing.T) { + client := New(config.HomeConfig{Enabled: true, Host: "next.example.com", Port: 8327}) + client.recoveryState.Store(uint32(recoveryStateSwitching)) + errRelease := client.PushConcurrencyRelease(context.Background(), ConcurrencyReleaseFrame{CredentialID: "cred-a", Model: "model-a", ReleaseSeq: 1}) + if !errors.Is(errRelease, ErrNotConnected) { + t.Fatalf("PushConcurrencyRelease() error = %v, want %v", errRelease, ErrNotConnected) + } + client.mu.Lock() + releaseClient := client.release + client.mu.Unlock() + if releaseClient != nil { + t.Fatal("release client was opened before the switched target became ready") + } +} + +func TestAmbiguousDispatchSuppressesTakeoverForNextLifetime(t *testing.T) { + client := New(config.HomeConfig{Enabled: true, Host: "next.example.com", Port: 8327}) + client.recoveryState.Store(uint32(recoveryStateSwitchingTakeover)) + client.AbortAmbiguousDispatch() + if !client.AmbiguousDispatch() { + t.Fatal("ambiguous dispatch was not recorded") + } + client.SuppressTakeover() + next := client.NewLifetime() + if got := recoveryState(next.recoveryState.Load()); got != recoveryStateSwitching { + t.Fatalf("next recovery state = %d, want %d", got, recoveryStateSwitching) + } +} + +func TestMembershipTakeoverUnavailableError(t *testing.T) { + for _, message := range []string{ + "ERR membership_takeover_unavailable", + "ERR wrong number of arguments for 'subscribe' command", + } { + if !IsMembershipTakeoverUnavailableError(errors.New(message)) { + t.Fatalf("takeover unavailable error %q was not recognized", message) + } + } + if IsMembershipTakeoverUnavailableError(errors.New("ERR connection refused")) { + t.Fatal("unrelated error was recognized as takeover unavailable") + } +} + func TestBuildKVSetArgs(t *testing.T) { args, errArgs := buildKVSetArgs("key", []byte("value"), KVSetOptions{EX: 2 * time.Second, NX: true}) if errArgs != nil { @@ -1105,6 +1174,12 @@ func TestConfigSubscriberUsesAppliedLifecycleRevisionAndRebuildsCommands(t *test if timeout != 4*time.Second { t.Fatalf("receive timeout = %s", timeout) } + client.recoveryState.Store(uint32(recoveryStateSwitchingTakeover)) + args, _ = client.subscriptionParameters() + if !reflect.DeepEqual(args, []string{"config", "9", "takeover"}) { + t.Fatalf("takeover subscribe args = %#v", args) + } + client.recoveryState.Store(uint32(recoveryStateStable)) client.promoteSubscription() client.mu.Lock() commandClient := client.cmd @@ -1177,8 +1252,9 @@ func TestRunConfigSubscriberLifetimeReturnsAfterHeartbeatLoss(t *testing.T) { client.mu.Lock() client.clusterNodes = []clusterNode{{IP: "failover.example.com", Port: 8327}} client.mu.Unlock() + client.recoveryState.Store(uint32(recoveryStateSwitchingTakeover)) - ready := make(chan struct{}, 1) + ready := make(chan bool, 1) errRun := client.RunConfigSubscriberLifetime(context.Background(), func(raw []byte) error { parsed, errParse := config.ParseConfigBytes(raw) if errParse != nil { @@ -1188,12 +1264,15 @@ func TestRunConfigSubscriberLifetimeReturnsAfterHeartbeatLoss(t *testing.T) { return errSet } return nil - }, func() { ready <- struct{}{} }) + }, func() { ready <- recoveryState(client.recoveryState.Load()) == recoveryStateStable }) if errRun == nil { t.Fatal("RunConfigSubscriberLifetime() error = nil after heartbeat loss") } select { - case <-ready: + case cleared := <-ready: + if !cleared { + t.Fatal("successful subscription ACK and command probe did not clear takeover state") + } default: t.Fatalf("RunConfigSubscriberLifetime() did not invoke onReady after subscription ACK: %v; commands=%#v", errRun, commands.All()) } @@ -1203,6 +1282,9 @@ func TestRunConfigSubscriberLifetimeReturnsAfterHeartbeatLoss(t *testing.T) { if got, _ := client.addr(); got != "failover.example.com:8327" { t.Fatalf("addr() = %q, want failover.example.com:8327 after heartbeat timeout", got) } + if got := recoveryState(client.recoveryState.Load()); got != recoveryStateSwitchingTakeover { + t.Fatalf("recovery state = %d, want %d", got, recoveryStateSwitchingTakeover) + } client.mu.Lock() commandClient, subscriptionClient := client.cmd, client.sub client.mu.Unlock() @@ -1215,8 +1297,8 @@ func TestRunConfigSubscriberLifetimeReturnsAfterHeartbeatLoss(t *testing.T) { if count := commands.CountCommandKey("SUBSCRIBE", redisChannelConfig); count != 1 { t.Fatalf("SUBSCRIBE config count = %d, want 1", count) } - if got := findRedisCommand(commands.All(), "SUBSCRIBE"); !reflect.DeepEqual(got, []string{"subscribe", "config", "1"}) { - t.Fatalf("SUBSCRIBE wire command = %#v, want []string{\"subscribe\", \"config\", \"1\"}", got) + if got := findRedisCommand(commands.All(), "SUBSCRIBE"); !reflect.DeepEqual(got, []string{"subscribe", "config", "1", "takeover"}) { + t.Fatalf("SUBSCRIBE wire command = %#v, want []string{\"subscribe\", \"config\", \"1\", \"takeover\"}", got) } } diff --git a/internal/home/concurrency_release.go b/internal/home/concurrency_release.go index 1717bcd8d..160aeae6d 100644 --- a/internal/home/concurrency_release.go +++ b/internal/home/concurrency_release.go @@ -61,6 +61,17 @@ func (f *releaseFlusher) SetConfigProvider(provider func() internalconfig.Creden f.signal() } +// SetSender replaces the Home lifetime used for subsequent release attempts. +func (f *releaseFlusher) SetSender(send func(context.Context, ConcurrencyReleaseFrame) error) { + if f == nil { + return + } + f.mu.Lock() + f.send = send + f.mu.Unlock() + f.signal() +} + // MarkDirty records the latest cumulative sequence for one release group and // returns a ticket completed when Home acknowledges that sequence. func (f *releaseFlusher) MarkDirty(group executionregistry.ReleaseGroup, sequence int64) *executionregistry.ReleaseTicket { @@ -174,11 +185,12 @@ func (f *releaseFlusher) timings() releaseFlusherTimings { } func (f *releaseFlusher) flush(ctx context.Context) bool { - if f == nil || f.send == nil { + if f == nil { return false } f.mu.Lock() + send := f.send pending := make(map[executionregistry.ReleaseGroup]int64, len(f.groups)) for group, state := range f.groups { if state.Latest > state.Acked { @@ -186,10 +198,13 @@ func (f *releaseFlusher) flush(ctx context.Context) bool { } } f.mu.Unlock() + if send == nil { + return false + } failed := false for group, sequence := range pending { - errSend := f.send(ctx, ConcurrencyReleaseFrame{ + errSend := send(ctx, ConcurrencyReleaseFrame{ CredentialID: group.CredentialID, Model: group.Model, ReleaseSeq: sequence, diff --git a/internal/home/concurrency_release_test.go b/internal/home/concurrency_release_test.go index 094d22d09..984ca5bca 100644 --- a/internal/home/concurrency_release_test.go +++ b/internal/home/concurrency_release_test.go @@ -474,3 +474,32 @@ func TestScopeEndBlocksDrainUntilReleaseSinkFlushesFinalSequence(t *testing.T) { t.Fatalf("final flushed sequence = %d, want 1", got) } } + +func TestReleaseFlusherSenderReplacementPreservesTicket(t *testing.T) { + flusher := newReleaseFlusher(time.Hour, time.Hour, func(context.Context, ConcurrencyReleaseFrame) error { + return errors.New("old Home unavailable") + }) + group := executionregistry.ReleaseGroup{CredentialID: "cred-1", Model: "gpt"} + ticket := flusher.MarkDirty(group, 1) + if ticket == nil { + t.Fatal("MarkDirty() ticket = nil") + } + if failed := flusher.flush(context.Background()); !failed { + t.Fatal("old sender release attempt did not fail") + } + + flusher.SetSender(func(_ context.Context, frame ConcurrencyReleaseFrame) error { + if frame.CredentialID != group.CredentialID || frame.Model != group.Model || frame.ReleaseSeq != 1 { + t.Fatalf("replacement sender frame = %#v", frame) + } + return nil + }) + if failed := flusher.flush(context.Background()); failed { + t.Fatal("replacement sender release attempt failed") + } + waitCtx, cancelWait := context.WithTimeout(context.Background(), time.Second) + defer cancelWait() + if errWait := ticket.Wait(waitCtx); errWait != nil { + t.Fatalf("ticket did not survive sender replacement: %v", errWait) + } +} diff --git a/sdk/cliproxy/executionregistry/registry.go b/sdk/cliproxy/executionregistry/registry.go index 9275a1cd8..adc688740 100644 --- a/sdk/cliproxy/executionregistry/registry.go +++ b/sdk/cliproxy/executionregistry/registry.go @@ -154,6 +154,30 @@ func (r *Registry) BeginDispatch() (*PendingDispatch, error) { return pending, nil } +// WaitPending waits until every dispatch with an unresolved Home response has ended or been installed. +func (r *Registry) WaitPending(ctx context.Context) error { + if r == nil { + return ErrRegistryClosed + } + if ctx == nil { + ctx = context.Background() + } + + r.mu.Lock() + for len(r.pending) != 0 { + changed := r.changed + r.mu.Unlock() + select { + case <-ctx.Done(): + return ctx.Err() + case <-changed: + } + r.mu.Lock() + } + r.mu.Unlock() + return nil +} + // End releases a dispatch token that was not installed. func (p *PendingDispatch) End() { if p == nil || p.registry == nil { diff --git a/sdk/cliproxy/executionregistry/registry_test.go b/sdk/cliproxy/executionregistry/registry_test.go index a1fb7c8aa..4595abfda 100644 --- a/sdk/cliproxy/executionregistry/registry_test.go +++ b/sdk/cliproxy/executionregistry/registry_test.go @@ -91,6 +91,42 @@ func TestDrainWaitsForPendingDispatch(t *testing.T) { } } +func TestWaitPendingDoesNotDrainActiveScope(t *testing.T) { + registry := New() + activePending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scope, errInstall := registry.Install(activePending, ScopeSpec{}) + if errInstall != nil { + t.Fatal(errInstall) + } + defer scope.End("test cleanup") + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + done := make(chan error, 1) + go func() { done <- registry.WaitPending(ctx) }() + select { + case errWait := <-done: + t.Fatalf("WaitPending() returned before pending dispatch ended: %v", errWait) + case <-time.After(20 * time.Millisecond): + } + pending.End() + if errWait := <-done; errWait != nil { + t.Fatalf("WaitPending() error = %v", errWait) + } + nextPending, errNext := registry.BeginDispatch() + if errNext != nil { + t.Fatalf("WaitPending() stopped registry acceptance: %v", errNext) + } + nextPending.End() +} + func TestDrainReturnsWhenBlockingResourceCloseExceedsContext(t *testing.T) { registry := New() pending, errBegin := registry.BeginDispatch() diff --git a/sdk/cliproxy/service_executionregistry_test.go b/sdk/cliproxy/service_executionregistry_test.go index 8c85077cf..8219d939d 100644 --- a/sdk/cliproxy/service_executionregistry_test.go +++ b/sdk/cliproxy/service_executionregistry_test.go @@ -1250,7 +1250,7 @@ func TestServiceExplicitReplacementCancelsRunWhenDrainTimesOut(t *testing.T) { scope.End("test cleanup") } -func TestServiceReplacesRegistryOnlyAfterNewSubscriptionAck(t *testing.T) { +func TestServiceKeepsRegistryAcrossHeartbeatFailoverAndExposesOnlyAfterNewACK(t *testing.T) { listener, errListen := net.Listen("tcp", "127.0.0.1:0") if errListen != nil { t.Fatalf("listen: %v", errListen) @@ -1325,15 +1325,15 @@ func TestServiceReplacesRegistryOnlyAfterNewSubscriptionAck(t *testing.T) { close(allowSecondAck) secondRegistry := waitForServiceRegistry(t, service, time.Second) - if secondRegistry == firstRegistry { - t.Fatal("replacement subscription reused the old registry") + if secondRegistry != firstRegistry { + t.Fatal("heartbeat failover replaced the execution registry") } if home.Current() == nil { t.Fatal("replacement client was not exposed after the replacement ACK") } } -func TestServiceDrainsBeforePreAckRetriesAndExposesOnlyAfterNewAck(t *testing.T) { +func TestServicePreservesActiveScopeDuringPreACKFailoverRetries(t *testing.T) { listener, errListen := net.Listen("tcp", "127.0.0.1:0") if errListen != nil { t.Fatalf("listen: %v", errListen) @@ -1393,16 +1393,9 @@ func TestServiceDrainsBeforePreAckRetriesAndExposesOnlyAfterNewAck(t *testing.T) close(loseFirst) select { case <-resourceClosed: - case <-time.After(time.Second): - t.Fatal("heartbeat loss did not close the active scope resource") - } - select { - case <-preAckAttempts: - t.Fatal("pre-ACK retry started before the active scope owner ended") + t.Fatal("heartbeat failover drained the active scope") case <-time.After(50 * time.Millisecond): } - scope.End("canceled") - firstPreAck := <-preAckAttempts secondPreAck := <-preAckAttempts if retryDelay := secondPreAck.Sub(firstPreAck); retryDelay < 75*time.Millisecond { @@ -1423,12 +1416,13 @@ func TestServiceDrainsBeforePreAckRetriesAndExposesOnlyAfterNewAck(t *testing.T) close(allowFinalAck) secondRegistry := waitForServiceRegistry(t, service, time.Second) - if secondRegistry == firstRegistry || home.Current() == nil { + if secondRegistry != firstRegistry || home.Current() == nil { t.Fatal("new Home lifetime was not exposed only after its subscription ACK") } + scope.End("completed") } -func TestServiceCancelsRunWhenBlockingScopeExceedsDrainBound(t *testing.T) { +func TestServiceHeartbeatFailoverDoesNotDrainBlockingScope(t *testing.T) { listener, errListen := net.Listen("tcp", "127.0.0.1:0") if errListen != nil { t.Fatalf("listen: %v", errListen) @@ -1507,30 +1501,243 @@ func TestServiceCancelsRunWhenBlockingScopeExceedsDrainBound(t *testing.T) { close(loseFirst) select { case <-started: - case <-time.After(time.Second): - t.Fatal("drain did not start closing the blocking scope") + t.Fatal("heartbeat failover started draining the blocking scope") + case <-time.After(50 * time.Millisecond): } select { case <-secondSubscribe: - t.Fatal("new subscription started before the old registry drained") - case <-time.After(50 * time.Millisecond): + case <-time.After(time.Second): + t.Fatal("new subscription did not start while the old scope remained active") } service.homeMu.Lock() exposedRegistry := service.homeRegistry service.homeMu.Unlock() if exposedRegistry != nil { - t.Fatal("new registry was exposed before the old registry drained") + t.Fatal("registry was exposed before the replacement ACK") + } + close(allowSecondAck) + if nextRegistry := waitForServiceRegistry(t, service, time.Second); nextRegistry != registry { + t.Fatal("heartbeat failover replaced the registry containing the active scope") } select { case <-serviceCtx.Done(): - case <-time.After(time.Second): - t.Fatal("service run was not canceled after drain timeout") + t.Fatal("heartbeat failover canceled the service run") + case <-time.After(50 * time.Millisecond): } close(release) scope.End("test cleanup") } +func TestServiceShutdownDrainsDetachedRegistryDuringRetry(t *testing.T) { + listener, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatalf("listen: %v", errListen) + } + firstAck := make(chan struct{}) + loseFirst := make(chan struct{}) + secondSubscribe := make(chan struct{}) + var secondSubscribeOnce sync.Once + allowSecondAck := make(chan struct{}) + stop := make(chan struct{}) + var subscriptionMu sync.Mutex + subscriptions := 0 + serverDone := make(chan struct{}) + go func() { + defer close(serverDone) + for { + conn, errAccept := listener.Accept() + if errAccept != nil { + return + } + go serveRegistryTestHomeConnection(conn, &subscriptionMu, &subscriptions, firstAck, loseFirst, secondSubscribe, &secondSubscribeOnce, allowSecondAck, stop) + } + }() + t.Cleanup(func() { + close(stop) + _ = listener.Close() + <-serverDone + home.ClearCurrent() + }) + + service := newRegistryTestService(t, listener) + serviceCtx, cancelService := context.WithCancel(context.Background()) + t.Cleanup(cancelService) + service.homeMu.Lock() + service.runCancel = cancelService + service.homeMu.Unlock() + service.startHomeSubscriber(serviceCtx) + + select { + case <-firstAck: + case <-time.After(time.Second): + t.Fatal("first subscription was not acknowledged") + } + registry := waitForServiceRegistry(t, service, time.Second) + pendingRetry, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + pendingScope, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scope, errInstall := registry.Install(pendingScope, executionregistry.ScopeSpec{}) + if errInstall != nil { + t.Fatal(errInstall) + } + resourceClosed := make(chan struct{}) + if errBind := scope.Bind(func() error { + close(resourceClosed) + go scope.End("shutdown") + return nil + }); errBind != nil { + t.Fatal(errBind) + } + t.Cleanup(func() { + pendingRetry.End() + scope.End("test cleanup") + }) + + service.homeMu.Lock() + client := service.homeClient + service.homeMu.Unlock() + if client == nil { + t.Fatal("ready Home client is unavailable") + } + close(loseFirst) + deadline := time.After(time.Second) + for { + errRelease := client.PushConcurrencyRelease(context.Background(), home.ConcurrencyReleaseFrame{CredentialID: "cred-a", Model: "model-a", ReleaseSeq: 1}) + if errors.Is(errRelease, home.ErrDispatchFenced) { + break + } + select { + case <-deadline: + t.Fatal("subscriber retry did not close the previous Home client") + case <-time.After(time.Millisecond): + } + } + + shutdownDone := make(chan error, 1) + go func() { + shutdownDone <- service.Shutdown(context.Background()) + }() + pendingRetry.End() + + select { + case <-resourceClosed: + case <-time.After(time.Second): + t.Fatal("shutdown did not drain the detached execution registry") + } + select { + case errShutdown := <-shutdownDone: + if errShutdown != nil { + t.Fatalf("Shutdown() error = %v", errShutdown) + } + case <-time.After(time.Second): + t.Fatal("Shutdown() did not complete after draining the detached registry") + } +} + +func TestServiceAmbiguousDispatchDrainsRegistryBeforeRetry(t *testing.T) { + listener, errListen := net.Listen("tcp", "127.0.0.1:0") + if errListen != nil { + t.Fatalf("listen: %v", errListen) + } + firstAck := make(chan struct{}) + loseFirst := make(chan struct{}) + secondSubscribe := make(chan struct{}) + var secondSubscribeOnce sync.Once + allowSecondAck := make(chan struct{}) + stop := make(chan struct{}) + var subscriptionMu sync.Mutex + subscriptions := 0 + serverDone := make(chan struct{}) + go func() { + defer close(serverDone) + for { + conn, errAccept := listener.Accept() + if errAccept != nil { + return + } + go serveRegistryTestHomeConnection(conn, &subscriptionMu, &subscriptions, firstAck, loseFirst, secondSubscribe, &secondSubscribeOnce, allowSecondAck, stop) + } + }() + t.Cleanup(func() { + close(stop) + _ = listener.Close() + <-serverDone + home.ClearCurrent() + }) + + service := newRegistryTestService(t, listener) + serviceCtx, cancelService := context.WithCancel(context.Background()) + t.Cleanup(cancelService) + service.homeMu.Lock() + service.runCancel = cancelService + service.homeMu.Unlock() + service.startHomeSubscriber(serviceCtx) + + select { + case <-firstAck: + case <-time.After(time.Second): + t.Fatal("first subscription was not acknowledged") + } + registry := waitForServiceRegistry(t, service, time.Second) + pending, errBegin := registry.BeginDispatch() + if errBegin != nil { + t.Fatal(errBegin) + } + scope, errInstall := registry.Install(pending, executionregistry.ScopeSpec{}) + if errInstall != nil { + t.Fatal(errInstall) + } + resourceClosed := make(chan struct{}) + if errBind := scope.Bind(func() error { + close(resourceClosed) + go scope.End("ambiguous dispatch") + return nil + }); errBind != nil { + t.Fatal(errBind) + } + + service.homeMu.Lock() + client := service.homeClient + service.homeMu.Unlock() + if client == nil { + t.Fatal("ready Home client is unavailable") + } + client.AbortAmbiguousDispatch() + select { + case <-resourceClosed: + case <-time.After(time.Second): + t.Fatal("ambiguous dispatch did not drain the active registry") + } + select { + case <-secondSubscribe: + case <-time.After(time.Second): + t.Fatal("subscriber did not retry after ambiguous dispatch drain") + } + service.homeMu.Lock() + exposedRegistry := service.homeRegistry + service.homeMu.Unlock() + if exposedRegistry != nil { + t.Fatal("replacement registry was exposed before its subscription ACK") + } + + close(allowSecondAck) + nextRegistry := waitForServiceRegistry(t, service, time.Second) + if nextRegistry == registry { + t.Fatal("ambiguous dispatch reused the drained execution registry") + } + select { + case <-serviceCtx.Done(): + t.Fatal("successful ambiguous dispatch recovery canceled the service run") + case <-time.After(50 * time.Millisecond): + } +} + func TestServiceBacksOffAfterRepeatedPreAckFailures(t *testing.T) { listener, errListen := net.Listen("tcp", "127.0.0.1:0") if errListen != nil { @@ -1585,7 +1792,7 @@ func TestServiceBacksOffAfterRepeatedPreAckFailures(t *testing.T) { } } -func TestServiceHeartbeatLossCancelsBlockedConfigFinalizationBeforeDrain(t *testing.T) { +func TestServiceHeartbeatLossCancelsBlockedConfigFinalizationWithoutDrainingRegistry(t *testing.T) { listener, errListen := net.Listen("tcp", "127.0.0.1:0") if errListen != nil { t.Fatalf("listen: %v", errListen) @@ -1648,8 +1855,8 @@ func TestServiceHeartbeatLossCancelsBlockedConfigFinalizationBeforeDrain(t *test } select { case <-resourceClosed: - case <-time.After(500 * time.Millisecond): - t.Fatal("heartbeat loss did not cancel the worker and drain the active execution") + t.Fatal("heartbeat loss drained the active execution") + case <-time.After(200 * time.Millisecond): } select { case <-secondConfig: @@ -1663,6 +1870,7 @@ func TestServiceHeartbeatLossCancelsBlockedConfigFinalizationBeforeDrain(t *test if currentRegistry != nil || currentClient != nil || home.Current() != nil { t.Fatal("heartbeat-lost lifetime left a published Home client or registry") } + scope.End("completed") } func TestServiceConfigWorkerFinalizesRapidUpdatesInOrder(t *testing.T) { diff --git a/sdk/cliproxy/service_home.go b/sdk/cliproxy/service_home.go index 10f8fdf92..d27b2f64e 100644 --- a/sdk/cliproxy/service_home.go +++ b/sdk/cliproxy/service_home.go @@ -492,6 +492,25 @@ func (s *Service) runHomeSubscriber(homeCtx context.Context, parentCtx context.C }() var previousClient *home.Client + registry := executionregistry.New() + cancelBound := atomic.Int64{} + cancelBound.Store(int64(internalconfig.CredentialConcurrencyConfig{}.WithDefaults().CPACancelBound)) + releaseFlusher := home.NewReleaseFlusher(nil, nil) + registry.SetReleaseSink(releaseFlusher.MarkDirty) + defer func() { + registry.SetReleaseSink(nil) + drainBound := time.Duration(cancelBound.Load()) + if drainBound <= 0 { + drainBound = internalconfig.CredentialConcurrencyConfig{}.WithDefaults().CPACancelBound + } + drainCtx, cancelDrain := context.WithTimeout(context.WithoutCancel(parentCtx), drainBound) + errDrain := registry.Drain(drainCtx) + cancelDrain() + if errDrain != nil && !errors.Is(errDrain, executionregistry.ErrRegistryClosed) && parentCtx.Err() == nil { + log.WithError(errDrain).Error("failed to drain detached Home execution registry") + s.cancelServiceRun() + } + }() for homeCtx.Err() == nil { supervisor.setPublisherCompletion(nil) client := previousClient @@ -501,18 +520,15 @@ func (s *Service) runHomeSubscriber(homeCtx context.Context, parentCtx context.C client = client.NewLifetime() } client.SetManagedLifetime(true) - registry := executionregistry.New() releaseCtx, releaseCancel := context.WithCancel(context.WithoutCancel(homeCtx)) - releaseFlusher := home.NewReleaseFlusher(client.LimiterConfig, client.PushConcurrencyRelease) - registry.SetReleaseSink(releaseFlusher.MarkDirty) + releaseFlusher.SetConfigProvider(client.LimiterConfig) + releaseFlusher.SetSender(client.PushConcurrencyRelease) releaseDone := make(chan struct{}) go func() { defer close(releaseDone) releaseFlusher.Run(releaseCtx) }() lifetimeCtx, lifetimeCancel := context.WithCancel(homeCtx) - cancelBound := atomic.Int64{} - cancelBound.Store(int64(internalconfig.CredentialConcurrencyConfig{}.WithDefaults().CPACancelBound)) queue := newHomeConfigWorkQueue() ready := make(chan struct{}) var readyOnce sync.Once @@ -552,6 +568,44 @@ func (s *Service) runHomeSubscriber(homeCtx context.Context, parentCtx context.C } s.detachHomeSubscriberLifetime(client, registry) + retry := errRun != nil && homeCtx.Err() == nil + if retry { + releaseCancel() + <-releaseDone + client.Close() + + settleBound := time.Duration(cancelBound.Load()) + 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() + return + } + if client.AmbiguousDispatch() || home.IsMembershipTakeoverUnavailableError(errRun) { + registry.SetReleaseSink(nil) + drainCtx, cancelDrain := context.WithTimeout(context.WithoutCancel(parentCtx), settleBound) + errDrain := registry.Drain(drainCtx) + cancelDrain() + if errDrain != nil { + log.WithError(errDrain).Error("failed to drain Home executions after unsafe subscriber replacement") + s.cancelServiceRun() + return + } + client.SuppressTakeover() + registry = executionregistry.New() + releaseFlusher = home.NewReleaseFlusher(nil, nil) + registry.SetReleaseSink(releaseFlusher.MarkDirty) + } + log.WithError(errRun).Warn("home config subscription lifetime ended") + if !published.Load() && !waitForHomeSubscriberRetry(homeCtx, homeSubscriberPreAckRetryBackoff) { + return + } + previousClient = client + continue + } + drainBound := time.Duration(cancelBound.Load()) drainCtx, cancelDrain := context.WithTimeout(context.WithoutCancel(parentCtx), drainBound) errDrain := registry.Drain(drainCtx) @@ -577,13 +631,7 @@ func (s *Service) runHomeSubscriber(homeCtx context.Context, parentCtx context.C } return } - if errRun != nil && homeCtx.Err() == nil { - log.WithError(errRun).Warn("home config subscription lifetime ended") - } - if !published.Load() && errRun != nil && !waitForHomeSubscriberRetry(homeCtx, homeSubscriberPreAckRetryBackoff) { - return - } - previousClient = client + return } }