Refactor context propagation to work with async

context/opencensus-context/CHANGELOG.md

@@ -0,0 +1,5 @@
+# Changelog
+
+## Unreleased
+
+- Add this changelog.

context/opencensus-context/README.rst

@@ -0,0 +1,7 @@
+OpenCensus Runtime Context
+============================================================================
+
+|pypi|
+
+.. |pypi| image:: https://badge.fury.io/py/opencensus-context.svg
+   :target: https://pypi.org/project/opencensus-context/

context/opencensus-context/examples/async_span.py

@@ -0,0 +1,55 @@
+# Copyright 2019, OpenCensus Authors
+#
+# 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.
+
+import asyncio
+from opencensus.common.runtime_context import RuntimeContext
+
+RuntimeContext.register_slot('current_span', None)
+
+
+class Span(object):
+    def __init__(self, name):
+        self.name = name
+        self.parent = RuntimeContext.current_span
+
+    def __repr__(self):
+        return ('{}(name={}, parent={})'
+                .format(
+                    type(self).__name__,
+                    self.name,
+                    self.parent,
+                ))
+
+    async def __aenter__(self):
+        RuntimeContext.current_span = self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        RuntimeContext.current_span = self.parent
+
+
+async def main():
+    print(RuntimeContext)
+    async with Span('foo'):
+        print(RuntimeContext)
+        await asyncio.sleep(0.1)
+        async with Span('bar'):
+            print(RuntimeContext)
+            await asyncio.sleep(0.1)
+        print(RuntimeContext)
+        await asyncio.sleep(0.1)
+    print(RuntimeContext)
+
+
+if __name__ == '__main__':
+    asyncio.run(main())

context/opencensus-context/examples/explicit_threading.py

@@ -0,0 +1,48 @@
+# Copyright 2019, OpenCensus Authors
+#
+# 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.
+
+from threading import Thread
+from opencensus.common.runtime_context import RuntimeContext
+
+RuntimeContext.register_slot('operation_id', '<empty>')
+
+
+def work(name):
+    print('Entering worker:', RuntimeContext)
+    RuntimeContext.operation_id = name
+    print('Exiting worker:', RuntimeContext)
+
+
+if __name__ == '__main__':
+    print('Main thread:', RuntimeContext)
+    RuntimeContext.operation_id = 'main'
+
+    print('Main thread:', RuntimeContext)
+
+    # by default context is not propagated to worker thread
+    thread = Thread(target=work, args=('foo',))
+    thread.start()
+    thread.join()
+
+    print('Main thread:', RuntimeContext)
+
+    # user can propagate context explicitly
+    thread = Thread(
+        target=RuntimeContext.with_current_context(work),
+        args=('bar',),
+    )
+    thread.start()
+    thread.join()
+
+    print('Main thread:', RuntimeContext)

context/opencensus-context/examples/py27.py

@@ -0,0 +1,34 @@
+# Copyright 2019, OpenCensus Authors
+#
+# 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.
+
+from opencensus.common.runtime_context import RuntimeContext
+
+RuntimeContext.register_slot('correlation_context', lambda: {})
+
+
+def hello(name):
+    correlation_context = RuntimeContext.correlation_context.copy()
+    correlation_context['name'] = name
+    RuntimeContext.correlation_context = correlation_context
+
+    print(RuntimeContext)
+
+
+if __name__ == '__main__':
+    print(RuntimeContext)
+    RuntimeContext.correlation_context['test'] = True
+    print(RuntimeContext)
+    hello('hello')
+    RuntimeContext.clear()
+    print(RuntimeContext)

context/opencensus-context/examples/py37.py

@@ -0,0 +1,51 @@
+# Copyright 2019, OpenCensus Authors
+#
+# 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.
+
+import asyncio
+
+from opencensus.common.runtime_context import RuntimeContext
+
+RuntimeContext.register_slot('correlation_context', lambda: dict())
+
+
+async def hello(name):
+    correlation_context = RuntimeContext.correlation_context.copy()
+    correlation_context['name'] = name
+    RuntimeContext.correlation_context = correlation_context
+
+    for i in range(3):
+        print('Hello {} {} {}'.format(
+            name,
+            i,
+            RuntimeContext,
+        ))
+        await asyncio.sleep(0.1)
+
+
+async def main():
+    print(RuntimeContext)
+    RuntimeContext.correlation_context['test'] = True
+    print(RuntimeContext)
+    await asyncio.gather(
+        hello('foo'),
+        hello('bar'),
+        hello('baz'),
+    )
+    print(RuntimeContext)
+    RuntimeContext.clear()
+    print(RuntimeContext)
+
+
+if __name__ == '__main__':
+    asyncio.run(main())

context/opencensus-context/examples/span.py

@@ -0,0 +1,56 @@
+# Copyright 2019, OpenCensus Authors
+#
+# 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.
+
+from opencensus.common.runtime_context import RuntimeContext
+
+RuntimeContext.register_slot('current_span', None)
+
+
+class Span(object):
+    def __init__(self, name):
+        self.name = name
+        self.parent = RuntimeContext.current_span
+
+    def __repr__(self):
+        return ('{}({})'.format(type(self).__name__, self.name))
+
+    def __enter__(self):
+        RuntimeContext.current_span = self
+
+    def __exit__(self, type, value, traceback):
+        RuntimeContext.current_span = self.parent
+
+    def start(self):
+        RuntimeContext.current_span = self
+
+    def end(self):
+        RuntimeContext.current_span = self.parent
+
+
+if __name__ == '__main__':
+    print(RuntimeContext)
+    with Span('foo'):
+        print(RuntimeContext)
+        with Span('bar'):
+            print(RuntimeContext)
+        print(RuntimeContext)
+    print(RuntimeContext)
+
+    # explicit start/end span
+    span = Span('baz')
+    print(RuntimeContext)
+    span.start()
+    print(RuntimeContext)
+    span.end()
+    print(RuntimeContext)

context/opencensus-context/examples/thread_pool.py

@@ -0,0 +1,49 @@
+# Copyright 2019, OpenCensus Authors
+#
+# 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.
+
+from multiprocessing.dummy import Pool as ThreadPool
+import time
+import threading
+from opencensus.common.runtime_context import RuntimeContext
+
+RuntimeContext.register_slot('operation_id', '<empty>')
+_console_lock = threading.Lock()
+
+
+def println(msg):
+    with _console_lock:
+        print(msg)
+
+
+def work(name):
+    println('Entering worker[{}]: {}'.format(name, RuntimeContext))
+    RuntimeContext.operation_id = name
+    time.sleep(0.01)
+    println('Exiting worker[{}]: {}'.format(name, RuntimeContext))
+
+
+if __name__ == "__main__":
+    println('Main thread: {}'.format(RuntimeContext))
+    RuntimeContext.operation_id = 'main'
+    pool = ThreadPool(2)  # create a thread pool with 2 threads
+    pool.map(RuntimeContext.with_current_context(work), [
+        'bear',
+        'cat',
+        'dog',
+        'horse',
+        'rabbit',
+    ])
+    pool.close()
+    pool.join()
+    println('Main thread: {}'.format(RuntimeContext))

context/opencensus-context/opencensus/__init__.py

@@ -0,0 +1 @@
+__path__ = __import__('pkgutil').extend_path(__path__, __name__)

context/opencensus-context/opencensus/common/__init__.py

@@ -0,0 +1 @@
+__path__ = __import__('pkgutil').extend_path(__path__, __name__)