From f1612f28954e38cd470260795bb57c544bf269d6 Mon Sep 17 00:00:00 2001
From: Mark McLoughlin <markmc@redhat.com>
Date: Fri, 9 Aug 2013 11:15:00 +0100
Subject: [PATCH] Add thread-local store of request context

Oslo's logging code has some useful support for including bits of the
request context in log messages. While this isn't exclusively about the
request context in a dispatching RPC method, it seems useful for
oslo.messaging to support the concept for at least this use case simply
by recording the context in a thread local store before dispatching an
endpoint method and immediately clearing it when the method returns.

Note, we don't need to store weak refs in our store because we will
clear the reference in all cases rather than ever leaving a stale
reference around in the store.

Change-Id: I70ac06ed3a2a891a7a7b388b1823a0f3b08f2dd1
---
 doc/source/server.rst                 |  2 +
 oslo/messaging/__init__.py            |  1 +
 oslo/messaging/_drivers/amqpdriver.py |  2 -
 oslo/messaging/localcontext.py        | 55 +++++++++++++++++++++++++++
 oslo/messaging/rpc/dispatcher.py      |  7 +++-
 5 files changed, 64 insertions(+), 3 deletions(-)
 create mode 100644 oslo/messaging/localcontext.py

diff --git a/doc/source/server.rst b/doc/source/server.rst
index 933829f3a..b7ac48650 100644
--- a/doc/source/server.rst
+++ b/doc/source/server.rst
@@ -16,3 +16,5 @@ Server
 .. autofunction:: expected_exceptions
 
 .. autoexception:: ExpectedException
+
+.. autofunction:: get_local_context
diff --git a/oslo/messaging/__init__.py b/oslo/messaging/__init__.py
index 83529c5b4..453a73ea2 100644
--- a/oslo/messaging/__init__.py
+++ b/oslo/messaging/__init__.py
@@ -14,6 +14,7 @@
 #    under the License.
 
 from .exceptions import *
+from .localcontext import *
 from .notify import *
 from .rpc import *
 from .serializer import *
diff --git a/oslo/messaging/_drivers/amqpdriver.py b/oslo/messaging/_drivers/amqpdriver.py
index 40c1fcc82..599b117ea 100644
--- a/oslo/messaging/_drivers/amqpdriver.py
+++ b/oslo/messaging/_drivers/amqpdriver.py
@@ -72,8 +72,6 @@ class AMQPListener(base.Listener):
         self.incoming = []
 
     def __call__(self, message):
-        # FIXME(markmc): del local.store.context
-
         # FIXME(markmc): logging isn't driver specific
         rpc_common._safe_log(LOG.debug, 'received %s', message)
 
diff --git a/oslo/messaging/localcontext.py b/oslo/messaging/localcontext.py
new file mode 100644
index 000000000..f7da49f63
--- /dev/null
+++ b/oslo/messaging/localcontext.py
@@ -0,0 +1,55 @@
+
+# Copyright 2013 Red Hat, Inc.
+#
+#    Licensed 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.
+
+__all__ = [
+    'get_local_context',
+    'set_local_context',
+    'clear_local_context',
+]
+
+import threading
+import uuid
+
+_KEY = '_%s_%s' % (__name__.replace('.', '_'), uuid.uuid4().hex)
+_STORE = threading.local()
+
+
+def get_local_context(ctxt):
+    """Retrieve the RPC endpoint request context for the current thread.
+
+    This method allows any code running in the context of a dispatched RPC
+    endpoint method to retrieve the context for this request.
+
+    This is commonly used for logging so that, for example, you can include the
+    request ID, user and tenant in every message logged from a RPC endpoint
+    method.
+
+    :returns: the context for the retuest dispatched in the current thread
+    """
+    return getattr(_STORE, _KEY, None)
+
+
+def set_local_context(ctxt):
+    """Set the request context for the current thread.
+
+    :param ctxt: a deserialized request context
+    :type ctxt: dict
+    """
+    setattr(_STORE, _KEY, ctxt)
+
+
+def clear_local_context():
+    """Clear the request context for the current thread."""
+    delattr(_STORE, _KEY)
diff --git a/oslo/messaging/rpc/dispatcher.py b/oslo/messaging/rpc/dispatcher.py
index bcf0864a8..5dea41ca2 100644
--- a/oslo/messaging/rpc/dispatcher.py
+++ b/oslo/messaging/rpc/dispatcher.py
@@ -26,6 +26,7 @@ __all__ = [
 import logging
 
 from oslo.messaging import _utils as utils
+from oslo.messaging import localcontext
 from oslo.messaging import serializer as msg_serializer
 from oslo.messaging import server as msg_server
 from oslo.messaging import target
@@ -117,7 +118,11 @@ class RPCDispatcher(object):
                 continue
 
             if hasattr(endpoint, method):
-                return self._dispatch(endpoint, method, ctxt, args)
+                localcontext.set_local_context(ctxt)
+                try:
+                    return self._dispatch(endpoint, method, ctxt, args)
+                finally:
+                    localcontext.clear_local_context()
 
             found_compatible = True