286 lines
9.2 KiB
Python
286 lines
9.2 KiB
Python
import socket
|
|
import time
|
|
import re
|
|
from select import select
|
|
|
|
|
|
class StreamDead(Exception):
|
|
"""raised when stream is dead"""
|
|
|
|
|
|
def parse_uri(uri):
|
|
scheme, _, hostport = uri.rpartition('://')
|
|
scheme = scheme or 'tcp'
|
|
if scheme != 'tcp':
|
|
raise ValueError(f'scheme {scheme!r} not supported')
|
|
host, _, port = hostport.rpartition(':')
|
|
host = host or 'localhost'
|
|
return host, int(port)
|
|
|
|
|
|
NETBUFSIZE = 4092
|
|
INF = float('inf')
|
|
|
|
|
|
class Base:
|
|
select_dict = {}
|
|
|
|
|
|
def short_hostname(host):
|
|
"""psi/lin/se special
|
|
|
|
- treat case where -129129xxxx is appended
|
|
"""
|
|
host = socket.gethostbyaddr(host)[0]
|
|
match = re.match(r'([^.-]+)(?:-129129\d{6}|(-[~.]*|)).psi.ch', host)
|
|
if match:
|
|
host = match.group(1) + (match.group(2) or '')
|
|
return host
|
|
|
|
|
|
class Stream(Base):
|
|
_last_time = None
|
|
dead = False
|
|
max_offline = 20
|
|
ping_timeout = 3
|
|
_next_ping = INF
|
|
_ping_deadline = 0
|
|
_deadline = INF
|
|
_next_connect = 0
|
|
|
|
def __init__(self, uri, name=None, timeout=5, encoding='latin-1'):
|
|
self.name = name or uri
|
|
self.uri = uri
|
|
self.tags = {}
|
|
self.streamname = self.uri
|
|
self.encoding = encoding
|
|
self.timeout = timeout
|
|
self.socket = None
|
|
self.cache = {}
|
|
self.errors = {}
|
|
self.start_time = time.time()
|
|
self.next_hour = (self.start_time // 3600 + 1) * 3600
|
|
self.generator = self.event_generator()
|
|
try:
|
|
self.connect()
|
|
self.init()
|
|
except Exception as e:
|
|
print(self.uri, repr(e))
|
|
raise
|
|
|
|
def connect(self):
|
|
self.socket = socket.create_connection(parse_uri(self.uri))
|
|
self.select_dict[self.socket.fileno()] = self
|
|
self.settimeout(self.timeout)
|
|
host, _, port = self.uri.partition(':')
|
|
# try to convert uri to host name
|
|
host = short_hostname(host)
|
|
self.tags['stream'] = f'{host}:{port}'
|
|
print(self.uri, '=', self.tags['stream'], 'connected')
|
|
self._buffer = []
|
|
self._deadline = INF
|
|
self._next_connect = 0
|
|
self._pinged = False
|
|
|
|
def init(self):
|
|
raise NotImplementedError
|
|
|
|
def ping(self):
|
|
raise NotImplementedError
|
|
|
|
def pong(self):
|
|
self._alive = True
|
|
|
|
def close(self):
|
|
"""Do our best to close a socket."""
|
|
if self.socket is None:
|
|
return
|
|
self.select_dict.pop(self.socket.fileno(), None)
|
|
print(self.uri, 'close socket')
|
|
try:
|
|
self.socket.shutdown(socket.SHUT_RDWR)
|
|
except socket.error:
|
|
pass
|
|
try:
|
|
self.socket.close()
|
|
except socket.error:
|
|
pass
|
|
self.socket = None
|
|
|
|
def is_offline(self):
|
|
"""try restart if needed"""
|
|
if self.socket is None:
|
|
now = time.time()
|
|
if now < self._next_connect:
|
|
return True
|
|
try:
|
|
self.connect()
|
|
self.init()
|
|
except Exception as e:
|
|
if now > self._deadline:
|
|
self.close()
|
|
self.dead = min(now, self._last_live + 1)
|
|
return True
|
|
if self._deadline == INF:
|
|
print(f'error "{e}" connecting to {self.uri} retrying for {self.max_offline} sec')
|
|
self._deadline = now + self.max_offline
|
|
else:
|
|
print('.', end='', flush=True)
|
|
self._next_connect = now + 0.5
|
|
return True
|
|
return False
|
|
|
|
def settimeout(self, timeout):
|
|
self.timeout = timeout
|
|
if self.socket:
|
|
self.socket.settimeout(timeout)
|
|
|
|
def notimeout(self):
|
|
if self.socket:
|
|
self.socket.settimeout(0)
|
|
|
|
def send(self, line):
|
|
try:
|
|
self.socket.sendall(line.encode(self.encoding))
|
|
self.socket.sendall(b'\n')
|
|
except Exception as e:
|
|
self.close()
|
|
e.args = ('send:' + e.args[0],)
|
|
raise
|
|
|
|
def get_lines(self):
|
|
"""generator returning lines as bytes"""
|
|
if self.is_offline():
|
|
return
|
|
now = time.time()
|
|
if now > self._next_ping:
|
|
if self._next_ping == self._ping_deadline:
|
|
print(self.uri, 'no pong')
|
|
self.close()
|
|
return
|
|
self.ping()
|
|
self._last_live = self._next_ping - self.ping_timeout
|
|
self._next_ping = self._ping_deadline = now + self._next_ping
|
|
buffer = self._buffer
|
|
while 1:
|
|
try:
|
|
received = self.socket.recv(NETBUFSIZE)
|
|
if not received:
|
|
raise ConnectionAbortedError('connection closed by other end')
|
|
except (TimeoutError, BlockingIOError):
|
|
break
|
|
except Exception:
|
|
self._last_live = now
|
|
self.close()
|
|
raise
|
|
splitted = received.split(b'\n')
|
|
if len(splitted) == 1:
|
|
buffer.append(received)
|
|
else:
|
|
self._next_ping = now + self.ping_timeout
|
|
buffer.append(splitted[0])
|
|
splitted[0] = b''.join(buffer)
|
|
buffer[:] = [splitted.pop()]
|
|
for line in splitted:
|
|
yield line.decode(self.encoding)
|
|
if len(received) < NETBUFSIZE and self.timeout == 0:
|
|
break
|
|
|
|
def event_generator(self):
|
|
raise NotImplementedError
|
|
|
|
def get_tags(self, key):
|
|
"""get tags for key"""
|
|
raise NotImplementedError
|
|
|
|
def finish_events(self, events, end_time):
|
|
for key in list(self.cache):
|
|
self.cache.pop(key)
|
|
dbkey = '.'.join(key)
|
|
events.append((True, None, end_time, dbkey, 'float', self.tags))
|
|
events.append((False, 'END', end_time, dbkey, 'error', self.tags))
|
|
|
|
def get_events(self, events, maxevents):
|
|
"""get available events
|
|
|
|
:param events: a list events to be appended to
|
|
:param maxevents: hint for max number of events to be joined
|
|
there might be more after a full hour or when the stream is dying
|
|
:return: True when maxevents is reached
|
|
"""
|
|
for key, value, error, ts, tags in self.generator:
|
|
ts = max(self.start_time, min(ts or INF, time.time()))
|
|
if ts >= self.next_hour:
|
|
ts_ = (ts // 3600) * 3600
|
|
for key_, value_ in self.cache.items():
|
|
events.append((True, value_, ts_, '.'.join(key_), 'float', self.get_tags(key_)))
|
|
for key_, error_ in self.errors.items():
|
|
events.append((False, error_, ts_, '.'.join(key_), 'error', self.get_tags(key_)))
|
|
self.next_hour = ts_ + 3600
|
|
if value != self.cache.get(key, None) or error != self.errors.get(key, None):
|
|
dbkey = '.'.join(key)
|
|
events.append((True, value, ts, dbkey, 'float', tags))
|
|
self.cache[key] = value
|
|
if error and self.errors.get(key) != error:
|
|
events.append((False, error, ts, dbkey, 'error', tags))
|
|
self.errors[key] = error
|
|
elif len(events) >= maxevents:
|
|
return True
|
|
else:
|
|
if self.dead:
|
|
self.finish_events(events, self.dead)
|
|
raise StreamDead()
|
|
self.generator = self.event_generator()
|
|
return False
|
|
|
|
|
|
class EventStream:
|
|
return_on_wait = False
|
|
# return_on_wait = True: stop generator when no more streams have buffered content
|
|
# note: a stream with buffered content might not be ready to emit any event, because
|
|
# of filtering
|
|
|
|
def __init__(self, *udp, **streams):
|
|
self.streams = streams
|
|
self.udp = {v.socket.fileno(): v for v in udp}
|
|
|
|
def wait_ready(self, timeout):
|
|
ready = select(Stream.select_dict, [], [], timeout)[0]
|
|
return [Stream.select_dict[f] for f in ready]
|
|
|
|
def get_events(self, maxevents=20):
|
|
"""return events from all streams
|
|
|
|
:param maxevents: hint for max number of events to be joined
|
|
there might be more after a full hour or when the stream is dying
|
|
:return: list of events
|
|
|
|
wait for at least one event
|
|
"""
|
|
events = []
|
|
while 1:
|
|
for stream in self.wait_ready(1):
|
|
if not isinstance(stream, Stream):
|
|
for streamcls, uri, *args in stream.events():
|
|
if uri not in self.streams:
|
|
print('add stream', uri, *args)
|
|
self.streams[uri] = streamcls(uri, *args)
|
|
for name, stream in self.streams.items():
|
|
try:
|
|
if stream.get_events(events, maxevents):
|
|
return events
|
|
except StreamDead:
|
|
self.streams.pop(name)
|
|
if events:
|
|
return events
|
|
if events:
|
|
return events
|
|
|
|
def finish(self):
|
|
events = []
|
|
end_time = time.time()
|
|
for stream in self.streams.values():
|
|
stream.close()
|
|
stream.finish_events(events, end_time)
|
|
return events
|