Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Refactor net gevent #18

Closed
wants to merge 1 commit into from
Closed

Refactor net gevent #18

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
180 changes: 27 additions & 153 deletions rebirthdb/gevent_net/net_gevent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,29 +15,31 @@
# This file incorporates work covered by the following copyright:
# Copyright 2010-2016 RethinkDB, all rights reserved.

import errno
import ssl
import struct

import gevent
import gevent.socket as socket
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the reason I noted waiting on tests below. The documentation for this class says it will only block one greenlet and that is essential to being async. I would rather wait for tests. Otherwise socket wrapper should work like connection where it operates on an internal instance object which can be overridden in the subclass constructor. Long term this pattern is a good place for metaclass factories.

from gevent.event import AsyncResult, Event
from gevent.lock import Semaphore

from rebirthdb import net, ql2_pb2
from rebirthdb.errors import ReqlAuthError, ReqlCursorEmpty, ReqlDriverError, ReqlTimeoutError, RqlDriverError, \
RqlTimeoutError
from rebirthdb.logger import default_logger
from rebirthdb.errors import ReqlCursorEmpty, RqlDriverError, RqlTimeoutError
from rebirthdb.net import SocketWrapper


__all__ = ['Connection']

pResponse = ql2_pb2.Response.ResponseType
pQuery = ql2_pb2.Query.QueryType
PROTO_RESPONSE_TYPE = ql2_pb2.Response.ResponseType
PROTO_QUERY_TYPE = ql2_pb2.Query.QueryType


class GeventCursorEmpty(ReqlCursorEmpty, StopIteration):
pass


class Connection(net.Connection):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This would still need the old init where it provides the instance class.

pass


# TODO: allow users to set sync/async?
class GeventCursor(net.Cursor):
def __init__(self, *args, **kwargs):
Expand Down Expand Up @@ -68,147 +70,23 @@ def _get_next(self, timeout):
return self.items.popleft()


# TODO: would be nice to share this code with net.py
# TODO(grandquista): code seems to already be a duplicate of superclass
# revisit this after testing is inplace.
class SocketWrapper(net.SocketWrapper):
def __init__(self, parent):
self.host = parent._parent.host
self.port = parent._parent.port
self._read_buffer = None
self._socket = None
self.ssl = parent._parent.ssl

try:
self._socket = socket.create_connection((self.host, self.port))
self._socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)

if len(self.ssl) > 0:
try:
if hasattr(ssl, 'SSLContext'): # Python2.7 and 3.2+, or backports.ssl
ssl_context = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
if hasattr(ssl_context, "options"):
ssl_context.options |= getattr(ssl, "OP_NO_SSLv2", 0)
ssl_context.options |= getattr(ssl, "OP_NO_SSLv3", 0)
self.ssl_context.verify_mode = ssl.CERT_REQUIRED
self.ssl_context.check_hostname = True # redundant with match_hostname
self.ssl_context.load_verify_locations(self.ssl["ca_certs"])
self._socket = ssl_context.wrap_socket(self._socket, server_hostname=self.host)
else: # this does not disable SSLv2 or SSLv3
self._socket = ssl.wrap_socket(
self._socket, cert_reqs=ssl.CERT_REQUIRED, ssl_version=ssl.PROTOCOL_SSLv23,
ca_certs=self.ssl["ca_certs"])
except IOError as exc:
self._socket.close()
raise ReqlDriverError("SSL handshake failed (see server log for more information): %s" % str(exc))
try:
ssl.match_hostname(self._socket.getpeercert(), hostname=self.host)
except ssl.CertificateError:
self._socket.close()
raise

parent._parent.handshake.reset()
response = None
while True:
request = parent._parent.handshake.next_message(response)
if request is None:
break
# This may happen in the `V1_0` protocol where we send two requests as
# an optimization, then need to read each separately
if request is not "":
self.sendall(request)

# The response from the server is a null-terminated string
response = b''
while True:
char = self.recvall(1)
if char == b'\0':
break
response += char
except (ReqlAuthError, ReqlTimeoutError):
self.close()
raise
except ReqlDriverError as ex:
self.close()
error = str(ex) \
.replace('receiving from', 'during handshake with') \
.replace('sending to', 'during handshake with')
raise ReqlDriverError(error)
except Exception as ex:
self.close()
raise ReqlDriverError("Could not connect to %s:%s. Error: %s" %
(self.host, self.port, ex))

def close(self):
if self._socket is not None:
try:
self._socket.shutdown(socket.SHUT_RDWR)
self._socket.close()
except Exception as ex:
default_logger.error(ex.message)
finally:
self._socket = None

