Make Socket._wait_for follow njsmith's original a little closer

I'd ignored what would happen if there were multiple send tasks waiting

Author: remleduff

.gitignore

@@ -0,0 +1,2 @@
+venv
+__pycache__

pubsub.py

@@ -0,0 +1,41 @@
+
+import trzmq
+import trio
+import zmq
+
+
+async def pub(n):
+    context = zmq.Context()
+    socket = context.socket(zmq.PUB)
+    socket.connect("tcp://0.0.0.0:5556")
+
+    s = trzmq.Socket(socket)
+    i = 1
+    while True:
+        topic = b"ZMQ-Test"
+        message = f"Hello, from {n}: {i}..."
+        await s.send(b"%b %b" % (topic, message.encode()))
+        print("%s %s" % (topic, message))
+        i += 1
+        await trio.sleep(1)
+
+
+async def sub():
+    context = zmq.Context()
+    socket = context.socket(zmq.SUB)
+    socket.bind("tcp://0.0.0.0:5556")
+    socket.setsockopt_string(zmq.SUBSCRIBE, '')
+    s = trzmq.Socket(socket)
+    while True:
+        string = await s.recv()
+        print(string)
+
+
+async def run():
+    async with trio.open_nursery() as nursery:
+        nursery.start_soon(sub)
+        nursery.start_soon(pub, 1)
+        nursery.start_soon(pub, 2)
+
+
+trio.run(run)

trzmq/_socket.py

@@ -45,18 +45,20 @@ async def _wait_for(self, flag):
             await trio.lowlevel.checkpoint()
             return
 
-        if not self._someone_is_watching_trigger:
-            self._someone_is_watching_trigger = True
-            try:
-                while not self._events & flag:
+        while not self._events & flag:
+            if not self._someone_is_watching_trigger:
+                self._someone_is_watching_trigger = True
+                try:
                     await trio.lowlevel.wait_readable(self._zmq_sock.fd)
+                finally:
+                    self._someone_is_watching_trigger = False
                     self._update_events()
-            finally:
-                self._someone_is_watching_trigger = False
-        else:
-            # Someone else is watching the socket, so we just need to
-            # wait to be woken up.
-            await self._wake_events[flag].wait()
+            else:
+                # Someone else is watching the socket, so we just need to
+                # wait to be woken up
+                await self._wake_events[flag].wait()
+                # self._events will be uptodate because if there was another
+                # sender/receiver they ran _update_events after send/receive
 
     # XX needs conflict detection
     # ...or does it? each send() call is atomic, like a DGRAM socket.