fix: Using queue instead of semaphore

Author: Bhargav Dodla

sdk/python/feast/infra/online_stores/contrib/cassandra_online_store/cassandra_online_store.py

@@ -19,6 +19,7 @@
 """
 
 import logging
+import queue
 import threading
 import time
 from datetime import datetime
@@ -399,24 +400,17 @@ def online_write_batch(
                       display progress.
         """
 
-        def on_success(result, semaphore, futures, future, lock):
-            semaphore.release()
-            with lock:
-                futures.remove(future)
+        def on_success(result, concurrent_queue):
+            concurrent_queue.get_nowait()
 
-        def on_failure(exc, semaphore, futures, future, lock):
-            semaphore.release()
-            with lock:
-                futures.remove(future)
+        def on_failure(exc, concurrent_queue):
+            concurrent_queue.get_nowait()
             logger.exception(f"Error writing a batch: {exc}")
             print(f"Error writing a batch: {exc}")
             raise Exception("Error writing a batch") from exc
 
         write_concurrency = config.online_store.write_concurrency
-
-        # rate_limiter = RateLimiter(rate=write_concurrency, interval=1)
-        semaphore = threading.Semaphore(write_concurrency)
-        lock = threading.Lock()
+        concurrent_queue: queue.Queue = queue.Queue(maxsize=write_concurrency)
 
         project = config.project
         ttl = (
@@ -428,7 +422,6 @@ def on_failure(exc, semaphore, futures, future, lock):
         keyspace: str = self._keyspace
         fqtable = CassandraOnlineStore._fq_table_name(keyspace, project, table)
 
-        futures = []
         insert_cql = self._get_cql_statement(
             config,
             "insert4",
@@ -451,58 +444,29 @@ def on_failure(exc, semaphore, futures, future, lock):
                     timestamp,
                 )
                 batch.add(insert_cql, params)
-            # with rate_limiter:
-            semaphore.acquire()
             future = session.execute_async(batch)
-            futures.append(future)
+            concurrent_queue.put(future)
             future.add_callbacks(
                 partial(
                     on_success,
-                    semaphore=semaphore,
-                    futures=futures,
-                    future=future,
-                    lock=lock,
+                    concurrent_queue=concurrent_queue,
                 ),
                 partial(
                     on_failure,
-                    semaphore=semaphore,
-                    futures=futures,
-                    future=future,
-                    lock=lock,
+                    concurrent_queue=concurrent_queue,
                 ),
             )
 
-            # TODO: Make this efficient by leveraging continuous writes rather
-            # than blocking until all writes are done. We may need to rate limit
-            # the writes to reduce the impact on read performance.
-            # if len(futures) >= write_concurrency:
-            #     # Raises exception if at least one of the batch fails
-            #     self._wait_for_futures(futures)
-            #     futures.clear()
-
             # this happens N-1 times, will be corrected outside:
             if progress:
                 progress(1)
-        while futures:
+        while not concurrent_queue.empty():
             time.sleep(0.001)
-        # if len(futures) > 0:
-        #     self._wait_for_futures(futures)
-        #     futures.clear()
 
         # correction for the last missing call to `progress`:
         if progress:
             progress(1)
 
-    def _wait_for_futures(self, futures):
-        try:
-            for future in futures:
-                future.result()
-            futures = []
-        except Exception as exc:
-            logger.error(f"Error writing a batch: {exc}")
-            print(f"Error writing a batch: {exc}")
-            raise Exception("Error writing a batch") from exc
-
     def online_read(
         self,
         config: RepoConfig,