def recvall(self, length):
res = b'' if self._read_buffer is None else self._read_buffer
while len(res) < length:
while True:
try:
chunk = self._socket.recv(length - len(res))
break
except ReqlTimeoutError:
raise
except IOError as ex:
if ex.errno == errno.ECONNRESET:
self.close()
raise ReqlDriverError("Connection is closed.")
elif ex.errno != errno.EINTR:
self.close()
raise ReqlDriverError(
'Connection interrupted receiving from %s:%s - %s' % (self.host, self.port, str(ex))
)
except Exception as ex:
self.close()
raise ReqlDriverError('Error receiving from %s:%s - %s' % (self.host, self.port, str(ex)))
if len(chunk) == 0:
self.close()
raise ReqlDriverError("Connection is closed.")
res += chunk
return res

def sendall(self, data):
offset = 0
while offset < len(data):
try:
offset += self._socket.send(data[offset:])
except IOError as ex:
if ex.errno == errno.ECONNRESET:
self.close()
raise ReqlDriverError("Connection is closed.")
elif ex.errno != errno.EINTR:
self.close()
raise ReqlDriverError(('Connection interrupted ' +
'sending to %s:%s - %s') %
(self.host, self.port, str(ex)))
except Exception as ex:
self.close()
raise ReqlDriverError('Error sending to %s:%s - %s' % (self.host, self.port, str(ex)))


class ConnectionInstance(object):
def __init__(self, parent, io_loop=None):
def __init__(self, parent):
self._parent = parent
self._closing = False
self._user_queries = {}
self._cursor_cache = {}

self._write_mutex = Semaphore()
self._socket = None
self.timeout = None

def connect(self, timeout):
if not self.timeout:
self.timeout = timeout

with gevent.Timeout(timeout, RqlTimeoutError(self._parent.host, self._parent.port)) as timeout:
self._socket = SocketWrapper(self)
self._socket = SocketWrapper(self, timeout)

# Start a parallel coroutine to perform reads
gevent.spawn(self._reader)
Expand All @@ -217,8 +95,9 @@ def connect(self, timeout):
def is_open(self):
return self._socket is not None and self._socket.is_open()

def close(self, noreply_wait=False, token=None, exception=None):
def close(self, no_reply_wait=False, token=None, exception=None):
self._closing = True

if exception is not None:
err_message = "Connection is closed (%s)." % str(exception)
else:
Expand All @@ -234,9 +113,9 @@ def close(self, noreply_wait=False, token=None, exception=None):
self._user_queries = {}
self._cursor_cache = {}

if noreply_wait:
noreply = net.Query(pQuery.NOREPLY_WAIT, token, None, None)
self.run_query(noreply, False)
if no_reply_wait:
no_reply = net.Query(PROTO_QUERY_TYPE.NOREPLY_WAIT, token, None, None)
self.run_query(no_reply, False)

try:
self._socket.close()
Expand Down Expand Up @@ -267,9 +146,9 @@ def run_query(self, query, noreply):
def _reader(self):
try:
while True:
buf = self._socket.recvall(12)
buf = self._socket.recvall(12, self.timeout)
(token, length,) = struct.unpack("<qL", buf)
buf = self._socket.recvall(length)
buf = self._socket.recvall(length, self.timeout)

cursor = self._cursor_cache.get(token)
if cursor is not None:
Expand All @@ -279,13 +158,13 @@ def _reader(self):
# we don't lose track of it in case of an exception
query, async_res = self._user_queries[token]
res = net.Response(token, buf, self._parent._get_json_decoder(query))
if res.type == pResponse.SUCCESS_ATOM:
if res.type == PROTO_RESPONSE_TYPE.SUCCESS_ATOM:
async_res.set(net.maybe_profile(res.data[0], res))
elif res.type in (pResponse.SUCCESS_SEQUENCE,
pResponse.SUCCESS_PARTIAL):
elif res.type in (PROTO_RESPONSE_TYPE.SUCCESS_SEQUENCE,
PROTO_RESPONSE_TYPE.SUCCESS_PARTIAL):
cursor = GeventCursor(self, query, res)
async_res.set(net.maybe_profile(cursor, res))
elif res.type == pResponse.WAIT_COMPLETE:
elif res.type == PROTO_RESPONSE_TYPE.WAIT_COMPLETE:
async_res.set(None)
else:
async_res.set_exception(res.make_error(query))
Expand All @@ -295,8 +174,3 @@ def _reader(self):
except Exception as ex:
if not self._closing:
self.close(exception=ex)


class Connection(net.Connection):
def __init__(self, *args, **kwargs):
super(Connection, self).__init__(ConnectionInstance, *args, **kwargs)