11import json
22
3- import grpc
3+ from grpc import StreamStreamMultiCallable , RpcError , StatusCode # type: ignore
44
55from dapr .clients .exceptions import StreamInactiveError
66from dapr .clients .grpc ._response import TopicEventResponse
77from dapr .proto import api_v1 , appcallback_v1
88import queue
99import threading
10+ from typing import Optional
1011
1112
1213def success ():
@@ -28,11 +29,11 @@ def __init__(self, stub, pubsub_name, topic, metadata=None, dead_letter_topic=No
2829 self .topic = topic
2930 self .metadata = metadata or {}
3031 self .dead_letter_topic = dead_letter_topic or ''
31- self ._stream = None
32- self ._response_thread = None
33- self ._send_queue = queue .Queue ()
34- self ._receive_queue = queue .Queue ()
35- self ._stream_active = False
32+ self ._stream : Optional [ StreamStreamMultiCallable ] = None # Type annotation for gRPC stream
33+ self ._response_thread : Optional [ threading . Thread ] = None # Type for thread
34+ self ._send_queue : queue . Queue = queue .Queue () # Type annotation for send queue
35+ self ._receive_queue : queue . Queue = queue .Queue () # Type annotation for receive queue
36+ self ._stream_active : bool = False
3637 self ._stream_lock = threading .Lock () # Protects _stream_active
3738
3839 def start (self ):
@@ -55,9 +56,8 @@ def outgoing_request_iterator():
5556 # Start sending back acknowledgement messages from the send queue
5657 while self ._is_stream_active ():
5758 try :
58- response = self ._send_queue .get ()
59- # The above blocks until a message is available or the stream is closed
60- # so that's why we need to check again if the stream is still active
59+ response = self ._send_queue .get (timeout = 1 )
60+ # Check again if the stream is still active
6161 if not self ._is_stream_active ():
6262 break
6363 yield response
@@ -76,17 +76,19 @@ def outgoing_request_iterator():
7676
7777 def _handle_incoming_messages (self ):
7878 try :
79- # The first message dapr sends on the stream is for signalling only, so discard it
80- next (self ._stream )
81-
82- # Read messages from the stream and put them in the receive queue
83- for message in self ._stream :
84- if self ._is_stream_active ():
85- self ._receive_queue .put (message .event_message )
86- else :
87- break
88- except grpc .RpcError as e :
89- if e .code () != grpc .StatusCode .CANCELLED :
79+ # Check if the stream is not None
80+ if self ._stream is not None :
81+ # The first message dapr sends on the stream is for signalling only, so discard it
82+ next (self ._stream )
83+
84+ # Read messages from the stream and put them in the receive queue
85+ for message in self ._stream :
86+ if self ._is_stream_active ():
87+ self ._receive_queue .put (message .event_message )
88+ else :
89+ break
90+ except RpcError as e :
91+ if e .code () != StatusCode .CANCELLED :
9092 print (f'gRPC error in stream: { e .details ()} , Status Code: { e .code ()} ' )
9193 except Exception as e :
9294 raise Exception (f'Error while handling responses: { e } ' )
@@ -157,8 +159,8 @@ def close(self):
157159 if self ._stream :
158160 try :
159161 self ._stream .cancel ()
160- except grpc . RpcError as e :
161- if e .code () != grpc . StatusCode .CANCELLED :
162+ except RpcError as e :
163+ if e .code () != StatusCode .CANCELLED :
162164 raise Exception (f'Error while closing stream: { e } ' )
163165 except Exception as e :
164166 raise Exception (f'Error while closing stream: { e } ' )
0 commit comments