THRIFT-1704: Tornado support (Python)
diff --git a/lib/py/src/TTornado.py b/lib/py/src/TTornado.py
new file mode 100644
index 0000000..af309c3
--- /dev/null
+++ b/lib/py/src/TTornado.py
@@ -0,0 +1,153 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+
+from cStringIO import StringIO
+import logging
+import socket
+import struct
+
+from thrift.transport import TTransport
+from thrift.transport.TTransport import TTransportException
+
+from tornado import gen
+from tornado import iostream
+from tornado import netutil
+
+
+class TTornadoStreamTransport(TTransport.TTransportBase):
+ """a framed, buffered transport over a Tornado stream"""
+ def __init__(self, host, port, stream=None):
+ self.host = host
+ self.port = port
+ self.is_queuing_reads = False
+ self.read_queue = []
+ self.__wbuf = StringIO()
+
+ # servers provide a ready-to-go stream
+ self.stream = stream
+ if self.stream is not None:
+ self._set_close_callback()
+
+ # not the same number of parameters as TTransportBase.open
+ def open(self, callback):
+ logging.debug('socket connecting')
+ sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
+ self.stream = iostream.IOStream(sock)
+
+ def on_close_in_connect(*_):
+ message = 'could not connect to {}:{}'.format(self.host, self.port)
+ raise TTransportException(
+ type=TTransportException.NOT_OPEN,
+ message=message)
+ self.stream.set_close_callback(on_close_in_connect)
+
+ def finish(*_):
+ self._set_close_callback()
+ callback()
+
+ self.stream.connect((self.host, self.port), callback=finish)
+
+ def _set_close_callback(self):
+ def on_close():
+ raise TTransportException(
+ type=TTransportException.END_OF_FILE,
+ message='socket closed')
+ self.stream.set_close_callback(self.close)
+
+ def close(self):
+ # don't raise if we intend to close
+ self.stream.set_close_callback(None)
+ self.stream.close()
+
+ def read(self, _):
+ # The generated code for Tornado shouldn't do individual reads -- only
+ # frames at a time
+ assert "you're doing it wrong" is True
+
+ @gen.engine
+ def readFrame(self, callback):
+ self.read_queue.append(callback)
+ logging.debug('read queue: %s', self.read_queue)
+
+ if self.is_queuing_reads:
+ # If a read is already in flight, then the while loop below should
+ # pull it from self.read_queue
+ return
+
+ self.is_queuing_reads = True
+ while self.read_queue:
+ next_callback = self.read_queue.pop()
+ result = yield gen.Task(self._readFrameFromStream)
+ next_callback(result)
+ self.is_queuing_reads = False
+
+ @gen.engine
+ def _readFrameFromStream(self, callback):
+ logging.debug('_readFrameFromStream')
+ frame_header = yield gen.Task(self.stream.read_bytes, 4)
+ frame_length, = struct.unpack('!i', frame_header)
+ logging.debug('received frame header, frame length = %i', frame_length)
+ frame = yield gen.Task(self.stream.read_bytes, frame_length)
+ logging.debug('received frame payload')
+ callback(frame)
+
+ def write(self, buf):
+ self.__wbuf.write(buf)
+
+ def flush(self, callback=None):
+ wout = self.__wbuf.getvalue()
+ wsz = len(wout)
+ # reset wbuf before write/flush to preserve state on underlying failure
+ self.__wbuf = StringIO()
+ # N.B.: Doing this string concatenation is WAY cheaper than making
+ # two separate calls to the underlying socket object. Socket writes in
+ # Python turn out to be REALLY expensive, but it seems to do a pretty
+ # good job of managing string buffer operations without excessive copies
+ buf = struct.pack("!i", wsz) + wout
+
+ logging.debug('writing frame length = %i', wsz)
+ self.stream.write(buf, callback)
+
+
+class TTornadoServer(netutil.TCPServer):
+ def __init__(self, processor, iprot_factory, oprot_factory=None,
+ *args, **kwargs):
+ super(TTornadoServer, self).__init__(*args, **kwargs)
+
+ self._processor = processor
+ self._iprot_factory = iprot_factory
+ self._oprot_factory = (oprot_factory if oprot_factory is not None
+ else iprot_factory)
+
+ def handle_stream(self, stream, address):
+ try:
+ host, port = address
+ trans = TTornadoStreamTransport(host=host, port=port, stream=stream)
+ oprot = self._oprot_factory.getProtocol(trans)
+
+ def next_pass():
+ if not trans.stream.closed():
+ self._processor.process(trans, self._iprot_factory, oprot,
+ callback=next_pass)
+
+ next_pass()
+
+ except Exception:
+ logging.exception('thrift exception in handle_stream')
+ trans.close()