broker

#!/usr/bin/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
#
#   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.
#
import optparse, socket
from connection import Connection
from session import Session, SessionError
from link import LinkError, Receiver, Sender, link
from selector import Acceptor, Selector
from queue import Queue, DESTRUCTIVE
from util import ConnectionSelectable

class Broker:

  def __init__(self):
    self.queues = {}
    self.cursors = {}
    self.init = {Sender.direction: self.init_sender,
                 Receiver.direction: self.init_receiver}
    self.process = {Sender.direction:self.process_sender,
                    Receiver.direction: self.process_receiver}

  def tick(self, connection):
    if connection.opening():
      connection.open()

    # XXX
    for ssn in connection.incoming.values():
      if ssn.beginning():
        ssn.begin()

    for ssn in connection.outgoing.values():
      for link in ssn.links.values():
        if link.opening():
          if self.init[link.direction](link):
            link.local = link.remote
            link.attach()
          else:
            link.local = None
            # XXX
            link.modified = True
            link.attach()
            link.detach()

        # XXX
        if link.local:
          self.process[link.direction](link)

        if link.closing():
          link.detach()

        if link.closed():
          ssn.remove(link)

      if ssn.ending():
        ssn.end()
        connection.remove(ssn)

    if connection.closing():
      connection.close()

    connection.tick()

  def init_sender(self, link):
    if link.remote.source in self.queues:
      q = self.queues[link.remote.source]
      cursor = q.cursor(DESTRUCTIVE)
      self.cursors[link] = cursor, {}
      return True
    else:
      return False

  def init_receiver(self, link):
    if link.remote.target in self.queues:
      link.flow(10)
      return True
    else:
      return False

  def process_sender(self, link):
    cursor, unacked = self.cursors[link]
    while link.capacity() > 0:
      entry = cursor.get()
      if entry is None:
        link.drained()
        break
      elif entry.item is None:
        continue
      else:
        xfr = entry.item
        tag = link.send(fragments = xfr.fragments)
        print "TRANSFER:", tag, xfr
        unacked[tag] = entry

    for t, r in link.get_remote():
      if r.settled or r.state is not None:
        e = unacked[t]
        xfr = e.item
        e.dequeue()
        print "DEQUEUED:", t, r
        del unacked[t]
        link.settle(t, "DEQUEUED")

  def process_receiver(self, link):
    while link.pending():
      xfr = link.get()
      self.queues[link.local.target].put(xfr)
      print "ENQUEUED:", xfr
      link.settle(xfr.delivery_tag, "ENQUEUED")
    if link.capacity() < 10: link.flow(10)

def bind(host, port):
  sock = socket.socket()
  sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  sock.bind((host, port))
  sock.listen(5)
  sock.setblocking(0)
  return sock

parser = optparse.OptionParser(usage="usage: %prog [options] QUEUE_1 ... QUEUE_n",
                               description="Prototype amqp broker.")
parser.add_option("-i", "--interface", default="0.0.0.0",
                  help="interface to listen on (default %default)")
parser.add_option("-p", "--port", type=int, default=5672,
                  help="port to listen on (default %default)")
parser.add_option("-t", "--trace", default="err",
                  help="enable tracing for specified categories")

opts, args = parser.parse_args()

try:
  broker = Broker()

  for a in args:
    broker.queues[a] = Queue()

  selector = Selector()
  sock = bind(opts.interface, opts.port)

  def handler(accepted):
    conn = Connection(lambda n: Session(n, link))
    conn.tracing = set(opts.trace.split())
    selector.register(ConnectionSelectable(accepted, conn, broker.tick))

  selector.register(Acceptor(sock, handler))
  selector.run()
except KeyboardInterrupt:
  pass