Fix reuse of various steps (#482)

* Fix `StreamTarget` and `KafkaTarget` init on flow reuse

[ML-5351](https://jira.iguazeng.com/browse/ML-5351)

* Fix `AsyncEmitSource` reuse

* Add regression test

* Move table creation to _init for proper step reuse

* Proper SyncEmitSource init in case of flow reuse (not a fix)

* Add test, overhaul closeables

* Revert "Add test, overhaul closeables"

This reverts commit da272dd.

* Revert changes related to table init

integration/test_flow_integration.py

@@ -261,6 +261,46 @@ def test_write_to_v3io_stream(assign_stream_teardown_test):
     asyncio.run(async_test_write_to_v3io_stream(assign_stream_teardown_test))
 
 
+async def async_test_write_to_v3io_stream_flow_reuse(stream_path, flow, iteration):
+
+    controller = flow.run()
+    for i in range(10):
+        await controller.emit(i)
+
+    await asyncio.sleep(5)
+
+    try:
+        shard0_data = await GetShardData().get_shard_data(f"{stream_path}/0")
+        assert shard0_data == [b"0", b"2", b"4", b"6", b"8"] * iteration
+        shard1_data = await GetShardData().get_shard_data(f"{stream_path}/1")
+        assert shard1_data == [b"1", b"3", b"5", b"7", b"9"] * iteration
+    finally:
+        await controller.terminate()
+        await controller.await_termination()
+
+
+# ML-5351
+def test_write_to_v3io_stream_flow_reuse(assign_stream_teardown_test):
+    stream_path = assign_stream_teardown_test
+    flow = build_flow(
+        [
+            AsyncEmitSource(),
+            Map(lambda x: str(x)),
+            StreamTarget(
+                V3ioDriver(),
+                stream_path,
+                sharding_func=lambda event: int(event.body),
+                batch_size=8,
+                shards=2,
+                full_event=False,
+            ),
+        ]
+    )
+
+    asyncio.run(async_test_write_to_v3io_stream_flow_reuse(stream_path, flow, 1))
+    asyncio.run(async_test_write_to_v3io_stream_flow_reuse(stream_path, flow, 2))
+
+
 async def async_test_write_to_v3io_stream_full_event_readback(
     setup_stream_teardown_test,
 ):

storey/sources.py

@@ -283,18 +283,19 @@ def __init__(
         super().__init__(**kwargs)
         if buffer_size <= 0:
             raise ValueError("Buffer size must be positive")
-        self._q = queue.Queue(buffer_size)
+        self._buffer_size = buffer_size
         self._key_field = key_field
         self._max_events_before_commit = max_events_before_commit or 20000
         self._max_time_before_commit = max_time_before_commit or 45
         self._max_wait_before_commit = max_wait_before_commit or 5
         self._explicit_ack = explicit_ack
-        self._termination_q = queue.Queue(1)
         self._ex = None
         self._closeables = []
 
     def _init(self):
         super()._init()
+        self._q = queue.Queue(self._buffer_size)
+        self._termination_q = queue.Queue(1)
         self._is_terminated = False
         self._outstanding_offsets = defaultdict(list)
 
@@ -583,12 +584,11 @@ def __init__(
     ):
         super().__init__(**kwargs)
         if buffer_size is None:
-            buffer_size = 8
+            self._buffer_size = 8
         elif buffer_size <= 0:
             raise ValueError("Buffer size must be positive")
         else:
-            kwargs["buffer_size"] = buffer_size
-        self._q = asyncio.Queue(buffer_size)
+            kwargs["buffer_size"] = self._buffer_size
         self._key_field = key_field
         self._max_events_before_commit = max_events_before_commit or 20000
         self._max_time_before_commit = max_time_before_commit or 45
@@ -601,6 +601,7 @@ def _init(self):
         super()._init()
         self._is_terminated = False
         self._outstanding_offsets = defaultdict(list)
+        self._q = asyncio.Queue(self._buffer_size)
 
     async def _run_loop(self):
         committer = None

storey/targets.py

@@ -839,13 +839,13 @@ def __init__(
 
         self._shards = shards
         self._retention_period_hours = retention_period_hours
-        self._initialized = False
 
         self._full_event = full_event
 
     def _init(self):
         Flow._init(self)
         _Writer._init(self)
+        self._initialized = False
         self._worker_exited = False
 
     @staticmethod
@@ -1048,13 +1048,12 @@ def __init__(
         Flow.__init__(self, **kwargs)
         _Writer.__init__(self, columns, infer_columns_from_data, retain_dict=True)
 
-        self._initialized = False
-
         self._full_event = full_event
 
     def _init(self):
         Flow._init(self)
         _Writer._init(self)
+        self._initialized = False
 
     async def _lazy_init(self):
         from kafka import KafkaProducer