Turned dandelion stem into a module and corrected use of wait_for

src/gossip/client/dandelionstem/__init__.py

@@ -2,7 +2,6 @@ from queue import Queue
 from time import sleep
 from secrets import choice
 import traceback
-from asyncio import wait_for
 
 from typing import TYPE_CHECKING, Coroutine, Tuple, List
 
@@ -10,15 +9,17 @@ from onionrthreads import add_delayed_thread
 from blockdb import add_block_to_db
 import logger
 
-from ..constants import BLACKHOLE_EVADE_TIMER_SECS, OUTBOUND_DANDELION_EDGES
-from ..commands import GossipCommands, command_to_byte
-from .. import dandelion
+from ...constants import BLACKHOLE_EVADE_TIMER_SECS, OUTBOUND_DANDELION_EDGES
+from ...commands import GossipCommands, command_to_byte
+from ... import dandelion
+
+from .stemstream import do_stem_stream
 
 if TYPE_CHECKING:
     from ordered_set import OrderedSet
     from onionrblocks import Block
-    from ..peer import Peer
-    from ..dandelion.phase import DandelionPhase
+    from ...peer import Peer
+    from ...dandelion.phase import DandelionPhase
     import socket
 
 
@@ -27,21 +28,18 @@ class StemConnectionDenied(ConnectionRefusedError): pass  # noqa
 
 
 async def _setup_edge(
-        peer_set: OrderedSet['Peer'], exclude_set: OrderedSet['Peer']):
+        peer_set: "OrderedSet[Peer]", exclude_set: "OrderedSet[Peer]"):
     """Negotiate stem connection with random peer, add to exclu set if fail"""
     try:
         peer: 'Peer' = choice(peer_set - exclude_set)
     except IndexError:
         raise NotEnoughEdges
 
-    async def _get_sock():
-        return peer.get_socket()
-
     # If peer is good or bad, exclude it no matter what
     exclude_set.add(peer)
 
     try:
-        s = await wait_for(_get_sock(), 12)
+        s = peer.get_socket(12)
     except TimeoutError:
         logger.debug(f"{peer.transport_address} timed out when trying stemout")
     except Exception:
@@ -50,6 +48,7 @@ async def _setup_edge(
 
     try:
         s.sendall(command_to_byte(GossipCommands.PUT_BLOCKS))
+        s.settimeout(10)
         if s.recv(1) == dandelion.StemAcceptResult.DENY:
             raise StemConnectionDenied
     except StemConnectionDenied:
@@ -68,16 +67,9 @@ async def _setup_edge(
     s.close()
 
 
-async def _do_stem_stream(
-        peer_socket: 'socket.socket',
-        block_queue: Queue['Block'],
-        d_phase: 'DandelionPhase'):
-    return
-
-
 async def stem_out(
-        block_queues: Tuple[Queue['Block'], Queue['Block']],
-        peer_set: OrderedSet['Peer'],
+        block_queues: Tuple["Queue[Block]", "Queue[Block]"],
+        peer_set: "OrderedSet[Peer]",
         d_phase: 'DandelionPhase'):
 
     # don't bother if there are no possible outbound edges
@@ -97,16 +89,16 @@ async def stem_out(
 
     # Pick edges randomly
     # Using orderedset for the tried edges to ensure random pairing with queue
-    tried_edges: OrderedSet['Peer'] = OrderedSet()
+    tried_edges: "OrderedSet[Peer]" = OrderedSet()
 
     while len(peer_sockets) < OUTBOUND_DANDELION_EDGES:
         try:
             # Get a socket for stem out (makes sure they accept)
-            peer_sockets.append(
-                await wait_for(_setup_edge(peer_set, tried_edges), 15))
+            peer_sockets.append(await _setup_edge(peer_set, tried_edges))
         except NotEnoughEdges:
             # No possible edges at this point (edges < OUTBOUND_DANDELION_EDGE)
-            logger.warn("Not able to build enough peers for stemout.", terminal=True)
+            logger.warn("Not able to build enough peers for stemout.",
+                        terminal=True)
             break
         else:
             # Ran out of time for stem phase
@@ -122,20 +114,13 @@ async def stem_out(
 
     for count, peer_socket in enumerate(peer_sockets):
         stream_routines.append(
-            _do_stem_stream(peer_socket, block_queues[count], d_phase))
+            do_stem_stream(peer_socket, block_queues[count], d_phase))
 
     for routine in stream_routines:
         try:
-            await wait_for(routine, d_phase.remaining_time() + 1)
-        except TimeoutError:
-            # This is what should happen if everything went well
-            pass
+            await routine
         except Exception:
             logger.warn(traceback.format_exc())
         else:
             # stream routine exited early
             pass
-
-    # Clean up stem sockets
-    for sock in peer_sockets:
-        sock.close()