You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
196 lines
6.7 KiB
196 lines
6.7 KiB
import socket
|
|
import select
|
|
import sys
|
|
import logging
|
|
|
|
from .wire import decode_varint, encode
|
|
from .reader import BytesBuffer
|
|
from .msg import RequestDecoder, message_types
|
|
|
|
# hold the asyncronous state of a connection
|
|
# ie. we may not get enough bytes on one read to decode the message
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class Connection():
|
|
|
|
def __init__(self, fd, app):
|
|
self.fd = fd
|
|
self.app = app
|
|
self.recBuf = BytesBuffer(bytearray())
|
|
self.resBuf = BytesBuffer(bytearray())
|
|
self.msgLength = 0
|
|
self.decoder = RequestDecoder(self.recBuf)
|
|
self.inProgress = False # are we in the middle of a message
|
|
|
|
def recv(this):
|
|
data = this.fd.recv(1024)
|
|
if not data: # what about len(data) == 0
|
|
raise IOError("dead connection")
|
|
this.recBuf.write(data)
|
|
|
|
# ABCI server responds to messges by calling methods on the app
|
|
|
|
class ABCIServer():
|
|
|
|
def __init__(self, app, port=5410):
|
|
self.app = app
|
|
# map conn file descriptors to (app, reqBuf, resBuf, msgDecoder)
|
|
self.appMap = {}
|
|
|
|
self.port = port
|
|
self.listen_backlog = 10
|
|
|
|
self.listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
self.listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
|
self.listener.setblocking(0)
|
|
self.listener.bind(('', port))
|
|
|
|
self.listener.listen(self.listen_backlog)
|
|
|
|
self.shutdown = False
|
|
|
|
self.read_list = [self.listener]
|
|
self.write_list = []
|
|
|
|
def handle_new_connection(self, r):
|
|
new_fd, new_addr = r.accept()
|
|
new_fd.setblocking(0) # non-blocking
|
|
self.read_list.append(new_fd)
|
|
self.write_list.append(new_fd)
|
|
print('new connection to', new_addr)
|
|
|
|
self.appMap[new_fd] = Connection(new_fd, self.app)
|
|
|
|
def handle_conn_closed(self, r):
|
|
self.read_list.remove(r)
|
|
self.write_list.remove(r)
|
|
r.close()
|
|
print("connection closed")
|
|
|
|
def handle_recv(self, r):
|
|
# app, recBuf, resBuf, conn
|
|
conn = self.appMap[r]
|
|
while True:
|
|
try:
|
|
print("recv loop")
|
|
# check if we need more data first
|
|
if conn.inProgress:
|
|
if (conn.msgLength == 0 or conn.recBuf.size() < conn.msgLength):
|
|
conn.recv()
|
|
else:
|
|
if conn.recBuf.size() == 0:
|
|
conn.recv()
|
|
|
|
conn.inProgress = True
|
|
|
|
# see if we have enough to get the message length
|
|
if conn.msgLength == 0:
|
|
ll = conn.recBuf.peek()
|
|
if conn.recBuf.size() < 1 + ll:
|
|
# we don't have enough bytes to read the length yet
|
|
return
|
|
print("decoding msg length")
|
|
conn.msgLength = decode_varint(conn.recBuf)
|
|
|
|
# see if we have enough to decode the message
|
|
if conn.recBuf.size() < conn.msgLength:
|
|
return
|
|
|
|
# now we can decode the message
|
|
|
|
# first read the request type and get the particular msg
|
|
# decoder
|
|
typeByte = conn.recBuf.read(1)
|
|
typeByte = int(typeByte[0])
|
|
resTypeByte = typeByte + 0x10
|
|
req_type = message_types[typeByte]
|
|
|
|
if req_type == "flush":
|
|
# messages are length prefixed
|
|
conn.resBuf.write(encode(1))
|
|
conn.resBuf.write([resTypeByte])
|
|
conn.fd.send(conn.resBuf.buf)
|
|
conn.msgLength = 0
|
|
conn.inProgress = False
|
|
conn.resBuf = BytesBuffer(bytearray())
|
|
return
|
|
|
|
decoder = getattr(conn.decoder, req_type)
|
|
|
|
print("decoding args")
|
|
req_args = decoder()
|
|
print("got args", req_args)
|
|
|
|
# done decoding message
|
|
conn.msgLength = 0
|
|
conn.inProgress = False
|
|
|
|
req_f = getattr(conn.app, req_type)
|
|
if req_args is None:
|
|
res = req_f()
|
|
elif isinstance(req_args, tuple):
|
|
res = req_f(*req_args)
|
|
else:
|
|
res = req_f(req_args)
|
|
|
|
if isinstance(res, tuple):
|
|
res, ret_code = res
|
|
else:
|
|
ret_code = res
|
|
res = None
|
|
|
|
print("called", req_type, "ret code:", ret_code, 'res:', res)
|
|
if ret_code != 0:
|
|
print("non-zero retcode:", ret_code)
|
|
|
|
if req_type in ("echo", "info"): # these dont return a ret code
|
|
enc = encode(res)
|
|
# messages are length prefixed
|
|
conn.resBuf.write(encode(len(enc) + 1))
|
|
conn.resBuf.write([resTypeByte])
|
|
conn.resBuf.write(enc)
|
|
else:
|
|
enc, encRet = encode(res), encode(ret_code)
|
|
# messages are length prefixed
|
|
conn.resBuf.write(encode(len(enc) + len(encRet) + 1))
|
|
conn.resBuf.write([resTypeByte])
|
|
conn.resBuf.write(encRet)
|
|
conn.resBuf.write(enc)
|
|
except IOError as e:
|
|
print("IOError on reading from connection:", e)
|
|
self.handle_conn_closed(r)
|
|
return
|
|
except Exception as e:
|
|
logger.exception("error reading from connection")
|
|
self.handle_conn_closed(r)
|
|
return
|
|
|
|
def main_loop(self):
|
|
while not self.shutdown:
|
|
r_list, w_list, _ = select.select(
|
|
self.read_list, self.write_list, [], 2.5)
|
|
|
|
for r in r_list:
|
|
if (r == self.listener):
|
|
try:
|
|
self.handle_new_connection(r)
|
|
# undo adding to read list ...
|
|
except NameError as e:
|
|
print("Could not connect due to NameError:", e)
|
|
except TypeError as e:
|
|
print("Could not connect due to TypeError:", e)
|
|
except:
|
|
print("Could not connect due to unexpected error:", sys.exc_info()[0])
|
|
else:
|
|
self.handle_recv(r)
|
|
|
|
def handle_shutdown(self):
|
|
for r in self.read_list:
|
|
r.close()
|
|
for w in self.write_list:
|
|
try:
|
|
w.close()
|
|
except Exception as e:
|
|
print(e) # TODO: add logging
|
|
self.shutdown = True
|