This repository has been archived by the owner on Nov 18, 2022. It is now read-only.
/
influx.py
151 lines (119 loc) · 4.86 KB
/
influx.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
"""
Handle the Influxdb transactions and publications that need to occur.
"""
from general import read_credentials
from settings import env
from typing import List, Dict, Union, Any
from influxdb import InfluxDBClient
from contextlib import AbstractContextManager
from queue import LifoQueue
from threading import Thread
from time import sleep
from requests.exceptions import ConnectionError
measureT = List[Dict[str, Union[str, Dict[str, Union[Dict[str, float], str]]]]]
class InfluxPublisher(AbstractContextManager):
"""
Manage the publication of data points to Influxdb.
Synchronous and potentially lossy method of publishing points via InfluxDBClient.
"""
def __init__(self, credentials: str) -> None:
self._client = InfluxDBClient(**read_credentials(credentials))
def publish(self, js: measureT) -> None:
"""
Publish a metric to the configured Influx DB.
Args:
js: the data to publish (a list of measurements).
Returns:
Nothing.
"""
self._client.write_points(js)
def __enter__(self) -> 'InfluxPublisher':
return self
def __exit__(self, *args: Any) -> None:
self._client.close()
class DaemonInfluxPublisher:
"""
Provide an interface to a daemon that periodically flushes a queue of data points.
"""
def __init__(self, credentials: str) -> None:
self.queue = LifoQueue(maxsize=env['QUEUE_MAX'])
self._n_threads = 0
# Start a non-blocking daemon thread that's created and periodically reads
# from the queue, as well as spins up additional threads.
self._publish_daemon = Thread(target=self._d_publisher, daemon=True)
self._publish_daemon.start()
# Create a client connection, with which we intend to publish measurements.
self._client = InfluxDBClient(**read_credentials(credentials))
def publish(self, js: measureT) -> None:
"""
Add a metric to the queue for the daemon to flush periodically. I've kept the
interface the same as `publish` on `InfluxPublisher`, above.
Args:
js: the data to publish (a list of measurements).
Returns:
Nothing.
"""
self.queue.put(js, timeout=env['QUEUE_TIMEOUT'])
def _d_publisher(self) -> None:
"""
'Daemonized' publisher that has the ability to spin up other threads.
Periodically wakes up and tries to empty its queue of measurements.
This method is not meant to be called directly.
Returns:
Nothing.
"""
while True:
if not self.queue.empty():
# No reason to create more threads than necessary, here.
if env['MAX_THREADS'] > self.queue.qsize():
self._n_threads = self.queue.qsize()
division = 1
remainder = 0
else:
self._n_threads = env['MAX_THREADS']
division = self.queue.qsize() // self._n_threads
remainder = self.queue.qsize() % self._n_threads
self._spawn_threads(division, remainder)
sleep(env['DAEMON_WAKEUP'])
def _spawn_threads(self, division: int, remainder: int =0) -> None:
"""
Make publication threads, and return if there's an error during the request.
Args:
division: number of thread creation/join cycles to make during drain.
remainder: number of threads to create in the very last cycle.
Returns:
Nothing.
"""
for _i in [division, remainder]:
for _ in range(_i):
q_size_start = self.queue.qsize()
threads = []
for _ in range(self._n_threads):
new_thread = Thread(target=self.t_publish, args=())
new_thread.start()
threads.append(new_thread)
for thread in threads:
thread.join()
# If the ending queue size is the same as we started, just wait until
# the next publish cycle to try again, they're obviously failing.
if self.queue.qsize() == q_size_start:
return None
def t_publish(self) -> None:
"""
Publish a list of measurements to the open client connection on an influxdb.
Args:
p: the measures to publish.
Returns:
Nothing.
"""
p: measureT = self.queue.get()
try:
self._client.write_points(p)
except ConnectionError:
# Just put the item back on the queue lol.
self.queue.put(p)
def __enter__(self) -> 'DaemonInfluxPublisher':
return self
def __exit__(self, *args: Any) -> None:
# TODO: figure out the correct way to kill the daemon safely?
pass