From aff761df35c401ba673d33b6c9775f68da637ae9 Mon Sep 17 00:00:00 2001 From: ChickenBenny Date: Sat, 24 Oct 2026 11:54:15 +0800 Subject: [PATCH 1/3] refactor(runs): track watch subscribers as subscription structs Signed-off-by: ChickenBenny --- runs/repository/impl/action.go | 102 +++++++++++----------------- runs/repository/impl/action_test.go | 26 +++---- 2 files changed, 50 insertions(+), 78 deletions(-) diff --git a/runs/repository/impl/action.go b/runs/repository/impl/action.go index 4abc9b4d69..9525ab1b8a 100644 --- a/runs/repository/impl/action.go +++ b/runs/repository/impl/action.go @@ -124,8 +124,8 @@ type actionRepo struct { listenerMinReconnectInterval time.Duration // Subscriber management for LISTEN/NOTIFY - runSubscribers map[chan string]bool - actionSubscribers map[chan string]bool + runSubscribers map[*subscription]struct{} + actionSubscribers map[*subscription]struct{} mu sync.RWMutex notifyMu sync.Mutex @@ -158,8 +158,8 @@ func NewActionRepo( repo := &actionRepo{ db: db, dsn: dsn, - runSubscribers: make(map[chan string]bool), - actionSubscribers: make(map[chan string]bool), + runSubscribers: make(map[*subscription]struct{}), + actionSubscribers: make(map[*subscription]struct{}), notificationConfig: notificationConfig, notificationMetrics: notificationMetrics, } @@ -780,27 +780,37 @@ func (r *actionRepo) WatchStateUpdates(ctx context.Context, updates chan<- *comm <-ctx.Done() } -// WatchRunUpdates watches for run updates via LISTEN/NOTIFY -func (r *actionRepo) WatchRunUpdates(ctx context.Context, runID *common.RunIdentifier, updates chan<- *models.Run, errs chan<- error) { - runKey := fmt.Sprintf("%s/%s/%s", runID.Project, runID.Domain, runID.Name) - notifCh := make(chan string, r.notificationConfig.watchBufferSize) +// subscription is one Watch* subscriber. +type subscription struct { + notify chan string +} +func (r *actionRepo) subscribe(subs map[*subscription]struct{}) *subscription { + sub := &subscription{notify: make(chan string, r.notificationConfig.watchBufferSize)} r.mu.Lock() - r.runSubscribers[notifCh] = true + subs[sub] = struct{}{} r.mu.Unlock() + return sub +} - defer func() { - r.mu.Lock() - delete(r.runSubscribers, notifCh) - close(notifCh) - r.mu.Unlock() - }() +func (r *actionRepo) unsubscribe(subs map[*subscription]struct{}, sub *subscription) { + r.mu.Lock() + delete(subs, sub) + close(sub.notify) + r.mu.Unlock() +} + +// WatchRunUpdates watches for run updates via LISTEN/NOTIFY +func (r *actionRepo) WatchRunUpdates(ctx context.Context, runID *common.RunIdentifier, updates chan<- *models.Run, errs chan<- error) { + runKey := fmt.Sprintf("%s/%s/%s", runID.Project, runID.Domain, runID.Name) + sub := r.subscribe(r.runSubscribers) + defer r.unsubscribe(r.runSubscribers, sub) for { select { case <-ctx.Done(): return - case notifPayload := <-notifCh: + case notifPayload := <-sub.notify: if notifPayload == runKey { run, err := r.GetRun(ctx, runID) if err != nil { @@ -822,24 +832,14 @@ func (r *actionRepo) WatchRunUpdates(ctx context.Context, runID *common.RunIdent // WatchAllRunUpdates watches for all run updates via LISTEN/NOTIFY func (r *actionRepo) WatchAllRunUpdates(ctx context.Context, updates chan<- *models.Run, errs chan<- error) { - notifCh := make(chan string, r.notificationConfig.watchBufferSize) - - r.mu.Lock() - r.runSubscribers[notifCh] = true - r.mu.Unlock() - - defer func() { - r.mu.Lock() - delete(r.runSubscribers, notifCh) - close(notifCh) - r.mu.Unlock() - }() + sub := r.subscribe(r.runSubscribers) + defer r.unsubscribe(r.runSubscribers, sub) for { select { case <-ctx.Done(): return - case notifPayload := <-notifCh: + case notifPayload := <-sub.notify: parts := strings.Split(notifPayload, "/") if len(parts) != 3 { logger.Warnf(ctx, "Invalid run notification payload: %s", notifPayload) @@ -875,24 +875,14 @@ func (r *actionRepo) WatchAllRunUpdates(ctx context.Context, updates chan<- *mod // WatchAllActionUpdates watches for all action updates for a run via LISTEN/NOTIFY func (r *actionRepo) WatchAllActionUpdates(ctx context.Context, runID *common.RunIdentifier, updates chan<- *models.Action, errs chan<- error) { runPrefix := fmt.Sprintf("%s/%s/%s/", runID.Project, runID.Domain, runID.Name) - notifCh := make(chan string, r.notificationConfig.watchBufferSize) - - r.mu.Lock() - r.actionSubscribers[notifCh] = true - r.mu.Unlock() - - defer func() { - r.mu.Lock() - delete(r.actionSubscribers, notifCh) - close(notifCh) - r.mu.Unlock() - }() + sub := r.subscribe(r.actionSubscribers) + defer r.unsubscribe(r.actionSubscribers, sub) for { select { case <-ctx.Done(): return - case notifPayload := <-notifCh: + case notifPayload := <-sub.notify: // Collect this notification plus any others already queued, // deduplicating by action name so we only query the DB once // per action during a burst of updates. @@ -904,7 +894,7 @@ func (r *actionRepo) WatchAllActionUpdates(ctx context.Context, runID *common.Ru drain: for { select { - case extra := <-notifCh: + case extra := <-sub.notify: if len(extra) > len(runPrefix) && extra[:len(runPrefix)] == runPrefix { pending[extra[len(runPrefix):]] = true } @@ -943,24 +933,14 @@ func (r *actionRepo) WatchAllActionUpdates(ctx context.Context, runID *common.Ru func (r *actionRepo) WatchActionUpdates(ctx context.Context, actionID *common.ActionIdentifier, updates chan<- *models.Action, errs chan<- error) { targetPayload := fmt.Sprintf("%s/%s/%s/%s", actionID.Run.Project, actionID.Run.Domain, actionID.Run.Name, actionID.Name) - notifCh := make(chan string, r.notificationConfig.watchBufferSize) - - r.mu.Lock() - r.actionSubscribers[notifCh] = true - r.mu.Unlock() - - defer func() { - r.mu.Lock() - delete(r.actionSubscribers, notifCh) - close(notifCh) - r.mu.Unlock() - }() + sub := r.subscribe(r.actionSubscribers) + defer r.unsubscribe(r.actionSubscribers, sub) for { select { case <-ctx.Done(): return - case notifPayload := <-notifCh: + case notifPayload := <-sub.notify: if notifPayload != targetPayload { continue } @@ -969,7 +949,7 @@ func (r *actionRepo) WatchActionUpdates(ctx context.Context, actionID *common.Ac for { // Prevent continuous update from the action select { - case extra := <-notifCh: + case extra := <-sub.notify: if extra != targetPayload { continue } @@ -1054,9 +1034,9 @@ func (r *actionRepo) processNotifications() { case "run_updates": // Broadcast to all run subscribers r.mu.RLock() - for ch := range r.runSubscribers { + for sub := range r.runSubscribers { select { - case ch <- notif.Extra: + case sub.notify <- notif.Extra: default: // Channel full, skip this subscriber logger.Warnf(context.Background(), "Run subscriber channel full, dropping notification") @@ -1067,9 +1047,9 @@ func (r *actionRepo) processNotifications() { case "action_updates": // Broadcast to all action subscribers r.mu.RLock() - for ch := range r.actionSubscribers { + for sub := range r.actionSubscribers { select { - case ch <- notif.Extra: + case sub.notify <- notif.Extra: default: logger.Warnf(context.Background(), "Action subscriber channel full, dropping notification payload=%s", notif.Extra) } diff --git a/runs/repository/impl/action_test.go b/runs/repository/impl/action_test.go index 712418a3e7..d9ff42ffe5 100644 --- a/runs/repository/impl/action_test.go +++ b/runs/repository/impl/action_test.go @@ -84,8 +84,8 @@ func newNotifyRepoWithDB(t *testing.T) (*actionRepo, *sql.DB, *sql.Conn) { r := &actionRepo{ db: db, dsn: database.GetPostgresDsn(context.Background(), testDbConfig.Postgres), - runSubscribers: make(map[chan string]bool), - actionSubscribers: make(map[chan string]bool), + runSubscribers: make(map[*subscription]struct{}), + actionSubscribers: make(map[*subscription]struct{}), pendingActions: make(map[string]struct{}, defaultNotificationBufferSize), pendingActionQueue: make([]string, 0, defaultNotificationBufferSize), pendingRuns: make(map[string]struct{}, defaultNotificationBufferSize), @@ -1256,8 +1256,8 @@ func TestProcessNotifications_SurvivesPingFailureDuringOutage(t *testing.T) { dsn: database.GetPostgresDsn(ctx, testDbConfig.Postgres) + " application_name=" + appName, listenerPingInterval: 100 * time.Millisecond, listenerMinReconnectInterval: 100 * time.Millisecond, - runSubscribers: make(map[chan string]bool), - actionSubscribers: make(map[chan string]bool), + runSubscribers: make(map[*subscription]struct{}), + actionSubscribers: make(map[*subscription]struct{}), } require.NoError(t, r.startPostgresListener()) t.Cleanup(func() { r.listener.Close() }) //nolint:errcheck @@ -1304,16 +1304,16 @@ func TestProcessNotifications_SurvivesPingFailureDuringOutage(t *testing.T) { // subscribeActions registers a raw subscriber the way WatchActionUpdates does. func subscribeActions(t *testing.T, r *actionRepo, size int) chan string { t.Helper() - ch := make(chan string, size) + sub := &subscription{notify: make(chan string, size)} r.mu.Lock() - r.actionSubscribers[ch] = true + r.actionSubscribers[sub] = struct{}{} r.mu.Unlock() t.Cleanup(func() { r.mu.Lock() - delete(r.actionSubscribers, ch) + delete(r.actionSubscribers, sub) r.mu.Unlock() }) - return ch + return sub.notify } func notifyTestActionID(name string) *common.ActionIdentifier { @@ -1435,15 +1435,7 @@ func TestNotifyPump_ConcurrentWritersDeliverEveryAction(t *testing.T) { // Subscribe the way WatchActionUpdates does, so we observe what actually // came back through pg_notify and the listener. - delivered := make(chan string, 4*total) - repo.mu.Lock() - repo.actionSubscribers[delivered] = true - repo.mu.Unlock() - defer func() { - repo.mu.Lock() - delete(repo.actionSubscribers, delivered) - repo.mu.Unlock() - }() + delivered := subscribeActions(t, repo, 4*total) ctx := context.Background() var wg sync.WaitGroup From a94be3e45f81b7bde391d01fc9d59d7a17d742b0 Mon Sep 17 00:00:00 2001 From: ChickenBenny Date: Sat, 24 Oct 2026 11:54:33 +0800 Subject: [PATCH 2/3] fix(runs): resync watch subscribers after the listener reconnects Signed-off-by: ChickenBenny --- runs/repository/impl/action.go | 82 +++++++++---- runs/repository/impl/action_test.go | 184 ++++++++++++++++++++++++---- 2 files changed, 216 insertions(+), 50 deletions(-) diff --git a/runs/repository/impl/action.go b/runs/repository/impl/action.go index 9525ab1b8a..1215499d95 100644 --- a/runs/repository/impl/action.go +++ b/runs/repository/impl/action.go @@ -783,10 +783,23 @@ func (r *actionRepo) WatchStateUpdates(ctx context.Context, updates chan<- *comm // subscription is one Watch* subscriber. type subscription struct { notify chan string + // resync asks the subscriber to re-read its state from the DB after the listener reconnects. + resync chan struct{} +} + +// requestResync never blocks, and repeated requests coalesce into one re-read. +func (s *subscription) requestResync() { + select { + case s.resync <- struct{}{}: + default: + } } func (r *actionRepo) subscribe(subs map[*subscription]struct{}) *subscription { - sub := &subscription{notify: make(chan string, r.notificationConfig.watchBufferSize)} + sub := &subscription{ + notify: make(chan string, r.notificationConfig.watchBufferSize), + resync: make(chan struct{}, 1), + } r.mu.Lock() subs[sub] = struct{}{} r.mu.Unlock() @@ -810,28 +823,32 @@ func (r *actionRepo) WatchRunUpdates(ctx context.Context, runID *common.RunIdent select { case <-ctx.Done(): return + case <-sub.resync: case notifPayload := <-sub.notify: - if notifPayload == runKey { - run, err := r.GetRun(ctx, runID) - if err != nil { - if ctx.Err() != nil { - return - } - errs <- err - return - } - select { - case updates <- run: - case <-ctx.Done(): - return - } + if notifPayload != runKey { + continue } } + + run, err := r.GetRun(ctx, runID) + if err != nil { + if ctx.Err() != nil { + return + } + errs <- err + return + } + select { + case updates <- run: + case <-ctx.Done(): + return + } } } // WatchAllRunUpdates watches for all run updates via LISTEN/NOTIFY func (r *actionRepo) WatchAllRunUpdates(ctx context.Context, updates chan<- *models.Run, errs chan<- error) { + // TODO(#8193): resync after a listener reconnect. sub := r.subscribe(r.runSubscribers) defer r.unsubscribe(r.runSubscribers, sub) @@ -875,6 +892,7 @@ func (r *actionRepo) WatchAllRunUpdates(ctx context.Context, updates chan<- *mod // WatchAllActionUpdates watches for all action updates for a run via LISTEN/NOTIFY func (r *actionRepo) WatchAllActionUpdates(ctx context.Context, runID *common.RunIdentifier, updates chan<- *models.Action, errs chan<- error) { runPrefix := fmt.Sprintf("%s/%s/%s/", runID.Project, runID.Domain, runID.Name) + // TODO(#8193): resync after a listener reconnect. sub := r.subscribe(r.actionSubscribers) defer r.unsubscribe(r.actionSubscribers, sub) @@ -940,6 +958,7 @@ func (r *actionRepo) WatchActionUpdates(ctx context.Context, actionID *common.Ac select { case <-ctx.Done(): return + case <-sub.resync: case notifPayload := <-sub.notify: if notifPayload != targetPayload { continue @@ -957,21 +976,21 @@ func (r *actionRepo) WatchActionUpdates(ctx context.Context, actionID *common.Ac break drain } } + } - action, err := r.GetAction(ctx, actionID) - if err != nil { - if ctx.Err() != nil { - return - } - logger.Errorf(ctx, "Failed to get action from notification: %v", err) - continue - } - - select { - case updates <- action: - case <-ctx.Done(): + action, err := r.GetAction(ctx, actionID) + if err != nil { + if ctx.Err() != nil { return } + logger.Errorf(ctx, "Failed to get action from notification: %v", err) + continue + } + + select { + case updates <- action: + case <-ctx.Done(): + return } } } @@ -1027,6 +1046,15 @@ func (r *actionRepo) processNotifications() { } if notif == nil { // pq sends nil after it reconnects; notifications sent while it was down may be lost. + logger.Infof(context.Background(), "PostgreSQL listener reconnected, resyncing subscribers") + r.mu.RLock() + for sub := range r.runSubscribers { + sub.requestResync() + } + for sub := range r.actionSubscribers { + sub.requestResync() + } + r.mu.RUnlock() continue } diff --git a/runs/repository/impl/action_test.go b/runs/repository/impl/action_test.go index d9ff42ffe5..a2285239dc 100644 --- a/runs/repository/impl/action_test.go +++ b/runs/repository/impl/action_test.go @@ -1246,6 +1246,32 @@ func TestRunNotifyLoop_RetriesUndeliveredPayloads(t *testing.T) { } // A listener ping that fails during a Postgres outage must not stop notification delivery once pq reconnects. +// dropListenerConnection keeps new connections to the test database refused and terminates the +// connection named appName. The returned func lets connections back in. +func dropListenerConnection(ctx context.Context, t *testing.T, appName string) (restore func()) { + t.Helper() + // A database cannot disallow its own connections, so run the outage SQL from the maintenance database. + adminCfg := testDbConfig.Postgres + adminCfg.DbName = "postgres" + admin, err := sql.Open("postgres", database.GetPostgresDsn(ctx, adminCfg)) + require.NoError(t, err) + t.Cleanup(func() { admin.Close() }) //nolint:errcheck + allow := func(allowed bool) error { + _, err := admin.ExecContext(ctx, fmt.Sprintf("ALTER DATABASE %q WITH ALLOW_CONNECTIONS %t", testDbConfig.Postgres.DbName, allowed)) + return err + } + + require.NoError(t, allow(false)) + t.Cleanup(func() { _ = allow(true) }) + var terminated int + require.NoError(t, admin.QueryRowContext(ctx, + "SELECT count(*) FROM (SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE application_name = $1) t", + appName).Scan(&terminated)) + require.Equal(t, 1, terminated, "expected to terminate exactly the listener connection") + + return func() { require.NoError(t, allow(true)) } +} + func TestProcessNotifications_SurvivesPingFailureDuringOutage(t *testing.T) { const appName = "flyte_listener_outage_test" db := setupActionDB(t) @@ -1263,30 +1289,10 @@ func TestProcessNotifications_SurvivesPingFailureDuringOutage(t *testing.T) { t.Cleanup(func() { r.listener.Close() }) //nolint:errcheck updates := subscribeActions(t, r, 16) - // A database cannot disallow its own connections, so run the outage SQL from the maintenance database. - adminCfg := testDbConfig.Postgres - adminCfg.DbName = "postgres" - admin, err := sql.Open("postgres", database.GetPostgresDsn(ctx, adminCfg)) - require.NoError(t, err) - t.Cleanup(func() { admin.Close() }) //nolint:errcheck - dbName := testDbConfig.Postgres.DbName - - // Keep pq's reconnect attempts failing, then drop the LISTEN connection. - _, err = admin.ExecContext(ctx, fmt.Sprintf("ALTER DATABASE %q WITH ALLOW_CONNECTIONS false", dbName)) - require.NoError(t, err) - t.Cleanup(func() { - _, _ = admin.ExecContext(ctx, fmt.Sprintf("ALTER DATABASE %q WITH ALLOW_CONNECTIONS true", dbName)) - }) - var terminated int - require.NoError(t, admin.QueryRowContext(ctx, - "SELECT count(*) FROM (SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE application_name = $1) t", - appName).Scan(&terminated)) - require.Equal(t, 1, terminated, "expected to terminate exactly the listener connection") - + restore := dropListenerConnection(ctx, t, appName) // Hold the outage across several ping ticks so Ping fails while pq has no connection. time.Sleep(10 * r.listenerPingInterval) - _, err = admin.ExecContext(ctx, fmt.Sprintf("ALTER DATABASE %q WITH ALLOW_CONNECTIONS true", dbName)) - require.NoError(t, err) + restore() // pq retries with backoff; keep notifying until the reader forwards one. require.Eventually(t, func() bool { @@ -1301,10 +1307,142 @@ func TestProcessNotifications_SurvivesPingFailureDuringOutage(t *testing.T) { }, 10*time.Second, 100*time.Millisecond, "notifications must reach subscribers after the listener reconnects") } +// A nil notification (pq reconnected) must reach every subscriber, even one with a full buffer. +func TestProcessNotifications_ResyncsSubscribersOnReconnect(t *testing.T) { + notify := make(chan *pq.Notification) + r := &actionRepo{ + listener: &pq.Listener{Notify: notify}, + listenerPingInterval: time.Hour, + runSubscribers: make(map[*subscription]struct{}), + actionSubscribers: make(map[*subscription]struct{}), + notificationConfig: NewNotificationConfig(testNotificationBufferLimit, notifyRetryMinBackoff, notifyRetryMaxBackoff, 1), + } + runSub := r.subscribe(r.runSubscribers) + actionSub := r.subscribe(r.actionSubscribers) + actionSub.notify <- "fills-the-buffer" + + done := make(chan struct{}) + go func() { + r.processNotifications() + close(done) + }() + notify <- nil + notify <- nil + close(notify) + <-done + + assert.Len(t, runSub.resync, 1, "run subscriber must be asked to resync") + assert.Len(t, actionSub.resync, 1, "a full notify buffer must not drop the resync") +} + +func TestWatchActionUpdates_ResyncSendsCurrentState(t *testing.T) { + db := setupActionDB(t) + repoIface, err := NewActionRepo(db, testDbConfig, testNotificationConfig) + require.NoError(t, err) + r := repoIface.(*actionRepo) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + actionID := notifyTestActionID("resync-direct") + _, err = r.CreateAction(ctx, models.NewActionModel(actionID), false) + require.NoError(t, err) + + updates := make(chan *models.Action, 1) + go r.WatchActionUpdates(ctx, actionID, updates, make(chan error, 1)) + require.Eventually(t, func() bool { + r.mu.RLock() + defer r.mu.RUnlock() + return len(r.actionSubscribers) == 1 + }, 5*time.Second, 10*time.Millisecond) + + r.mu.RLock() + for sub := range r.actionSubscribers { + sub.requestResync() + } + r.mu.RUnlock() + + select { + case got := <-updates: + assert.Equal(t, actionID.Name, got.Name) + case <-time.After(5 * time.Second): + t.Fatal("resync must re-read the action and send it") + } +} + +// An update whose NOTIFY is lost while the listener is down must still reach the watcher. +func TestWatchActionUpdates_RecoversNotificationLostDuringOutage(t *testing.T) { + const appName = "flyte_listener_resync_test" + db := setupActionDB(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + r := &actionRepo{ + db: db, + dsn: database.GetPostgresDsn(ctx, testDbConfig.Postgres) + " application_name=" + appName, + listenerPingInterval: 100 * time.Millisecond, + listenerMinReconnectInterval: 100 * time.Millisecond, + runSubscribers: make(map[*subscription]struct{}), + actionSubscribers: make(map[*subscription]struct{}), + pendingActions: make(map[string]struct{}, defaultNotificationBufferSize), + pendingActionQueue: make([]string, 0, defaultNotificationBufferSize), + pendingRuns: make(map[string]struct{}, defaultNotificationBufferSize), + pendingRunQueue: make([]string, 0, defaultNotificationBufferSize), + pendingCh: make(chan struct{}, 1), + notificationConfig: testNotificationConfig, + } + actionID := notifyTestActionID("resync-outage") + _, err := r.CreateAction(ctx, models.NewActionModel(actionID), false) + require.NoError(t, err) + + require.NoError(t, r.startPostgresListener()) + t.Cleanup(func() { r.listener.Close() }) //nolint:errcheck + + updates := make(chan *models.Action, 16) + go r.WatchActionUpdates(ctx, actionID, updates, make(chan error, 1)) + require.Eventually(t, func() bool { + r.mu.RLock() + defer r.mu.RUnlock() + return len(r.actionSubscribers) == 1 + }, 5*time.Second, 10*time.Millisecond) + + // Hold a pooled connection so the update below still works while new connections are refused. + writer, err := db.Conn(ctx) + require.NoError(t, err) + t.Cleanup(func() { writer.Close() }) //nolint:errcheck + + restore := dropListenerConnection(ctx, t, appName) + + // Update the action and notify while nobody is listening: this NOTIFY is lost. + succeeded := int32(common.ActionPhase_ACTION_PHASE_SUCCEEDED) + _, err = writer.ExecContext(ctx, + "UPDATE actions SET phase = $1 WHERE project = $2 AND domain = $3 AND run_name = $4 AND name = $5", + succeeded, actionID.Run.Project, actionID.Run.Domain, actionID.Run.Name, actionID.Name) + require.NoError(t, err) + payload := fmt.Sprintf("%s/%s/%s/%s", actionID.Run.Project, actionID.Run.Domain, actionID.Run.Name, actionID.Name) + _, err = writer.ExecContext(ctx, "SELECT pg_notify('action_updates', $1)", payload) + require.NoError(t, err) + + time.Sleep(5 * r.listenerPingInterval) + restore() + + // No further NOTIFY is sent: only the reconnect resync can deliver the new phase. + deadline := time.After(15 * time.Second) + for { + select { + case got := <-updates: + if got.Phase == succeeded { + return + } + case <-deadline: + t.Fatal("the watcher must see the update made while the listener was down") + } + } +} + // subscribeActions registers a raw subscriber the way WatchActionUpdates does. func subscribeActions(t *testing.T, r *actionRepo, size int) chan string { t.Helper() - sub := &subscription{notify: make(chan string, size)} + sub := &subscription{notify: make(chan string, size), resync: make(chan struct{}, 1)} r.mu.Lock() r.actionSubscribers[sub] = struct{}{} r.mu.Unlock() From 3ccb3667d3e6d0457433d25a3c2d47f993fd42fb Mon Sep 17 00:00:00 2001 From: ChickenBenny Date: Sat, 24 Oct 2026 12:11:29 +0800 Subject: [PATCH 3/3] fix(runs): resync list watchers after the listener reconnects Signed-off-by: ChickenBenny --- runs/repository/impl/action.go | 14 +++++- runs/repository/impl/action_test.go | 42 ++++++++++++++++++ runs/repository/interfaces/action.go | 1 + runs/service/run_service.go | 50 ++++++++++++++++----- runs/service/run_service_test.go | 65 ++++++++++++++++++++++++++++ 5 files changed, 158 insertions(+), 14 deletions(-) diff --git a/runs/repository/impl/action.go b/runs/repository/impl/action.go index 1215499d95..8341119842 100644 --- a/runs/repository/impl/action.go +++ b/runs/repository/impl/action.go @@ -848,7 +848,6 @@ func (r *actionRepo) WatchRunUpdates(ctx context.Context, runID *common.RunIdent // WatchAllRunUpdates watches for all run updates via LISTEN/NOTIFY func (r *actionRepo) WatchAllRunUpdates(ctx context.Context, updates chan<- *models.Run, errs chan<- error) { - // TODO(#8193): resync after a listener reconnect. sub := r.subscribe(r.runSubscribers) defer r.unsubscribe(r.runSubscribers, sub) @@ -856,6 +855,12 @@ func (r *actionRepo) WatchAllRunUpdates(ctx context.Context, updates chan<- *mod select { case <-ctx.Done(): return + case <-sub.resync: + select { + case updates <- nil: + case <-ctx.Done(): + return + } case notifPayload := <-sub.notify: parts := strings.Split(notifPayload, "/") if len(parts) != 3 { @@ -892,7 +897,6 @@ func (r *actionRepo) WatchAllRunUpdates(ctx context.Context, updates chan<- *mod // WatchAllActionUpdates watches for all action updates for a run via LISTEN/NOTIFY func (r *actionRepo) WatchAllActionUpdates(ctx context.Context, runID *common.RunIdentifier, updates chan<- *models.Action, errs chan<- error) { runPrefix := fmt.Sprintf("%s/%s/%s/", runID.Project, runID.Domain, runID.Name) - // TODO(#8193): resync after a listener reconnect. sub := r.subscribe(r.actionSubscribers) defer r.unsubscribe(r.actionSubscribers, sub) @@ -900,6 +904,12 @@ func (r *actionRepo) WatchAllActionUpdates(ctx context.Context, runID *common.Ru select { case <-ctx.Done(): return + case <-sub.resync: + select { + case updates <- nil: + case <-ctx.Done(): + return + } case notifPayload := <-sub.notify: // Collect this notification plus any others already queued, // deduplicating by action name so we only query the DB once diff --git a/runs/repository/impl/action_test.go b/runs/repository/impl/action_test.go index a2285239dc..0b9be619e5 100644 --- a/runs/repository/impl/action_test.go +++ b/runs/repository/impl/action_test.go @@ -1335,6 +1335,48 @@ func TestProcessNotifications_ResyncsSubscribersOnReconnect(t *testing.T) { assert.Len(t, actionSub.resync, 1, "a full notify buffer must not drop the resync") } +func TestWatchAllUpdates_SendNilOnResync(t *testing.T) { + r := &actionRepo{ + runSubscribers: make(map[*subscription]struct{}), + actionSubscribers: make(map[*subscription]struct{}), + notificationConfig: testNotificationConfig, + } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + runs := make(chan *models.Run, 1) + actions := make(chan *models.Action, 1) + go r.WatchAllRunUpdates(ctx, runs, make(chan error, 1)) + go r.WatchAllActionUpdates(ctx, notifyTestActionID("a1").Run, actions, make(chan error, 1)) + require.Eventually(t, func() bool { + r.mu.RLock() + defer r.mu.RUnlock() + return len(r.runSubscribers) == 1 && len(r.actionSubscribers) == 1 + }, 5*time.Second, 10*time.Millisecond) + + r.mu.RLock() + for sub := range r.runSubscribers { + sub.requestResync() + } + for sub := range r.actionSubscribers { + sub.requestResync() + } + r.mu.RUnlock() + + select { + case got := <-runs: + assert.Nil(t, got) + case <-time.After(5 * time.Second): + t.Fatal("WatchAllRunUpdates must send nil on resync") + } + select { + case got := <-actions: + assert.Nil(t, got) + case <-time.After(5 * time.Second): + t.Fatal("WatchAllActionUpdates must send nil on resync") + } +} + func TestWatchActionUpdates_ResyncSendsCurrentState(t *testing.T) { db := setupActionDB(t) repoIface, err := NewActionRepo(db, testDbConfig, testNotificationConfig) diff --git a/runs/repository/interfaces/action.go b/runs/repository/interfaces/action.go index 0672d6cb0f..f4eba089cd 100644 --- a/runs/repository/interfaces/action.go +++ b/runs/repository/interfaces/action.go @@ -40,6 +40,7 @@ type ActionRepo interface { // Watch operations (for streaming) WatchRunUpdates(ctx context.Context, runID *common.RunIdentifier, updates chan<- *models.Run, errs chan<- error) + // The WatchAll* watchers send nil when updates may have been lost; the caller should re-list. WatchAllRunUpdates(ctx context.Context, updates chan<- *models.Run, errs chan<- error) WatchAllActionUpdates(ctx context.Context, runID *common.RunIdentifier, updates chan<- *models.Action, errs chan<- error) WatchActionUpdates(ctx context.Context, actionID *common.ActionIdentifier, updates chan<- *models.Action, errs chan<- error) diff --git a/runs/service/run_service.go b/runs/service/run_service.go index 45feda7566..c641db5b57 100644 --- a/runs/service/run_service.go +++ b/runs/service/run_service.go @@ -1338,22 +1338,26 @@ func (s *RunService) WatchRuns( // Step 2: Send existing runs that match filter listInput := s.convertWatchRequestToListInput(req.Msg) - - runs, err := s.repo.ActionRepo().ListActions(ctx, listInput) - if err != nil { - logger.Errorf(ctx, "Failed to list runs: %v", err) - // Continue even if list fails - still watch for new updates - } else if len(runs) > 0 { + sendRuns := func() error { + runs, err := s.repo.ActionRepo().ListActions(ctx, listInput) + if err != nil { + logger.Errorf(ctx, "Failed to list runs: %v", err) + // Continue even if list fails - still watch for new updates + return nil + } + if len(runs) == 0 { + return nil + } protoRuns := make([]*workflow.Run, len(runs)) for i, run := range runs { protoRuns[i] = s.convertRunToProto(run) } - - if err := stream.Send(&workflow.WatchRunsResponse{ + return stream.Send(&workflow.WatchRunsResponse{ Runs: protoRuns, - }); err != nil { - return err - } + }) + } + if err := sendRuns(); err != nil { + return err } for { @@ -1364,6 +1368,13 @@ func (s *RunService) WatchRuns( logger.Errorf(ctx, "Error watching runs: %v", err) return err case run := <-updatesCh: + if run == nil { + // Updates may have been lost; send the list again. + if err := sendRuns(); err != nil { + return err + } + continue + } // Filter the run based on the watch request criteria if !s.runMatchesFilter(run, req.Msg) { continue @@ -1414,6 +1425,13 @@ func (s *RunService) WatchActions( if !ok { return nil } + if updated == nil { + // Updates may have been lost; re-list the run. + if err := s.listAndSendAllActions(ctx, runID, rsm, stream); err != nil { + logger.Errorf(ctx, "Failed to re-list actions: %v", err) + } + continue + } updates, err := rsm.upsertActions(ctx, []*models.Action{updated}) if err != nil { return err @@ -1468,7 +1486,15 @@ func (s *RunService) listAndSendAllActions( batch = batch[:pageSize] } - updates, err := rsm.upsertActions(ctx, batch) + // On a re-list, skip actions the stream already has: upserting an existing one marks it + // and its ancestors as changed even when nothing differs. + fresh := make([]*models.Action, 0, len(batch)) + for _, action := range batch { + if known := rsm.GetActionTreeNodeByName(action.Name); known == nil || !known.Action.UpdatedAt.Equal(action.UpdatedAt) { + fresh = append(fresh, action) + } + } + updates, err := rsm.upsertActions(ctx, fresh) if err != nil { return err } diff --git a/runs/service/run_service_test.go b/runs/service/run_service_test.go index 7696f804ba..c9beae9c21 100644 --- a/runs/service/run_service_test.go +++ b/runs/service/run_service_test.go @@ -2105,3 +2105,68 @@ func TestConvertNodeUpdateToEnrichedProto_ConditionMetadata(t *testing.T) { // type check before signaling. assert.Equal(t, core.SimpleType_STRING, condMeta.GetType().GetSimple()) } + +func TestWatchRuns_RelistsOnResync(t *testing.T) { + actionRepo, _, svc := newTestService(t) + client := newRunServiceTestClient(t, svc) + + run := func(phase common.ActionPhase) *models.Run { + return &models.Run{Project: "p", Domain: "d", RunName: "r1", Name: "a0", Phase: int32(phase)} + } + actionRepo.On("WatchAllRunUpdates", mock.Anything, mock.Anything, mock.Anything). + Run(func(args mock.Arguments) { + args.Get(1).(chan<- *models.Run) <- nil + <-args.Get(0).(context.Context).Done() + }).Once() + actionRepo.On("ListActions", mock.Anything, mock.Anything). + Return([]*models.Run{run(common.ActionPhase_ACTION_PHASE_RUNNING)}, nil).Once() + actionRepo.On("ListActions", mock.Anything, mock.Anything). + Return([]*models.Run{run(common.ActionPhase_ACTION_PHASE_SUCCEEDED)}, nil).Once() + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + stream, err := client.WatchRuns(ctx, connect.NewRequest(&workflow.WatchRunsRequest{})) + require.NoError(t, err) + + require.True(t, stream.Receive()) + assert.Equal(t, common.ActionPhase_ACTION_PHASE_RUNNING, stream.Msg().GetRuns()[0].GetAction().GetStatus().GetPhase()) + require.True(t, stream.Receive()) + assert.Equal(t, common.ActionPhase_ACTION_PHASE_SUCCEEDED, stream.Msg().GetRuns()[0].GetAction().GetStatus().GetPhase()) +} + +func TestWatchActions_RelistsOnResync(t *testing.T) { + actionRepo, _, svc := newTestService(t) + client := newRunServiceTestClient(t, svc) + + runID := &common.RunIdentifier{Project: "p", Domain: "d", Name: "r1"} + root := func(phase common.ActionPhase, updatedAt time.Time) *models.Action { + return &models.Action{Project: "p", Domain: "d", RunName: "r1", Name: "a0", Phase: int32(phase), UpdatedAt: updatedAt} + } + child := &models.Action{ + Project: "p", Domain: "d", RunName: "r1", Name: "a1", + ParentActionName: sql.NullString{String: "a0", Valid: true}, + Phase: int32(common.ActionPhase_ACTION_PHASE_RUNNING), + UpdatedAt: time.Unix(100, 0), + } + actionRepo.On("WatchAllActionUpdates", mock.Anything, matchRunID(runID), mock.Anything, mock.Anything). + Run(func(args mock.Arguments) { + args.Get(2).(chan<- *models.Action) <- nil + <-args.Get(0).(context.Context).Done() + }).Once() + actionRepo.On("ListActions", mock.Anything, mock.Anything). + Return([]*models.Action{root(common.ActionPhase_ACTION_PHASE_RUNNING, time.Unix(100, 0)), child}, nil).Once() + actionRepo.On("ListActions", mock.Anything, mock.Anything). + Return([]*models.Action{root(common.ActionPhase_ACTION_PHASE_SUCCEEDED, time.Unix(200, 0)), child}, nil).Once() + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + stream, err := client.WatchActions(ctx, connect.NewRequest(&workflow.WatchActionsRequest{RunId: runID})) + require.NoError(t, err) + + require.True(t, stream.Receive()) + assert.Len(t, stream.Msg().GetEnrichedActions(), 2) + require.True(t, stream.Receive()) + resent := stream.Msg().GetEnrichedActions() + require.Len(t, resent, 1, "the unchanged child must not be sent again") + assert.Equal(t, common.ActionPhase_ACTION_PHASE_SUCCEEDED, resent[0].GetAction().GetStatus().GetPhase()) +}