✨ feat: allow shutdown consume

Author: mic1on

src/usepy_plugin_redis/__init__.py

@@ -1,4 +1,5 @@
 import logging
+import threading
 
 import redis
 import time
@@ -65,8 +66,19 @@ def connection(self):
             del self.state.connection
 
 
-class RedisStreamMixin:
-    connection: redis.Redis
+class RedisStreamStore(RedisStore):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.__shutdown = False
+        self.__shutdown_event = threading.Event()
+
+    def __del__(self):
+        self.shutdown()
+
+    def shutdown(self):
+        self.__shutdown = True
+        self.__shutdown_event.set()
 
     def _create_group(self, stream, group):
         try:
@@ -90,8 +102,12 @@ def send(self, stream, message, **kwargs):
 
     def start_consuming(self, stream, group, consumer, callback, prefetch=1, **kwargs):
         """开始消费"""
+        self.__shutdown = False
+        self.__shutdown_event.clear()
         self._create_group(stream, group)
-        while True:
+        while not self.__shutdown:
+            if self.__shutdown:
+                break
             try:
                 messages = self.connection.xreadgroup(group, consumer, {stream: '>'}, count=prefetch, **kwargs)
                 for message in messages:
@@ -107,9 +123,5 @@ def start_consuming(self, stream, group, consumer, callback, prefetch=1, **kwarg
                 time.sleep(1)
 
 
-class RedisStreamStore(RedisStore, RedisStreamMixin):
-    pass
-
-
 useRedis = RedisStore
 useRedisStreamStore = RedisStreamStore