non-blocking `SocketService`
paulfantom committed
Jul 18, 2016 at 22:56 UTC
7c78678d6f065d874b8c37b64aeb75f37ba4a2d9
3 files changed
+146
-69
python.d/python_modules/base.py
+113
-62
@@ -6,6 +6,7 @@ import time
6
import sys
7
import os
8
import socket
9
+import select
10
try:
11
import urllib.request as urllib2
12
except ImportError:
@@ -479,92 +480,142 @@ class SocketService(SimpleService):
480
self.port = None
481
self.unix_socket = None
482
self.request = ""
483
+ self.__socket_config = None
484
SimpleService.__init__(self, configuration=configuration, name=name)
485
484
- def _get_raw_data(self):
486
+ def _connect(self):
487
"""
486
- Get raw data with low-level "socket" module.
487
- :return: str
488
+ Recreate socket and connect to it since they cannot be reused
489
+ Available configurations are IPv6, IPv4 or UNIX socket
490
+ :return:
491
"""
489
- # Recreate socket since they cannot be reused
490
- if self._sock is None:
491
- try:
492
- if self.unix_socket is None:
493
- self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
494
- self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
495
- #sock.setsockopt(socket.SOL_SOCKET, socket.TCP_NODELAY, 1)
496
- #sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
497
- #sock.settimeout(self.update_every)
498
- self._sock.settimeout(0.5)
499
- self._sock.connect((self.host, self.port))
500
- self._sock.settimeout(0.5) # Just to be sure
492
+ try:
493
+ if self.unix_socket is None:
494
+ if self.__socket_config is None:
495
+ # establish ipv6 or ipv4 connection.
496
+ for res in socket.getaddrinfo(self.host, self.port, socket.AF_UNSPEC, socket.SOCK_STREAM):
497
+ try:
498
+ af, socktype, proto, canonname, sa = res
499
+ self._sock = socket.socket(af, socktype, proto)
500
+ except socket.error as msg:
501
+ self._sock = None
502
+ continue
503
+ try:
504
+ self._sock.connect(sa)
505
+ except socket.error as msg:
506
+ self._disconnect()
507
+ continue
508
+ self.__socket_config = res
509
+ break
510
else:
502
- self._sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
503
- #self._sock.settimeout(self.update_every)
504
- self._sock.settimeout(0.05)
505
- self._sock.connect(self.unix_socket)
506
- self._sock.settimeout(0.05) # Just to be sure
507
- except Exception as e:
508
- self.error(str(e),
509
- "Cannot create socket with following configuration: host:", str(self.host),
510
- "port:", str(self.port),
511
- "socket:", str(self.unix_socket))
512
- self._sock = None
513
- return None
511
+ # connect to socket with previously established configuration
512
+ try:
513
+ af, socktype, proto, canonname, sa = self.__socket_config
514
+ self._sock = socket.socket(af, socktype, proto)
515
+ self._sock.connect(sa)
516
+ except socket.error as msg:
517
+ self._disconnect()
518
+ else:
519
+ # connect to unix socket
520
+ self._sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
521
+ self._sock.connect(self.unix_socket)
522
+ except Exception as e:
523
+ self.error(str(e),
524
+ "Cannot create socket with following configuration: host:", str(self.host),
525
+ "port:", str(self.port),
526
+ "socket:", str(self.unix_socket))
527
+ self._sock = None
528
+ self._sock.setblocking(0)
529
+
530
+ def _disconnect(self):
531
+ """
532
+ Close socket connection
533
+ :return:
534
+ """
535
+ try:
536
+ self._sock.shutdown(2) # 0 - read, 1 - write, 2 - all
537
+ self._sock.close()
538
+ except Exception:
539
+ pass
540
+ self._sock = None
541
542
+ def _send(self):
543
+ """
544
+ Send request.
545
+ :return: boolean
546
+ """
547
# Send request if it is needed
548
if self.request != "".encode():
549
try:
550
self._sock.send(self.request)
551
except Exception as e:
520
- try:
521
- self._sock.shutdown(0)
522
- self._sock.close()
523
- self._sock = None
524
- except:
525
- pass
552
+ self._disconnect()
553
self.error(str(e),
554
"used configuration: host:", str(self.host),
555
"port:", str(self.port),
556
"socket:", str(self.unix_socket))
530
- return None
531
-
532
- # Receive first two bytes from socket. This will test if there is anything to receive
533
- size = 2
534
- try:
535
- data = self._sock.recv(size).decode()
536
- except Exception as e:
537
- self.error(str(e),
538
- "used configuration: host:", str(self.host),
539
- "port:", str(self.port),
540
- "socket:", str(self.unix_socket))
541
- self._sock.shutdown(0)
542
- self._sock.close()
543
- self._sock = None
544
- return None
557
+ return False
558
+ return True
559
546
- # Receive the rest
560
+ def _receive(self):
561
+ """
562
+ Receive data from socket
563
+ :return: str
564
+ """
565
+ data = ""
566
while True:
548
- # implement something like TCP Window Scaling
549
- if size < 4096:
550
- size *= 2
551
- buf = self._sock.recv(size)
552
- data += buf.decode()
553
- if len(buf) == 0: break # precaution
554
- if len(buf) < size and not self._more_data_available():
567
+ try:
568
+ ready_to_read, _, in_error = select.select([self._sock], [], [], 0.01)
569
+ except Exception as e:
570
+ self.debug("SELECT", str(e))
571
+ self._disconnect()
572
break
573
+ if len(ready_to_read) > 0:
574
+ buf = self._sock.recv(4096)
575
+ if len(buf) == 0 or buf is None:
576
+ break
577
+ data += buf.decode()
578
else:
557
- size = 4
579
+ self._disconnect()
580
+ break
581
+
582
+ return data
583
+
584
+ def _get_raw_data(self):
585
+ """
586
+ Get raw data with low-level "socket" module.
587
+ :return: str
588
+ """
589
+ if self._sock is None:
590
+ self._connect()
591
+
592
+ # Send request if it is needed
593
+ if not self._send():
594
+ return None
595
+
596
+ finished = False
597
+ data = ""
598
+ prevent_infinite_loop = 1000000
599
+ while not finished:
600
+ data += self._receive()
601
+ finished = self._check_raw_data(data)
602
+ prevent_infinite_loop -= 1
603
+ if prevent_infinite_loop <= 0:
604
+ self.debug("Almost got into infinite loop while grabbing data. Is _check_raw_data() ok?")
605
+ break
606
607
if not self._keep_alive:
560
- self._sock.shutdown(0)
561
- self._sock.close()
562
- self._sock = None
608
+ self._disconnect()
609
610
return data
611
566
- def _more_data_available(self):
567
- return False
612
+ def _check_raw_data(self, data):
613
+ """
614
+ Check if all data has been gathered from socket
615
+ :param data: str
616
+ :return: boolean
617
+ """
618
+ return True
619
620
def _parse_config(self):
621
"""
python.d/redis.chart.py
+19
-1
@@ -100,6 +100,24 @@ class Service(SocketService):
100
else:
101
return data
102
103
+ def _check_raw_data(self, data):
104
+ """
105
+ Check if all data has been gathered from socket.
106
+ Parse first line containing message length and check against received message
107
+ :param data: str
108
+ :return: boolean
109
+ """
110
+ length = len(data)
111
+ supposed = data.split('\n')[0][1:]
112
+ offset = len(supposed) + 4 # 1 dollar sing, 1 new line character + 1 ending sequence '\r\n'
113
+ supposed = int(supposed)
114
+ if length - offset >= supposed:
115
+ return True
116
+ else:
117
+ return False
118
+
119
+ return False
120
+
121
def check(self):
122
"""
123
Parse configuration, check if redis is available, and dynamically create chart lines data
@@ -108,7 +126,7 @@ class Service(SocketService):
126
self._parse_config()
127
if self.name == "":
128
self.name = "local"
111
- self.chart_name += "_" + self.name
129
+ self.chart_name += "_" + self.name
130
data = self._get_data()
131
if data is None:
132
return False
python.d/squid.chart.py
+14
-6
@@ -72,22 +72,30 @@ class Service(SocketService):
72
self.error("invalid data received")
73
return None
74
75
+ self.debug("DATA:", str(len(raw)))
76
if len(data) == 0:
77
self.error("no data received")
78
return None
79
else:
80
return data
81
81
- def _more_data_available(self):
82
- try:
83
- ready_to_read, _, in_error = select.select([self._sock, ], [], [], 0.05)
84
- except Exception as e:
85
- self.debug("select returned exception", str(e))
86
- if len(ready_to_read) > 0:
82
+ def _check_raw_data(self, data):
83
+ if "Content-Length" not in data[:1024]: # assuming headers should be in first 1024 bytes (performance)
84
+ return True # "Content-Length" not found, assume everything is ok
85
+
86
+ # TODO write some parser of "Content-Length"
87
+ return True
88
+ length = len(data)
89
+
90
+
91
+ supposed = 0
92
+ if length >= supposed:
93
return True
94
else:
95
return False
96
97
+ return False
98
+
99
def check(self):
100
"""
101
Parse essential configuration, autodetect squid configuration (if needed), and check if data is available