From 63274fa2c7ff065437de54d74aa3afeded5d7bf2 Mon Sep 17 00:00:00 2001
From: Gabe Joseph <gjoseph92@gmail.com>
Date: Fri, 30 Sep 2022 01:56:16 -0600
Subject: [PATCH] fix `_next_in_linear_chain`

could have traversed into widely-shared deps
---
 distributed/scheduler.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/distributed/scheduler.py b/distributed/scheduler.py
index 989e2df18a..69491e7224 100644
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -8296,15 +8296,15 @@ def _next_in_linear_chain(ts: TaskState, cutoff: int) -> TaskState | None:
 
     # Check if this is part of a linear chain:
     # exactly 1 dependency, excluding widely-shared tasks.
-    non_widely_shared = 0
+    next: TaskState | None = None
     for dts in ts.dependencies:
         if len(dts.dependents) > cutoff:  # widely-shared; ignore it
             continue
-        if non_widely_shared:
+        if next:
             return None
-        non_widely_shared += 1
+        next = dts
 
-    return next(iter(ts.dependents))
+    return next
 
 
 def family(