add event hook into tchannel

add zipkin tracing

Author: Junchao Wu

examples/handler.py

@@ -30,7 +30,7 @@
 
 
 @tornado.gen.coroutine
-def say_hi(request, response, opts):
+def say_hi(request, response, proxy):
     arg2 = yield request.arg2()
     arg3 = yield request.arg3()
     response.argstreams = [
@@ -41,7 +41,7 @@ def say_hi(request, response, opts):
 
 
 @tornado.gen.coroutine
-def say_ok(request, response, opts):
+def say_ok(request, response, proxy):
     yield print_arg(request, 1)
     yield print_arg(request, 2)
 
@@ -52,9 +52,10 @@ def say_ok(request, response, opts):
 
 
 @tornado.gen.coroutine
-def echo(request, response, opts):
+def echo(request, response, proxy):
+    print "echo"
+    yield tornado.gen.sleep(1)
     # stream args right back to request side
-    print "streaming"
     response.argstreams = [
         InMemStream(request.endpoint),
         request.argstreams[1],
@@ -63,7 +64,7 @@ def echo(request, response, opts):
 
 
 @tornado.gen.coroutine
-def slow(request, response, opts):
+def slow(request, response, proxy):
     yield tornado.gen.sleep(random.random())
     response.argstreams = [
         InMemStream(),
@@ -80,7 +81,7 @@ def get_example_handler():
     dispatcher.register("slow", slow)
 
     @dispatcher.route("bye")
-    def say_bye(request, response, opts):
+    def say_bye(request, response, proxy):
         yield print_arg(request, 1)
         yield print_arg(request, 2)
 

examples/stream_client.py

@@ -20,9 +20,10 @@
 
 import os
 import sys
-from options import get_args
 import tornado
 import tornado.ioloop
+
+from options import get_args
 from tchannel.tornado import TChannel
 from tchannel.tornado.stream import InMemStream, PipeStream
 from tchannel.tornado.util import print_arg

examples/thrift/client.py examples/thrift_examples/client.py

@@ -0,0 +1,109 @@
+# Copyright (c) 2015 Uber Technologies, Inc.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+
+from __future__ import absolute_import
+from enum import IntEnum
+import logging
+
+log = logging.getLogger('tchannel')
+
+
+class EventType(IntEnum):
+    """Types to represent system events
+
+    Events:
+
+    send_request: before client sends request
+
+    send_response: after server sends response
+
+    receive_request: after server receive request
+
+    receive_response: after client receive response
+
+    """
+    send_request = 0x00,
+    send_response = 0x01,
+    receive_request = 0x02,
+    receive_response = 0x03
+
+
+class EventHook(object):
+    """provide all event hook interfaces
+
+    Customized Hook should should inherit from EventHook class and implement
+    the events' hooks that it wants to listen.
+
+    Example::
+
+        TraceHook(EventHook):
+            def send_request(self, context):
+                ....
+
+    """
+
+    def send_request(self, context):
+        """Event hook for sending request
+
+        :param context:
+            request object to send
+        """
+        pass
+
+    def send_response(self, context):
+        """Event hook for sending response
+
+        :param context:
+            response object sent
+        """
+        pass
+
+    def receive_request(self, context):
+        """Event hook for receiving request
+
+        :param context:
+            request object received
+        """
+        pass
+
+    def receive_response(self, context):
+        """Event hook for receiving response
+
+        :param context:
+            response object received
+        """
+        pass
+
+
+class EventEmitter(object):
+    def __init__(self):
+        self.hooks = []
+
+    def register_hook(self, hook):
+        self.hooks.append(hook)
+
+    def fire(self, event, *args, **kwargs):
+        # TODO find proper hook name
+        event_hook_name = event.name
+        for hook in self.hooks:
+            try:
+                getattr(hook, event_hook_name)(*args, **kwargs)
+            except Exception as e:
+                log.error(e.message)

examples/thrift/hello.thrift examples/thrift_examples/hello.thrift

@@ -24,6 +24,7 @@
 import os
 import socket
 import sys
+from tchannel.event import EventType
 
 import tornado.gen
 import tornado.iostream
@@ -75,13 +76,24 @@ class TornadoConnection(object):
     CALL_REQ_TYPES = frozenset([Types.CALL_REQ, Types.CALL_REQ_CONTINUE])
     CALL_RES_TYPES = frozenset([Types.CALL_RES, Types.CALL_RES_CONTINUE])
 
-    def __init__(self, connection):
+    def __init__(self, connection, tchannel=None):
         assert connection, "connection is required"
 
         self.closed = False
         self.connection = connection
 
-        self.remote_host = None
+        sockname = connection.socket.getsockname()
+        if len(sockname) == 2:
+            (self.remote_host,
+             self.remote_host_port) = sockname
+        elif len(sockname) == 1:
+            self.remote_host = sockname[0]
+            self.remote_host_port = 0
+        else:
+            self.remote_host = "0.0.0.0"
+            self.remote_host_port = 0
+
+        self.remote_host_port = int(self.remote_host_port)
         self.remote_process_name = None
         self.requested_version = PROTOCOL_VERSION
 
@@ -90,8 +102,10 @@ def __init__(self, connection):
 
         # We need to use two separate message factories to avoid message ID
         # collision while assembling fragmented messages.
-        self.request_message_factory = MessageFactory()
-        self.response_message_factory = MessageFactory()
+        self.request_message_factory = MessageFactory(self.remote_host,
+                                                      self.remote_host_port)
+        self.response_message_factory = MessageFactory(self.remote_host,
+                                                       self.remote_host_port)
 
         # Queue of unprocessed incoming calls.
         self._messages = queues.Queue()
@@ -103,6 +117,8 @@ def __init__(self, connection):
         # handshake has been performed.
         self._loop_running = False
 
+        self.tchannel = tchannel
+
         connection.set_close_callback(self._on_close)
 
     def next_message_id(self):
@@ -344,14 +360,16 @@ def _extract_handshake_headers(self, message):
                 'Missing required header: process_name'
             )
 
-        self.remote_host = message.host_port
+        (self.remote_host,
+            self.remote_host_port) = message.host_port.rsplit(':', 1)
+        self.remote_host_port = int(self.remote_host_port)
         self.remote_process_name = message.process_name
         self.requested_version = message.version
 
     @classmethod
     @tornado.gen.coroutine
     def outgoing(cls, hostport, process_name=None, serve_hostport=None,
-                 handler=None):
+                 handler=None, tchannel=None):
         """Initiate a new connection to the given host.
 
         :param hostport:
@@ -384,7 +402,7 @@ def outgoing(cls, hostport, process_name=None, serve_hostport=None,
                 "Couldn't connect to %s" % hostport, e
             )
 
-        connection = cls(stream)
+        connection = cls(stream, tchannel)
         log.debug("Performing handshake with %s", hostport)
         yield connection.initiate_handshake(headers={
             'host_port': serve_hostport,
@@ -521,12 +539,22 @@ def post_response(self, response):
         try:
             response.close_argstreams()
             yield self._stream(response, self.response_message_factory)
+
+            # event: send_response
+            if self.tchannel:
+                self.tchannel.event_emitter.fire(
+                    EventType.send_response, response)
         finally:
             response.close_argstreams(force=True)
 
     @tornado.gen.coroutine
     def post_request(self, request):
         """send the given request and response is not required"""
+
+        # event: send_request
+        if self.tchannel:
+            self.tchannel.event_emitter.fire(EventType.send_request, request)
+
         try:
             request.close_argstreams()
             yield self._stream(request, self.request_message_factory)
@@ -554,4 +582,21 @@ def send_request(self, request):
         future = tornado.gen.Future()
         self._outstanding[request.id] = future
         self.post_request(request)
-        return future
+
+        # the actual future that caller will yield
+        res_future = tornado.gen.Future()
+
+        def adapt_tracing(f):
+            # fetch the request tracing for response
+            f.result().tracing = request.tracing
+            res_future.set_result(f.result())
+            # event: receive_response
+            if self.tchannel:
+                self.tchannel.event_emitter.fire(
+                    EventType.receive_response, f.result())
+
+        tornado.ioloop.IOLoop.current().add_future(
+            future,
+            adapt_tracing)
+
+        return res_future