Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/google/cloud/pubsub_v1/subscriber/_protocol/messages_on_hold.py: 26%

1# Copyright 2020, Google LLC 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# https://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15import collections 

16import logging 

17import typing 

18from typing import Any, Callable, Iterable, Optional 

19 

20if typing.TYPE_CHECKING: # pragma: NO COVER 

21 from google.cloud.pubsub_v1 import subscriber 

22 

23 

24_LOGGER = logging.getLogger(__name__) 

25 

26 

27class MessagesOnHold(object): 

28 """Tracks messages on hold by ordering key. Not thread-safe.""" 

29 

30 def __init__(self): 

31 self._size = 0 

32 

33 # A FIFO queue for the messages that have been received from the server, 

34 # but not yet sent to the user callback. 

35 # Both ordered and unordered messages may be in this queue. Ordered 

36 # message state tracked in _pending_ordered_messages once ordered 

37 # messages are taken off this queue. 

38 # The tail of the queue is to the right side of the deque; the head is 

39 # to the left side. 

40 self._messages_on_hold = collections.deque() 

41 

42 # Dict of ordering_key -> queue of ordered messages that have not been 

43 # delivered to the user. 

44 # All ordering keys in this collection have a message in flight. Once 

45 # that one is acked or nacked, the next message in the queue for that 

46 # ordering key will be sent. 

47 # If the queue is empty, it means there's a message for that key in 

48 # flight, but there are no pending messages. 

49 self._pending_ordered_messages = {} 

50 

51 @property 

52 def size(self) -> int: 

53 """Return the number of messages on hold across ordered and unordered messages. 

54 

55 Note that this object may still store information about ordered messages 

56 in flight even if size is zero. 

57 

58 Returns: 

59 The size value. 

60 """ 

61 return self._size 

62 

63 def get(self) -> Optional["subscriber.message.Message"]: 

64 """Gets a message from the on-hold queue. A message with an ordering 

65 key wont be returned if there's another message with the same key in 

66 flight. 

67 

68 Returns: 

69 A message that hasn't been sent to the user yet or ``None`` if there are no 

70 messages available. 

71 """ 

72 while self._messages_on_hold: 

73 msg = self._messages_on_hold.popleft() 

74 

75 if msg.ordering_key: 

76 pending_queue = self._pending_ordered_messages.get(msg.ordering_key) 

77 if pending_queue is None: 

78 # Create empty queue to indicate a message with the 

79 # ordering key is in flight. 

80 self._pending_ordered_messages[ 

81 msg.ordering_key 

82 ] = collections.deque() 

83 self._size = self._size - 1 

84 return msg 

85 else: 

86 # Another message is in flight so add message to end of 

87 # queue for this ordering key. 

88 pending_queue.append(msg) 

89 else: 

90 # Unordered messages can be returned without any 

91 # restrictions. 

92 self._size = self._size - 1 

93 return msg 

94 

95 return None 

96 

97 def put(self, message: "subscriber.message.Message") -> None: 

98 """Put a message on hold. 

99 

100 Args: 

101 message: The message to put on hold. 

102 """ 

103 self._messages_on_hold.append(message) 

104 self._size = self._size + 1 

105 

106 def activate_ordering_keys( 

107 self, 

108 ordering_keys: Iterable[str], 

109 schedule_message_callback: Callable[["subscriber.message.Message"], Any], 

110 ) -> None: 

111 """Send the next message in the queue for each of the passed-in 

112 ordering keys, if they exist. Clean up state for keys that no longer 

113 have any queued messages. 

114 

115 See comment at streaming_pull_manager.activate_ordering_keys() for more 

116 detail about the impact of this method on load. 

117 

118 Args: 

119 ordering_keys: 

120 The ordering keys to activate. May be empty, or contain duplicates. 

121 schedule_message_callback: 

122 The callback to call to schedule a message to be sent to the user. 

123 """ 

124 for key in ordering_keys: 

125 pending_ordered_messages = self._pending_ordered_messages.get(key) 

126 if pending_ordered_messages is None: 

127 _LOGGER.warning( 

128 "No message queue exists for message ordering key: %s.", key 

129 ) 

130 continue 

131 next_msg = self._get_next_for_ordering_key(key) 

132 if next_msg: 

133 # Schedule the next message because the previous was dropped. 

134 # Note that this may overload the user's `max_bytes` limit, but 

135 # not their `max_messages` limit. 

136 schedule_message_callback(next_msg) 

137 else: 

138 # No more messages for this ordering key, so do clean-up. 

139 self._clean_up_ordering_key(key) 

140 

141 def _get_next_for_ordering_key( 

142 self, ordering_key: str 

143 ) -> Optional["subscriber.message.Message"]: 

144 """Get next message for ordering key. 

145 

146 The client should call clean_up_ordering_key() if this method returns 

147 None. 

148 

149 Args: 

150 ordering_key: Ordering key for which to get the next message. 

151 

152 Returns: 

153 The next message for this ordering key or None if there aren't any. 

154 """ 

155 queue_for_key = self._pending_ordered_messages.get(ordering_key) 

156 if queue_for_key: 

157 self._size = self._size - 1 

158 return queue_for_key.popleft() 

159 return None 

160 

161 def _clean_up_ordering_key(self, ordering_key: str) -> None: 

162 """Clean up state for an ordering key with no pending messages. 

163 

164 Args 

165 ordering_key: The ordering key to clean up. 

166 """ 

167 message_queue = self._pending_ordered_messages.get(ordering_key) 

168 if message_queue is None: 

169 _LOGGER.warning( 

170 "Tried to clean up ordering key that does not exist: %s", ordering_key 

171 ) 

172 return 

173 if len(message_queue) > 0: 

174 _LOGGER.warning( 

175 "Tried to clean up ordering key: %s with %d messages remaining.", 

176 ordering_key, 

177 len(message_queue), 

178 ) 

179 return 

180 del self._pending_ordered_messages[ordering_key]