fix: stop scheduling new futures after shutdown

Author: ringoldsdev

laygo/transformers/strategies/process.py

@@ -57,19 +57,40 @@ def _ordered_generator(
   ) -> Iterator[list[Out]]:
     """Generate results in their original order."""
     futures = deque()
+
+    # Pre-submit initial batch of futures
     for _ in range(self.max_workers + 1):
       try:
         chunk = next(chunks_iter)
         futures.append(executor.submit(_worker_process_chunk, transformer, context_handle, chunk))
       except StopIteration:
         break
+      except RuntimeError as e:
+        if "cannot schedule new futures after shutdown" in str(e):
+          break
+        raise
+
     while futures:
-      yield futures.popleft().result()
       try:
-        chunk = next(chunks_iter)
-        futures.append(executor.submit(_worker_process_chunk, transformer, context_handle, chunk))
-      except StopIteration:
-        continue
+        yield futures.popleft().result()
+
+        # Try to submit the next chunk
+        try:
+          chunk = next(chunks_iter)
+          futures.append(executor.submit(_worker_process_chunk, transformer, context_handle, chunk))
+        except StopIteration:
+          continue
+        except RuntimeError as e:
+          if "cannot schedule new futures after shutdown" in str(e):
+            # Executor is shut down, stop submitting new work
+            break
+          raise
+      except Exception:
+        # Cancel remaining futures and re-raise
+        for future in futures:
+          future.cancel()
+        futures.clear()
+        raise
 
   def _unordered_generator(
     self,
@@ -79,16 +100,38 @@ def _unordered_generator(
     context_handle: IContextHandle,
   ) -> Iterator[list[Out]]:
     """Generate results as they complete."""
-    futures = {
-      executor.submit(_worker_process_chunk, transformer, context_handle, chunk)
-      for chunk in itertools.islice(chunks_iter, self.max_workers + 1)
-    }
+    futures = set()
+
+    # Pre-submit initial batch
+    try:
+      for chunk in itertools.islice(chunks_iter, self.max_workers + 1):
+        futures.add(executor.submit(_worker_process_chunk, transformer, context_handle, chunk))
+    except RuntimeError as e:
+      if "cannot schedule new futures after shutdown" in str(e):
+        # If we can't submit any futures, there's nothing to process
+        return
+      raise
+
     while futures:
-      done, futures = wait(futures, return_when=FIRST_COMPLETED)
-      for future in done:
-        yield future.result()
-        try:
-          chunk = next(chunks_iter)
-          futures.add(executor.submit(_worker_process_chunk, transformer, context_handle, chunk))
-        except StopIteration:
-          continue
+      try:
+        done, futures = wait(futures, return_when=FIRST_COMPLETED)
+        for future in done:
+          yield future.result()
+
+          # Try to submit next chunk if available
+          try:
+            chunk = next(chunks_iter)
+            futures.add(executor.submit(_worker_process_chunk, transformer, context_handle, chunk))
+          except StopIteration:
+            continue
+          except RuntimeError as e:
+            if "cannot schedule new futures after shutdown" in str(e):
+              # Executor is shut down, stop submitting new work
+              break
+            raise
+      except Exception:
+        # Cancel remaining futures and re-raise
+        for future in futures:
+          future.cancel()
+        futures.clear()
+        raise