wehub-resource-sync
170 lines
4.6 KiB
Python
170 lines
4.6 KiB
Python
import threading
|
||
from time import time
|
||
from functools import lru_cache
|
||
from typing import Any
|
||
|
||
import zmq
|
||
|
||
from .common import HEARTBEAT_TOPIC, HEARTBEAT_TOLERANCE
|
||
|
||
|
||
class RemoteException(Exception):
|
||
"""
|
||
RPC remote exception
|
||
"""
|
||
|
||
def __init__(self, value: Any) -> None:
|
||
"""
|
||
Constructor
|
||
"""
|
||
self._value: Any = value
|
||
|
||
def __str__(self) -> str:
|
||
"""
|
||
Output error message
|
||
"""
|
||
return str(self._value)
|
||
|
||
|
||
class RpcClient:
|
||
""""""
|
||
|
||
def __init__(self) -> None:
|
||
"""Constructor"""
|
||
# zmq port related
|
||
self._context: zmq.Context = zmq.Context()
|
||
|
||
# Request socket (Request–reply pattern)
|
||
self._socket_req: zmq.Socket = self._context.socket(zmq.REQ)
|
||
|
||
# Subscribe socket (Publish–subscribe pattern)
|
||
self._socket_sub: zmq.Socket = self._context.socket(zmq.SUB)
|
||
|
||
# Set socket option to keepalive
|
||
for socket in [self._socket_req, self._socket_sub]:
|
||
socket.setsockopt(zmq.TCP_KEEPALIVE, 1)
|
||
socket.setsockopt(zmq.TCP_KEEPALIVE_IDLE, 60)
|
||
|
||
# Worker thread relate, used to process data pushed from server
|
||
self._active: bool = False # RpcClient status
|
||
self._thread: threading.Thread | None = None # RpcClient thread
|
||
self._lock: threading.Lock = threading.Lock()
|
||
|
||
self._last_received_ping: float = time()
|
||
|
||
@lru_cache(100) # noqa
|
||
def __getattr__(self, name: str) -> Any:
|
||
"""
|
||
Realize remote call function
|
||
"""
|
||
# Perform remote call task
|
||
def dorpc(*args: Any, **kwargs: Any) -> Any:
|
||
# Get timeout value from kwargs, default value is 30 seconds
|
||
timeout: int = kwargs.pop("timeout", 30000)
|
||
|
||
# Generate request
|
||
req: list = [name, args, kwargs]
|
||
|
||
# Send request and wait for response
|
||
with self._lock:
|
||
self._socket_req.send_pyobj(req)
|
||
|
||
# Timeout reached without any data
|
||
n: int = self._socket_req.poll(timeout)
|
||
if not n:
|
||
msg: str = f"Timeout of {timeout}ms reached for {req}"
|
||
raise RemoteException(msg)
|
||
|
||
rep = self._socket_req.recv_pyobj()
|
||
|
||
# Return response if successed; Trigger exception if failed
|
||
if rep[0]:
|
||
return rep[1]
|
||
else:
|
||
raise RemoteException(rep[1])
|
||
|
||
return dorpc
|
||
|
||
def start(
|
||
self,
|
||
req_address: str,
|
||
sub_address: str
|
||
) -> None:
|
||
"""
|
||
Start RpcClient
|
||
"""
|
||
if self._active:
|
||
return
|
||
|
||
# Connect zmq port
|
||
self._socket_req.connect(req_address)
|
||
self._socket_sub.connect(sub_address)
|
||
|
||
# Start RpcClient status
|
||
self._active = True
|
||
|
||
# Start RpcClient thread
|
||
self._thread = threading.Thread(target=self.run)
|
||
self._thread.start()
|
||
|
||
self._last_received_ping = time()
|
||
|
||
def stop(self) -> None:
|
||
"""
|
||
Stop RpcClient
|
||
"""
|
||
if not self._active:
|
||
return
|
||
|
||
# Stop RpcClient status
|
||
self._active = False
|
||
|
||
def join(self) -> None:
|
||
# Wait for RpcClient thread to exit
|
||
if self._thread and self._thread.is_alive():
|
||
self._thread.join()
|
||
self._thread = None
|
||
|
||
def run(self) -> None:
|
||
"""
|
||
Run RpcClient function
|
||
"""
|
||
pull_tolerance: int = HEARTBEAT_TOLERANCE * 1000
|
||
|
||
while self._active:
|
||
if not self._socket_sub.poll(pull_tolerance):
|
||
self.on_disconnected()
|
||
continue
|
||
|
||
# Receive data from subscribe socket
|
||
topic, data = self._socket_sub.recv_pyobj(flags=zmq.NOBLOCK)
|
||
|
||
if topic == HEARTBEAT_TOPIC:
|
||
self._last_received_ping = data
|
||
else:
|
||
# Process data by callable function
|
||
self.callback(topic, data)
|
||
|
||
# Close socket
|
||
self._socket_req.close()
|
||
self._socket_sub.close()
|
||
|
||
def callback(self, topic: str, data: Any) -> None:
|
||
"""
|
||
Callable function
|
||
"""
|
||
raise NotImplementedError
|
||
|
||
def subscribe_topic(self, topic: str) -> None:
|
||
"""
|
||
Subscribe data
|
||
"""
|
||
self._socket_sub.setsockopt_string(zmq.SUBSCRIBE, topic)
|
||
|
||
def on_disconnected(self) -> None:
|
||
"""
|
||
Callback when heartbeat is lost.
|
||
"""
|
||
msg: str = f"RpcServer has no response over {HEARTBEAT_TOLERANCE} seconds, please check you connection."
|
||
print(msg)
|