from twisted.internet.protocol import Protocol, ReconnectingClientFactory
import collector_pb2 as wire
import varint
from sys import stdout

class CollectorProtocol(Protocol):
    # CollectorProtocol is a Twisted implementation of appdash's protobuf-based
    # collector protocol.

    # writeMsg writes the protobuf message out to the protocol's transport. It
    # is preceded with the varint-encoded length of the data (i.e. a delimited
    # protobuf message).
    def writeMsg(self, msg):
        data = msg.SerializeToString()
        self.transport.write(varint.encode(len(data)))
        self.transport.write(data)

    def connectionMade(self):
        self._factory._log('connected!')

    def connectionLost(self, reason):
        self._factory._log("disconnected.", reason.getErrorMessage())

    def dataReceived(self, data):
        self._factory._log('got', len(data), 'bytes of unexpected data from server.')


class RemoteCollectorFactory(ReconnectingClientFactory):
    # RemoteCollectorFactory is a Twisted factory for remote collectors, which
    # collect spans and their annotations, sending them to a remote Go appdash
    # server for collection. After collection they can be viewed in appdash's
    # web user interface.

    _reactor = None
    _debug = False
    _remote = None
    _pending = []

    def __init__(self, reactor, debug=False):
        self._reactor = reactor
        self._debug = debug

    def _log(self, *args):
        if self._debug:
            print "appdash: %s" % (" ".join(args))

    # collect collects a packet to be flushed at a later time, when connection
    # has been made.
    #
    # TODO(slimsag): We shouldn't blindly take a protobuf CollectionPacket. It
    # is too low-level for clients.
    def collect(self, packet):
        self._pending.append(packet)
        if self._remote != None:
            self._reactor.callLater(1/2, self.__flush)

    # __flush is called internally after either a new collection has occured, or
    # after connection has been made with the remote server. It writes all the
    # pending messages out to the remote.
    def __flush(self):
        self._log("flushing", str(len(self._pending)), "messages")
        for p in self._pending:
            self._remote.writeMsg(p)
        self._pending = []

    def startedConnecting(self, connector):
        self._log('connecting..')

    def buildProtocol(self, addr):
        # Reset delay to reconnection -- otherwise it's exponential (which is
        # not a good match for us).
        self.resetDelay()

        # Create the protocol.
        p = CollectorProtocol()
        p._factory = self
        self._remote = p
        self._reactor.callLater(1/2, self.__flush)
        return p

    def clientConnectionFailed(self, connector, reason):
        self._log('connection failed:', reason.getErrorMessage())
        ReconnectingClientFactory.clientConnectionFailed(self, connector, reason)


@@ -1556,18 +1556,73 @@
 ts a
- packet to
+nnotations for the given spanID.%0A    #%0A    # The annotations will
  be 
@@ -1629,16 +1629,20 @@
 flushed 
+out 
 at a lat
@@ -1655,16 +1655,18 @@
 e, when 
+a 
 connecti
@@ -1678,174 +1678,660 @@
   # 
-has been made.%0A    #%0A    # TODO(slimsag): We shouldn't blindly take a protobuf 
+to the remote server has been made.%0A    def collect(self, spanID, *annotations):%0A        self._log(%22collecting%22, str(len(annotations)), %22annotations for%22, str(spanID))%0A%0A        # Create the protobuf message.%0A        p = wire.
 Collect
-ion
 Packet
-. It%0A    # is too low-level for clients.%0A    def collect(self, packet):
+()%0A        p.spanid.trace = spanID.trace%0A        p.spanid.span = spanID.span%0A        p.spanid.parent = spanID.parent%0A%0A        # Add each annotation to the message.%0A        for a in annotations:%0A            ap = p.annotation.add()%0A            ap.key = a.key%0A            ap.value = a.value%0A%0A        # Append the collection packet to the pending queue and flush later if%0A        # we already have a remote connection.
 %0A   
@@ -2357,21 +2357,16 @@
 append(p
-acket
 )%0A      

Make RemoteCollectorFactory.collect match Go.