Port to AsyncIO

Author: awilfox

aerospike_py/__init__.py

@@ -1 +1 @@
-__version__ = '0.1.0'
+__version__ = '0.2.0'

aerospike_py/client.py

@@ -1,7 +1,7 @@
-import socket
+import asyncio
 import hashlib
 
-from aerospike_py.connection import SocketConnection
+from aerospike_py.connection import AsyncConnection
 from aerospike_py.info import request_info_keys
 from aerospike_py.result_code import ASMSGProtocolException
 from aerospike_py.message import ASIOException
@@ -16,11 +16,11 @@ def hash_key(set='', key=''):
 
 
 class AerospikeClient:
-    def __init__(self, sck):
-        self.sck = sck
+    def __init__(self, conn):
+        self.conn = conn
 
     def info(self, keys):
-        return request_info_keys(self.sck, keys)[1]
+        return request_info_keys(self.conn, keys)[1]
 
     def _process_bucket(self, asmsg_ops):
         buckets = {}
@@ -29,10 +29,11 @@ def _process_bucket(self, asmsg_ops):
 
         return buckets
 
+    @asyncio.coroutine
     def _submit_message(self, envelope, retry_count=3):
         while retry_count:
             try:
-                outer, asmsg_hdr, asmsg_fields, asmsg_ops = aerospike_py.message.submit_message(self.sck, envelope)
+                outer, asmsg_hdr, asmsg_fields, asmsg_ops = yield from aerospike_py.message.submit_message(self.conn, envelope)
                 return self._process_bucket(asmsg_ops)
             except ASMSGProtocolException as e:
                 if e.result_code not in (14,):
@@ -45,10 +46,11 @@ def _submit_message(self, envelope, retry_count=3):
 
         return buckets
 
+    @asyncio.coroutine
     def _submit_batch(self, envelope, retry_count=3):
         while retry_count:
             try:
-                messages = aerospike_py.message.submit_multi_message(self.sck, envelope)
+                messages = yield from aerospike_py.message.submit_multi_message(self.conn, envelope)
                 return [self._process_bucket(x[3]) for x in messages]
             except ASMSGProtocolException as e:
                 if e.result_code not in (14,):
@@ -61,6 +63,7 @@ def _submit_batch(self, envelope, retry_count=3):
 
         return buckets
 
+    @asyncio.coroutine
     def get(self, namespace, set='', key='', bins=[], record_ttl=0, retry_count=3):
         flags = aerospike_py.message.AS_INFO1_READ
         if not bins:
@@ -77,6 +80,7 @@ def get(self, namespace, set='', key='', bins=[], record_ttl=0, retry_count=3):
 
         return self._submit_message(envelope, retry_count)
 
+    @asyncio.coroutine
     def mget(self, namespace, groups=[], bins={}, record_ttl=0, retry_count=3):
         flags = aerospike_py.message.AS_INFO1_READ # | aerospike_py.message.AS_INFO1_BATCH
         if not bins:
@@ -98,6 +102,7 @@ def mget(self, namespace, groups=[], bins={}, record_ttl=0, retry_count=3):
 
         return self._submit_batch(envelope, retry_count)
 
+    @asyncio.coroutine
     def put(self, namespace, set='', key='', bins={}, create_only=False, bin_create_only=False, record_ttl=0, retry_count=3):
         flags = aerospike_py.message.AS_INFO2_WRITE
         if create_only:
