fix: explicit shutdown

Author: ringoldsdev

laygo/transformers/strategies/process.py

@@ -57,38 +57,46 @@ def _ordered_generator(
   ) -> Iterator[list[Out]]:
     """Generate results in their original order."""
     futures = deque()
+    executor_shutdown = False
 
     # Pre-submit initial batch of futures
     for _ in range(self.max_workers + 1):
+      if executor_shutdown:
+        break
       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):
+          executor_shutdown = True
           break
         raise
 
     while futures:
       try:
         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
+        # Try to submit the next chunk only if executor is not shutdown
+        if not executor_shutdown:
+          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_shutdown = True
+              continue
+            raise
       except Exception:
         # Cancel remaining futures and re-raise
         for future in futures:
-          future.cancel()
+          try:
+            future.cancel()
+          except Exception:
+            pass  # Ignore cancellation errors
         futures.clear()
         raise
 
@@ -101,37 +109,44 @@ def _unordered_generator(
   ) -> Iterator[list[Out]]:
     """Generate results as they complete."""
     futures = set()
+    executor_shutdown = False
 
     # Pre-submit initial batch
-    try:
-      for chunk in itertools.islice(chunks_iter, self.max_workers + 1):
+    for chunk in itertools.islice(chunks_iter, self.max_workers + 1):
+      if executor_shutdown:
+        break
+      try:
         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
+      except RuntimeError as e:
+        if "cannot schedule new futures after shutdown" in str(e):
+          executor_shutdown = True
+          break
+        raise
 
     while futures:
       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
+          # Try to submit next chunk only if executor is not shutdown
+          if not executor_shutdown:
+            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_shutdown = True
+                continue
+              raise
       except Exception:
         # Cancel remaining futures and re-raise
         for future in futures:
-          future.cancel()
+          try:
+            future.cancel()
+          except Exception:
+            pass  # Ignore cancellation errors
         futures.clear()
         raise