collector: streamline tracking of polled fd-s

[loctrkd.git] / loctrkd / collector.py

diff --git a/loctrkd/collector.py b/loctrkd/collector.py
index 22bc9c3ffa2014346d74fd13b19de48d6d098472..6fc606a14ecabcbe41c6aba94c69ca7ce5bc55f4 100644 (file)
--- a/loctrkd/collector.py
+++ b/loctrkd/collector.py
@@ -14,7 +14,7 @@ from socket import (
 )
 from struct import pack
 from time import time
-from typing import Any, cast, Dict, List, Optional, Tuple, Union
+from typing import Any, cast, Dict, List, Optional, Set, Tuple, Union
 import zmq
 
 from . import common
@@ -68,7 +68,7 @@ pmods: List[ProtoModule] = []
 class Client:
     """Connected socket to the terminal plus buffer and metadata"""
 
-    def __init__(self, sock: socket, addr: Tuple[str, int]) -> None:
+    def __init__(self, sock: socket, addr: Any) -> None:
         self.sock = sock
         self.addr = addr
         self.pmod: Optional[ProtoModule] = None
@@ -87,7 +87,7 @@ class Client:
                 "%d bytes in buffer on close: %s", len(rest), rest[:64].hex()
             )
 
-    def recv(self) -> Optional[List[Tuple[float, Tuple[str, int], bytes]]]:
+    def recv(self) -> Optional[List[Tuple[float, Any, bytes]]]:
         """Read from the socket and parse complete messages"""
         try:
             segment = self.sock.recv(MAXBUFFER)
@@ -152,25 +152,32 @@ class Clients:
         self.by_fd: Dict[int, Client] = {}
         self.by_imei: Dict[str, Client] = {}
 
-    def add(self, clntsock: socket, clntaddr: Tuple[str, int]) -> int:
+    def fds(self) -> Set[int]:
+        return set(self.by_fd.keys())
+
+    def add(self, clntsock: socket, clntaddr: Any) -> int:
         fd = clntsock.fileno()
         log.info("Start serving fd %d from %s", fd, clntaddr)
         self.by_fd[fd] = Client(clntsock, clntaddr)
         return fd
 
     def stop(self, fd: int) -> None:
+        if fd not in self.by_fd:
+            log.debug("Fd %d is not served, ingore stop", fd)
+            return
         clnt = self.by_fd[fd]
         log.info("Stop serving fd %d (IMEI %s)", clnt.sock.fileno(), clnt.imei)
         clnt.close()
-        if clnt.imei:
+        if clnt.imei and self.by_imei[clnt.imei] == clnt:  # could be replaced
             del self.by_imei[clnt.imei]
         del self.by_fd[fd]
 
     def recv(
         self, fd: int
-    ) -> Optional[
-        List[Tuple[ProtoModule, Optional[str], float, Tuple[str, int], bytes]]
-    ]:
+    ) -> Optional[List[Tuple[ProtoModule, Optional[str], float, Any, bytes]]]:
+        if fd not in self.by_fd:
+            log.debug("Client at fd %d gone, ingore event", fd)
+            return None
         clnt = self.by_fd[fd]
         msgs = clnt.recv()
         if msgs is None:
@@ -185,18 +192,11 @@ class Clients:
                     clnt.imei = imei
                     oldclnt = self.by_imei.get(clnt.imei)
                     if oldclnt is not None:
-                        log.info(
-                            "Orphaning fd %d with the same IMEI",
-                            oldclnt.sock.fileno(),
-                        )
+                        oldfd = oldclnt.sock.fileno()
+                        log.info("Removing stale connection on fd %d", oldfd)
                         oldclnt.imei = None
+                        self.stop(oldfd)
                     self.by_imei[clnt.imei] = clnt
-                else:
-                    log.warning(
-                        "Login message from %s: %s, but client imei unfilled",
-                        peeraddr,
-                        packet,
-                    )
             result.append((clnt.pmod, clnt.imei, when, peeraddr, packet))
             log.debug(
                 "Received from %s (IMEI %s): %s",
@@ -239,11 +239,11 @@ def runserver(conf: ConfigParser, handle_hibernate: bool = True) -> None:
     poller.register(zpull, flags=zmq.POLLIN)
     poller.register(tcpfd, flags=zmq.POLLIN)
     clients = Clients()
+    pollingfds: Set[int] = set()
     try:
         while True:
-            tosend = []
-            topoll = []
-            tostop = []
+            tosend: List[Resp] = []
+            toadd: List[Tuple[socket, Any]] = []
             events = poller.poll(1000)
             for sk, fl in events:
                 if sk is zpull:
@@ -257,12 +257,12 @@ def runserver(conf: ConfigParser, handle_hibernate: bool = True) -> None:
                 elif sk == tcpfd:
                     clntsock, clntaddr = tcpl.accept()
                     clntsock.setsockopt(SOL_SOCKET, SO_KEEPALIVE, 1)
-                    topoll.append((clntsock, clntaddr))
+                    toadd.append((clntsock, clntaddr))
                 elif fl & zmq.POLLIN:
                     received = clients.recv(sk)
                     if received is None:
                         log.debug("Terminal gone from fd %d", sk)
-                        tostop.append(sk)
+                        clients.stop(sk)
                     else:
                         for pmod, imei, when, peeraddr, packet in received:
                             proto = pmod.proto_of_message(packet)
@@ -284,7 +284,7 @@ def runserver(conf: ConfigParser, handle_hibernate: bool = True) -> None:
                                     sk,
                                     imei,
                                 )
-                                tostop.append(sk)
+                                clients.stop(sk)
                             respmsg = pmod.inline_response(packet)
                             if respmsg is not None:
                                 tosend.append(
@@ -306,12 +306,13 @@ def runserver(conf: ConfigParser, handle_hibernate: bool = True) -> None:
                             packet=zmsg.packet,
                         ).packed
                     )
-            for fd in tostop:
+            for fd in pollingfds - clients.fds():
                 poller.unregister(fd)  # type: ignore
-                clients.stop(fd)
-            for clntsock, clntaddr in topoll:
+            for clntsock, clntaddr in toadd:
                 fd = clients.add(clntsock, clntaddr)
+            for fd in clients.fds() - pollingfds:
                 poller.register(fd, flags=zmq.POLLIN)
+            pollingfds = clients.fds()
     except KeyboardInterrupt:
         zpub.close()
         zpull.close()