@@ -118,6 +123,7 @@ def put(self, namespace, set='', key='', bins={}, create_only=False, bin_create_
 
         return self._submit_message(envelope, retry_count)
 
+    @asyncio.coroutine
     def delete(self, namespace, set='', key='', record_ttl=0, retry_count=3):
         envelope = aerospike_py.message.pack_asmsg(0, aerospike_py.message.AS_INFO2_WRITE | aerospike_py.message.AS_INFO2_DELETE, 0, 0, record_ttl, 0,
             [
@@ -129,6 +135,7 @@ def delete(self, namespace, set='', key='', record_ttl=0, retry_count=3):
 
         return self._submit_message(envelope, retry_count)
 
+    @asyncio.coroutine
     def incr(self, namespace, set='', key='', bin='', incr_by=0, record_ttl=0, retry_count=3):
         flags = aerospike_py.message.AS_INFO2_WRITE
 
@@ -143,6 +150,7 @@ def incr(self, namespace, set='', key='', bin='', incr_by=0, record_ttl=0, retry
 
         return self._submit_message(envelope, retry_count)
 
+    @asyncio.coroutine
     def _append_op(self, namespace, set='', key='', bin='', append_blob='', op=aerospike_py.message.AS_MSG_OP_APPEND, record_ttl=0, retry_count=3):
         flags = aerospike_py.message.AS_INFO2_WRITE
 
@@ -158,16 +166,19 @@ def _append_op(self, namespace, set='', key='', bin='', append_blob='', op=aeros
 
         return self._submit_message(envelope, retry_count)
 
+    @asyncio.coroutine
     def append(self, namespace, set='', key='', bin='', append_blob='', record_ttl=0):
         return self._append_op(namespace, set, key, bin, append_blob, aerospike_py.message.AS_MSG_OP_APPEND, record_ttl)
 
+    @asyncio.coroutine
     def prepend(self, namespace, set='', key='', bin='', append_blob='', record_ttl=0):
         return self._append_op(namespace, set, key, bin, append_blob, aerospike_py.message.AS_MSG_OP_PREPEND, record_ttl)
 
+    @asyncio.coroutine
     def touch(self, namespace, set, key, bin='', record_ttl=0):
         return self._append_op(namespace, set, key, bin, None, aerospike_py.message.AS_MSG_OP_TOUCH, record_ttl)
 
 
 def connect(host: str, port: int) -> AerospikeClient:
-    return AerospikeClient(SocketConnection(socket.create_connection((host, port))))
+    return AerospikeClient(AsyncConnection(host, port))
 

aerospike_py/connection.py

@@ -1,3 +1,10 @@
+import asyncio
+from logging import getLogger
+
+
+LOGGER = getLogger(__name__)
+
+
 class Connection:
     """Connection classes simply provide an interface specification for abstracting I/O.
     They can be used with Twisted, Eventlet, AsyncIO, etc. without problem.
@@ -9,19 +16,21 @@ def write(self, buf):
         pass
 
 
-class SocketConnection(Connection):
-    """SocketConnections are Connection instances which assume a functional Sockets API."""
-    def __init__(self, fd):
-        self._fd = fd
+class AsyncConnection(Connection):
+    """A Connection subclass which uses AsyncIO."""
+    def __init__(self, host: str, port: int):
+        conn = asyncio.open_connection(host, port)
+        loop = asyncio.get_event_loop()
+        try:
+            (self.reader, self.writer) = loop.run_until_complete(conn)
+        except OSError:
+            LOGGER.exception("Can't connect to Aerospike")
+            self.reader = self.writer = None
 
-    def read(self, length):
-        buf = bytearray(length)
-        view = memoryview(buf)
-        while length > 0:
-            nbytes = self._fd.recv_into(view, length)
-            view = view[nbytes:]
-            length -= nbytes
-        return buf
+    def read(self, length: int):
+        data = yield from self.reader.readexactly(length)
+        return data
 
     def write(self, buf):
-        return self._fd.send(buf)
+        self.writer.write(buf)
+        yield from self.writer.drain()

aerospike_py/info.py

@@ -1,15 +1,20 @@
+import asyncio
+
 from aerospike_py.connection import Connection
 from aerospike_py.message import pack_message, unpack_message, AerospikeOuterHeader
 
 
+@asyncio.coroutine
 def request_info_keys(conn: Connection, commands: list) -> (AerospikeOuterHeader, dict):
     payload = pack_message('\n'.join(commands).encode('UTF-8'), 1)
     conn.write(payload)
 
-    hdr_payload = conn.read(8)
+    hdr_payload = yield from conn.read(8)
     header, _ = unpack_message(hdr_payload)
 
-    header, payload = unpack_message(hdr_payload + conn.read(header.sz))
+    message = hdr_payload
+    message += yield from conn.read(header.sz)
+    header, payload = unpack_message(message)
     lines = payload.decode('UTF-8')
 
     infokeys = {}

aerospike_py/message.py

@@ -1,3 +1,4 @@
+import asyncio
 from collections import namedtuple
 import struct
 
@@ -223,18 +224,20 @@ def unpack_asmsg(data: bytes) -> (AerospikeASMSGHeader, list, list):
     return asmsg_hdr, fields, ops, data[pos:]
 
 
+@asyncio.coroutine
 def submit_message(conn: Connection, data: bytes) -> (AerospikeOuterHeader, AerospikeASMSGHeader, list, list):
     ohdr = AerospikeOuterHeader(2, 3, len(data))
     buf = pack_outer_header(ohdr) + data
-    conn.write(buf)
+    yield from conn.write(buf)
 
-    hdr_payload = conn.read(8)
+    hdr_payload = yield from conn.read(8)
     if not hdr_payload:
         raise ASIOException('read')
 
     header, _ = unpack_message(hdr_payload)
 
-    data = hdr_payload + conn.read(header.sz)
+    data = hdr_payload
+    data += yield from conn.read(header.sz)
 
     header, payload = unpack_message(data)
     asmsg_header, asmsg_fields, asmsg_ops, _ = unpack_asmsg(payload)
@@ -245,22 +248,24 @@ def submit_message(conn: Connection, data: bytes) -> (AerospikeOuterHeader, Aero
     return header, asmsg_header, asmsg_fields, asmsg_ops
 
 
+@asyncio.coroutine
 def submit_multi_message(conn: Connection, data: bytes) -> list:
     ohdr = AerospikeOuterHeader(2, 3, len(data))
     buf = pack_outer_header(ohdr) + data
-    conn.write(buf)
+    yield from conn.write(buf)
 
     not_last = True
     messages = []
 
     while not_last:
-        hdr_payload = conn.read(8)
+        hdr_payload = yield from conn.read(8)
         header, _ = unpack_message(hdr_payload)
 
         if not hdr_payload:
             raise ASIOException('read')
 
-        data = hdr_payload + conn.read(header.sz)
+        data = hdr_payload
+        data += yield from conn.read(header.sz)
 
         header, payload = unpack_message(data)
         while payload: