Make objects.vectors_to_send() a generator of chunks

This commit is contained in:
Lee Miller 2024-10-13 05:09:47 +03:00
parent 0ae0cf5b95
commit 912b7a152d
Signed by: lee.miller
GPG Key ID: 4F97A5EA88F4AB63
3 changed files with 17 additions and 11 deletions

View File

@ -278,14 +278,11 @@ class ConnectionBase(threading.Thread):
self.send_queue.put(message.Addr(addr)) self.send_queue.put(message.Addr(addr))
if shared.objects: if shared.objects:
to_send = shared.objects.vectors_to_send() for chunk in shared.objects.vectors_to_send(10000):
offset = 0
while offset < len(to_send):
# We limit size of inv messaged to 10000 entries # We limit size of inv messaged to 10000 entries
# because they might time out # because they might time out
# in very slow networks (I2P) # in very slow networks (I2P)
self.send_queue.put(message.Inv(to_send[offset:offset+10000])) self.send_queue.put(message.Inv(chunk))
offset += 10000
self.status = 'fully_established' self.status = 'fully_established'
def _process_queue(self): def _process_queue(self):

View File

@ -115,13 +115,20 @@ class Inventory():
self._pending.update(vectors) self._pending.update(vectors)
return vectors return vectors
def vectors_to_send(self, stream=None): def vectors_to_send(self, chunk_size=10000, stream=None):
if stream is None:
stream = shared.stream
now = int(time.time())
cur = self._db.execute( cur = self._db.execute(
'SELECT vector FROM objects WHERE expires > ? AND stream = ?' 'SELECT vector FROM objects WHERE expires > ? AND stream = ?'
' ORDER BY random()', ' ORDER BY random()', (now, stream)
(int(time.time()), stream or shared.stream)
) )
return [v for v, in cur] cur.arraysize = chunk_size
while True:
vectors = cur.fetchmany()
if not vectors:
return
yield [v for v, in vectors]
def get(self, vector, default=None): def get(self, vector, default=None):
try: try:

View File

@ -65,8 +65,10 @@ class TestObjects():
self.objects[obj.vector] = obj self.objects[obj.vector] = obj
needed.add(obj.vector) needed.add(obj.vector)
self.assertEqual(set(self.objects.vectors_to_send(4)), needed) self.assertEqual(
self.assertTrue(set(self.objects.vectors_to_send()).difference(needed)) set(next(self.objects.vectors_to_send(stream=4))), needed)
self.assertTrue(
set(next(self.objects.vectors_to_send())).difference(needed))
def test_filter(self): def test_filter(self):
"""Check the objects filtering""" """Check the objects filtering"""