Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit ef89f55

Browse files
plamutplamut
authored
fix: threads can skip the line in publisher flow controller (#422)
* Add publisher flow controller test for FIFO order * Simplify _run_in_daemon() test helper * Fix message slots not acquired in FIFO order * Unify the logic for distributing any free capacity * Use OrderedDict for the FIFO queue This allows to hold the queue of threads and their reservation data in a single structure, no need for the separate deque and reservations dict.
1 parent 0df7c96 commit ef89f55

File tree

2 files changed

+137
-77
lines changed

2 files changed

+137
-77
lines changed

‎google/cloud/pubsub_v1/publisher/flow_controller.py‎

Lines changed: 74 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15-
from collections import deque
15+
from collections import OrderedDict
1616
import logging
1717
import threading
1818
import warnings
@@ -24,12 +24,21 @@
2424
_LOGGER = logging.getLogger(__name__)
2525

2626

27-
class _QuantityReservation(object):
28-
"""A (partial) reservation of a quantifiable resource."""
27+
class _QuantityReservation:
28+
"""A (partial) reservation of quantifiable resources."""
2929

30-
def __init__(self, reserved, needed):
31-
self.reserved = reserved
32-
self.needed = needed
30+
def __init__(self, bytes_reserved: int, bytes_needed: int, has_slot: bool):
31+
self.bytes_reserved = bytes_reserved
32+
self.bytes_needed = bytes_needed
33+
self.has_slot = has_slot
34+
35+
def __repr__(self):
36+
return (
37+
f"{type(self).__name__}("
38+
f"bytes_reserved={self.bytes_reserved}, "
39+
f"bytes_needed={self.bytes_needed}, "
40+
f"has_slot={self.has_slot})"
41+
)
3342

3443

3544
class FlowController(object):
@@ -48,14 +57,13 @@ def __init__(self, settings):
4857
self._message_count = 0
4958
self._total_bytes = 0
5059

51-
# A FIFO queue of threads blocked on adding a message, from first to last.
60+
# A FIFO queue of threads blocked on adding a message that also tracks their
61+
# reservations of available flow control bytes and message slots.
5262
# Only relevant if the configured limit exceeded behavior is BLOCK.
53-
self._waiting = deque()
63+
self._waiting = OrderedDict()
5464

55-
# Reservations of available flow control bytes by the waiting threads.
56-
# Each value is a _QuantityReservation instance.
57-
self._byte_reservations = dict()
5865
self._reserved_bytes = 0
66+
self._reserved_slots = 0
5967

6068
# The lock is used to protect all internal state (message and byte count,
6169
# waiting threads to add, etc.).
@@ -131,11 +139,13 @@ def add(self, message):
131139
current_thread = threading.current_thread()
132140

133141
while self._would_overflow(message):
134-
if current_thread not in self._byte_reservations:
135-
self._waiting.append(current_thread)
136-
self._byte_reservations[current_thread] = _QuantityReservation(
137-
reserved=0, needed=message._pb.ByteSize()
142+
if current_thread not in self._waiting:
143+
reservation = _QuantityReservation(
144+
bytes_reserved=0,
145+
bytes_needed=message._pb.ByteSize(),
146+
has_slot=False,
138147
)
148+
self._waiting[current_thread] = reservation # Will be placed last.
139149

140150
_LOGGER.debug(
141151
"Blocking until there is enough free capacity in the flow - "
@@ -152,9 +162,9 @@ def add(self, message):
152162
# Message accepted, increase the load and remove thread stats.
153163
self._message_count += 1
154164
self._total_bytes += message._pb.ByteSize()
155-
self._reserved_bytes -= self._byte_reservations[current_thread].reserved
156-
del self._byte_reservations[current_thread]
157-
self._waiting.remove(current_thread)
165+
self._reserved_bytes -= self._waiting[current_thread].bytes_reserved
166+
self._reserved_slots -= 1
167+
del self._waiting[current_thread]
158168

159169
def release(self, message):
160170
"""Release a mesage from flow control.
@@ -180,39 +190,52 @@ def release(self, message):
180190
self._message_count = max(0, self._message_count)
181191
self._total_bytes = max(0, self._total_bytes)
182192

183-
self._distribute_available_bytes()
193+
self._distribute_available_capacity()
184194

185195
# If at least one thread waiting to add() can be unblocked, wake them up.
186196
if self._ready_to_unblock():
187197
_LOGGER.debug("Notifying threads waiting to add messages to flow.")
188198
self._has_capacity.notify_all()
189199

190-
def _distribute_available_bytes(self):
191-
"""Distribute availalbe free capacity among the waiting threads in FIFO order.
200+
def _distribute_available_capacity(self):
201+
"""Distribute available capacity among the waiting threads in FIFO order.
192202
193203
The method assumes that the caller has obtained ``_operational_lock``.
194204
"""
195-
available = self._settings.byte_limit - self._total_bytes - self._reserved_bytes
205+
available_slots = (
206+
self._settings.message_limit - self._message_count - self._reserved_slots
207+
)
208+
available_bytes = (
209+
self._settings.byte_limit - self._total_bytes - self._reserved_bytes
210+
)
211+
212+
for reservation in self._waiting.values():
213+
if available_slots <= 0 and available_bytes <= 0:
214+
break # Santa is now empty-handed, better luck next time.
196215

197-
for thread in self._waiting:
198-
if available <= 0:
199-
break
216+
# Distribute any free slots.
217+
if available_slots > 0 and not reservation.has_slot:
218+
reservation.has_slot = True
219+
self._reserved_slots += 1
220+
available_slots -= 1
200221

201-
reservation = self._byte_reservations[thread]
202-
still_needed = reservation.needed - reservation.reserved
222+
# Distribute any free bytes.
223+
if available_bytes <= 0:
224+
continue
203225

204-
# Sanity check for any internal inconsistencies.
205-
if still_needed < 0:
226+
bytes_still_needed = reservation.bytes_needed - reservation.bytes_reserved
227+
228+
if bytes_still_needed < 0: # Sanity check for any internal inconsistencies.
206229
msg = "Too many bytes reserved: {} / {}".format(
207-
reservation.reserved, reservation.needed
230+
reservation.bytes_reserved, reservation.bytes_needed
208231
)
209232
warnings.warn(msg, category=RuntimeWarning)
210-
still_needed = 0
233+
bytes_still_needed = 0
211234

212-
can_give = min(still_needed, available)
213-
reservation.reserved += can_give
235+
can_give = min(bytes_still_needed, available_bytes)
236+
reservation.bytes_reserved += can_give
214237
self._reserved_bytes += can_give
215-
available -= can_give
238+
available_bytes -= can_give
216239

217240
def _ready_to_unblock(self):
218241
"""Determine if any of the threads waiting to add a message can proceed.
@@ -225,10 +248,10 @@ def _ready_to_unblock(self):
225248
if self._waiting:
226249
# It's enough to only check the head of the queue, because FIFO
227250
# distribution of any free capacity.
228-
reservation = self._byte_reservations[self._waiting[0]]
251+
first_reservation = next(iter(self._waiting.values()))
229252
return (
230-
reservation.reserved >= reservation.needed
231-
and self._message_count < self._settings.message_limit
253+
first_reservation.bytes_reserved >= first_reservation.bytes_needed
254+
and first_reservation.has_slot
232255
)
233256

234257
return False
@@ -245,16 +268,22 @@ def _would_overflow(self, message):
245268
Returns:
246269
bool
247270
"""
248-
reservation = self._byte_reservations.get(threading.current_thread())
271+
reservation = self._waiting.get(threading.current_thread())
249272

250273
if reservation:
251-
enough_reserved = reservation.reserved >= reservation.needed
274+
enough_reserved = reservation.bytes_reserved >= reservation.bytes_needed
275+
has_slot = reservation.has_slot
252276
else:
253277
enough_reserved = False
278+
has_slot = False
254279

255280
bytes_taken = self._total_bytes + self._reserved_bytes + message._pb.ByteSize()
256281
size_overflow = bytes_taken > self._settings.byte_limit and not enough_reserved
257-
msg_count_overflow = self._message_count + 1 > self._settings.message_limit
282+
283+
msg_count_overflow = not has_slot and (
284+
(self._message_count + self._reserved_slots + 1)
285+
> self._settings.message_limit
286+
)
258287

259288
return size_overflow or msg_count_overflow
260289

@@ -275,18 +304,15 @@ def _load_info(self, message_count=None, total_bytes=None):
275304
Returns:
276305
str
277306
"""
278-
msg = "messages: {} / {}, bytes: {} / {} (reserved: {})"
279-
280307
if message_count is None:
281308
message_count = self._message_count
282309

283310
if total_bytes is None:
284311
total_bytes = self._total_bytes
285312

286-
return msg.format(
287-
message_count,
288-
self._settings.message_limit,
289-
total_bytes,
290-
self._settings.byte_limit,
291-
self._reserved_bytes,
313+
return (
314+
f"messages: {message_count} / {self._settings.message_limit} "
315+
f"(reserved: {self._reserved_slots}), "
316+
f"bytes: {total_bytes} / {self._settings.byte_limit} "
317+
f"(reserved: {self._reserved_bytes})"
292318
)

0 commit comments

Comments
 (0)