vnpy--vnpy
wehub-resource-sync
141 行
3.8 KiB
Python
141 行
3.8 KiB
Python
import threading
|
|
import traceback
|
|
from time import time
|
|
from collections.abc import Callable
|
|
|
|
import zmq
|
|
|
|
from .common import HEARTBEAT_TOPIC, HEARTBEAT_INTERVAL
|
|
|
|
|
|
class RpcServer:
|
|
""""""
|
|
|
|
def __init__(self) -> None:
|
|
"""
|
|
Constructor
|
|
"""
|
|
# Save functions dict: key is function name, value is function object
|
|
self._functions: dict[str, Callable] = {}
|
|
|
|
# Zmq port related
|
|
self._context: zmq.Context = zmq.Context()
|
|
|
|
# Reply socket (Request–reply pattern)
|
|
self._socket_rep: zmq.Socket = self._context.socket(zmq.REP)
|
|
|
|
# Publish socket (Publish–subscribe pattern)
|
|
self._socket_pub: zmq.Socket = self._context.socket(zmq.PUB)
|
|
|
|
# Worker thread related
|
|
self._active: bool = False # RpcServer status
|
|
self._thread: threading.Thread | None = None # RpcServer thread
|
|
self._lock: threading.Lock = threading.Lock()
|
|
|
|
# Heartbeat related
|
|
self._heartbeat_at: float | None = None
|
|
|
|
def is_active(self) -> bool:
|
|
""""""
|
|
return self._active
|
|
|
|
def start(
|
|
self,
|
|
rep_address: str,
|
|
pub_address: str,
|
|
) -> None:
|
|
"""
|
|
Start RpcServer
|
|
"""
|
|
if self._active:
|
|
return
|
|
|
|
# Bind socket address
|
|
self._socket_rep.bind(rep_address)
|
|
self._socket_pub.bind(pub_address)
|
|
|
|
# Start RpcServer status
|
|
self._active = True
|
|
|
|
# Start RpcServer thread
|
|
self._thread = threading.Thread(target=self.run)
|
|
self._thread.start()
|
|
|
|
# Init heartbeat publish timestamp
|
|
self._heartbeat_at = time() + HEARTBEAT_INTERVAL
|
|
|
|
def stop(self) -> None:
|
|
"""
|
|
Stop RpcServer
|
|
"""
|
|
if not self._active:
|
|
return
|
|
|
|
# Stop RpcServer status
|
|
self._active = False
|
|
|
|
def join(self) -> None:
|
|
# Wait for RpcServer thread to exit
|
|
if self._thread and self._thread.is_alive():
|
|
self._thread.join()
|
|
self._thread = None
|
|
|
|
def run(self) -> None:
|
|
"""
|
|
Run RpcServer functions
|
|
"""
|
|
while self._active:
|
|
# Poll response socket for 1 second
|
|
n: int = self._socket_rep.poll(1000)
|
|
self.check_heartbeat()
|
|
|
|
if not n:
|
|
continue
|
|
|
|
# Receive request data from Reply socket
|
|
req = self._socket_rep.recv_pyobj()
|
|
|
|
# Get function name and parameters
|
|
name, args, kwargs = req
|
|
|
|
# Try to get and execute callable function object; capture exception information if it fails
|
|
try:
|
|
func: Callable = self._functions[name]
|
|
r: object = func(*args, **kwargs)
|
|
rep: list = [True, r]
|
|
except Exception as e: # noqa
|
|
rep = [False, traceback.format_exc()]
|
|
|
|
# send callable response by Reply socket
|
|
self._socket_rep.send_pyobj(rep)
|
|
|
|
# Unbind socket address
|
|
self._socket_pub.close()
|
|
self._socket_rep.close()
|
|
|
|
def publish(self, topic: str, data: object) -> None:
|
|
"""
|
|
Publish data
|
|
"""
|
|
with self._lock:
|
|
self._socket_pub.send_pyobj([topic, data])
|
|
|
|
def register(self, func: Callable) -> None:
|
|
"""
|
|
Register function
|
|
"""
|
|
self._functions[func.__name__] = func
|
|
|
|
def check_heartbeat(self) -> None:
|
|
"""
|
|
Check whether it is required to send heartbeat.
|
|
"""
|
|
now: float = time()
|
|
|
|
if self._heartbeat_at and now >= self._heartbeat_at:
|
|
# Publish heartbeat
|
|
self.publish(HEARTBEAT_TOPIC, now)
|
|
|
|
# Update timestamp of next publish
|
|
self._heartbeat_at = now + HEARTBEAT_INTERVAL
|