From 7da9962efa675590f618e735eeba5165e290f80a Mon Sep 17 00:00:00 2001
From: quettabit <27509167+quettabit@users.noreply.github.com>
Date: Wed, 8 Apr 2026 01:28:26 -0600
Subject: [PATCH] initial commit

---
 src/s2_sdk/_producer.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/src/s2_sdk/_producer.py b/src/s2_sdk/_producer.py
index bbc0e26..c51bc90 100644
--- a/src/s2_sdk/_producer.py
+++ b/src/s2_sdk/_producer.py
@@ -112,11 +112,12 @@ async def close(self) -> None:
         if self._closed:
             return
         self._closed = True
-        await self._flush()
-        await self._session.close()
-        # Signal drain task to finish and wait for it
-        self._batch_ready.set()
-        await self._drain_task
+        try:
+            await self._flush()
+            await self._session.close()
+        finally:
+            self._batch_ready.set()
+            await self._drain_task
         if self._error is not None:
             raise self._error