lang/py/src/avro/txipc.py

#!/usr/bin/env python

##
# 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
#
#     https://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 __future__ import absolute_import, division, print_function

import io

from zope.interface import implements

import avro.io
from avro import ipc
from twisted.internet.defer import Deferred, maybeDeferred
from twisted.internet.protocol import Protocol
from twisted.web import resource, server
from twisted.web.client import Agent
from twisted.web.http_headers import Headers
from twisted.web.iweb import IBodyProducer


class TwistedRequestor(ipc.BaseRequestor):
  """A Twisted-compatible requestor. Returns a Deferred that will fire with the
     returning value, instead of blocking until the request completes."""
  def _process_handshake(self, call_response, message_name, request_datum):
    # process the handshake and call response
    buffer_decoder = avro.io.BinaryDecoder(io.BytesIO(call_response))
    call_response_exists = self.read_handshake_response(buffer_decoder)
    if call_response_exists:
      return self.read_call_response(message_name, buffer_decoder)
    else:
      return self.request(message_name, request_datum)

  def issue_request(self, call_request, message_name, request_datum):
    d = self.transceiver.transceive(call_request)
    d.addCallback(self._process_handshake, message_name, request_datum)
    return d

class RequestStreamingProducer(object):
  """A streaming producer for issuing requests with the Twisted.web Agent."""
  implements(IBodyProducer)

  paused = False
  stopped = False
  started = False

  def __init__(self, message):
    self._message = message
    self._length = len(message)
    # We need a buffer length header for every buffer and an additional
    # zero-length buffer as the message terminator
    self._length += (self._length / ipc.BUFFER_SIZE + 2) \
      * ipc.BUFFER_HEADER_LENGTH
    self._total_bytes_sent = 0
    self._deferred = Deferred()

  # read-only properties
  message = property(lambda self: self._message)
  length = property(lambda self: self._length)
  consumer = property(lambda self: self._consumer)
  deferred = property(lambda self: self._deferred)

  def _get_total_bytes_sent(self):
    return self._total_bytes_sent

  def _set_total_bytes_sent(self, bytes_sent):
    self._total_bytes_sent = bytes_sent

  total_bytes_sent = property(_get_total_bytes_sent, _set_total_bytes_sent)

  def startProducing(self, consumer):
    if self.started:
      return

    self.started = True
    self._consumer = consumer
    # Keep writing data to the consumer until we're finished,
    # paused (pauseProducing()) or stopped (stopProducing())
    while self.length - self.total_bytes_sent > 0 and \
      not self.paused and not self.stopped:
      self.write()
    # self.write will fire this deferred once it has written
    # the entire message to the consumer
    return self.deferred

  def resumeProducing(self):
    self.paused = False
    self.write(self)

  def pauseProducing(self):
    self.paused = True

  def stopProducing(self):
    self.stopped = True

  def write(self):
    if self.length - self.total_bytes_sent > ipc.BUFFER_SIZE:
      buffer_length = ipc.BUFFER_SIZE
    else:
      buffer_length = self.length - self.total_bytes_sent
    self.write_buffer(self.message[self.total_bytes_sent:
                              (self.total_bytes_sent + buffer_length)])
    self.total_bytes_sent += buffer_length
    # Make sure we wrote the entire message
    if self.total_bytes_sent == self.length and not self.stopped:
      self.stopProducing()
      # A message is always terminated by a zero-length buffer.
      self.write_buffer_length(0)
      self.deferred.callback(None)

  def write_buffer(self, chunk):
    buffer_length = len(chunk)
    self.write_buffer_length(buffer_length)
    self.consumer.write(chunk)

  def write_buffer_length(self, n):
    self.consumer.write(ipc.BIG_ENDIAN_INT_STRUCT.pack(n))

class AvroProtocol(Protocol):

  recvd = ''
  done = False

  def __init__(self, finished):
    self.finished = finished
    self.message = []

  def dataReceived(self, data):
    self.recvd = self.recvd + data
    while len(self.recvd) >= ipc.BUFFER_HEADER_LENGTH:
      buffer_length ,= ipc.BIG_ENDIAN_INT_STRUCT.unpack(
        self.recvd[:ipc.BUFFER_HEADER_LENGTH])
      if buffer_length == 0:
        response = ''.join(self.message)
        self.done = True
        self.finished.callback(response)
        break
      if len(self.recvd) < buffer_length + ipc.BUFFER_HEADER_LENGTH:
        break
      buffer = self.recvd[ipc.BUFFER_HEADER_LENGTH:buffer_length + ipc.BUFFER_HEADER_LENGTH]
      self.recvd = self.recvd[buffer_length + ipc.BUFFER_HEADER_LENGTH:]
      self.message.append(buffer)

  def connectionLost(self, reason):
    if not self.done:
      self.finished.errback(ipc.ConnectionClosedException("Reader read 0 bytes."))

class TwistedHTTPTransceiver(object):
  """This transceiver uses the Agent class present in Twisted.web >= 9.0
     for issuing requests to the remote endpoint."""
  def __init__(self, host, port, remote_name=None, reactor=None):
    self.url = "http://%s:%d/" % (host, port)

    if remote_name is None:
      # There's no easy way to get this peer's remote address
      # in Twisted so I use a random UUID to identify ourselves
      import uuid
      self.remote_name = uuid.uuid4()

    if reactor is None:
      from twisted.internet import reactor
    self.agent = Agent(reactor)

  def read_framed_message(self, response):
    finished = Deferred()
    response.deliverBody(AvroProtocol(finished))
    return finished

  def transceive(self, request):
    req_method = 'POST'
    req_headers = {
      'Content-Type': ['avro/binary'],
      'Accept-Encoding': ['identity'],
    }

    body_producer = RequestStreamingProducer(request)
    d = self.agent.request(
      req_method,
      self.url,
      headers=Headers(req_headers),
      bodyProducer=body_producer)
    return d.addCallback(self.read_framed_message)

class AvroResponderResource(resource.Resource):
  """This Twisted.web resource can be placed anywhere in a URL hierarchy
     to provide an Avro endpoint. Different Avro protocols can be served
     by the same web server as long as they are in different resources in
     a URL hierarchy."""
  isLeaf = True

  def __init__(self, responder):
    resource.Resource.__init__(self)
    self.responder = responder

  def cb_render_POST(self, resp_body, request):
    request.setResponseCode(200)
    request.setHeader('Content-Type', 'avro/binary')
    resp_writer = ipc.FramedWriter(request)
    resp_writer.write_framed_message(resp_body)
    request.finish()

  def render_POST(self, request):
    # Unfortunately, Twisted.web doesn't support incoming
    # streamed input yet, the whole payload must be kept in-memory
    request.content.seek(0, 0)
    call_request_reader = ipc.FramedReader(request.content)
    call_request = call_request_reader.read_framed_message()
    d = maybeDeferred(self.responder.respond, call_request)
    d.addCallback(self.cb_render_POST, request)
    return server.NOT_DONE_YET