Tag Propagation (#209)

Author: kcooperstein

opencensus/tags/propagation/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2018, 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.

opencensus/tags/propagation/binary_serializer.py

@@ -0,0 +1,108 @@
+# Copyright 2018, 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.
+
+# -*- coding: utf-8 -*-
+
+import logging
+import six
+
+from google.protobuf.internal.encoder import _VarintBytes
+
+from opencensus.tags import tag_map as tag_map_module
+
+# Used for decoding hex bytes to hex string.
+UTF8 = 'utf-8'
+
+VERSION_ID = 0
+TAG_FIELD_ID = 0
+TAG_MAP_SERIALIZED_SIZE_LIMIT = 8192
+
+
+class BinarySerializer(object):
+    def from_byte_array(self, binary):
+        if len(binary) <= 0:
+            logging.warning("Input byte[] cannot be empty/")
+            return tag_map_module.TagMap()
+        else:
+            buffer = memoryview(binary)
+            version_id = buffer[0]
+            if six.PY2:
+                version_id = ord(version_id)
+            if version_id != VERSION_ID:
+                raise ValueError("Invalid version id.")
+            return self._parse_tags(buffer)
+
+    def to_byte_array(self, tag_context):
+        encoded_bytes = b''
+        encoded_bytes += _VarintBytes(VERSION_ID)
+        total_chars = 0
+        for tag in tag_context.tags:
+            for tag_key, tag_value in tag.items():
+                total_chars += len(tag_key)
+                total_chars += len(tag_value)
+                encoded_bytes = self._encode_tag(
+                    tag_key, tag_value, encoded_bytes)
+        if total_chars <= TAG_MAP_SERIALIZED_SIZE_LIMIT:
+            return encoded_bytes
+        else:  # pragma: NO COVER
+            logging.warning("Size of the tag context exceeds the maximum size")
+
+    def _parse_tags(self, buffer):
+        tag_context = tag_map_module.TagMap()
+        limit = len(buffer)
+        total_chars = 0
+        i = 1
+        while i < limit:
+            field_id = buffer[i] if six.PY3 else ord(buffer[i])
+            if field_id == TAG_FIELD_ID:
+                i += 1
+                key = self._decode_string(buffer, i)
+                i += len(key)
+                total_chars += len(key)
+                i += 1
+                val = self._decode_string(buffer, i)
+                i += len(val)
+                total_chars += len(val)
+                i += 1
+                if total_chars > \
+                        TAG_MAP_SERIALIZED_SIZE_LIMIT:  # pragma: NO COVER
+                    logging.warning("Size of the tag context exceeds maximum")
+                    break
+                else:
+                    tag_context.insert(str(key), str(val))
+            else:
+                break
+        return tag_context
+
+    def _encode_tag(self, tag_key, tag_value, encoded_bytes):
+        encoded_bytes += _VarintBytes(TAG_FIELD_ID)
+        encoded_bytes = self._encode_string(tag_key, encoded_bytes)
+        encoded_bytes = self._encode_string(tag_value, encoded_bytes)
+        return encoded_bytes
+
+    def _encode_string(self, input_str, encoded_bytes):
+        encoded_bytes += _VarintBytes(len(input_str))
+        encoded_bytes += input_str.encode(UTF8)
+        return encoded_bytes
+
+    def _decode_string(self, buffer, pos):
+        length = buffer[pos] if six.PY3 else ord(buffer[pos])
+        builder = ""
+        i = 1
+        while i <= length:
+            bytes_to_decode = buffer[pos + i] if six.PY3 \
+                else ord(buffer[pos + i])
+            builder += _VarintBytes(bytes_to_decode).decode()
+            i += 1
+        return builder

opencensus/tags/tag_map.py

@@ -29,7 +29,7 @@ def __init__(self, tags=None):
                     self._map[tag_key] = tag_value
 
         else:
-            self.tags = {}
+            self._map = {}
 
     @property
     def map(self):

tests/unit/tags/propagation/test_binary_serializer.py

@@ -0,0 +1,72 @@
+# Copyright 2018, 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 unittest
+
+from opencensus.tags.propagation import binary_serializer
+
+
+class TestBinarySerializer(unittest.TestCase):
+    def test_from_byte_array_input_empty(self):
+        binary = bytearray(b'')
+        propagator = binary_serializer.BinarySerializer()
+        tag_context = propagator.from_byte_array(binary=binary)
+        expected_tags = {}
+
+        self.assertEqual(expected_tags, tag_context.map)
+
+    def test_from_byte_array_invalid_version_id(self):
+        binary = bytearray(b'\x04key1\x04val1')
+        propagator = binary_serializer.BinarySerializer()
+        with self.assertRaises(ValueError):
+            propagator.from_byte_array(binary=binary)
+
+    def test_from_byte_array(self):
+        binary = bytearray(b'\x00\x00\x04key1\x04val1\x00\x04key2\x04val2'
+                           b'\x00\x04key3\x04val3')
+        propagator = binary_serializer.BinarySerializer()
+        tag_context = propagator.from_byte_array(binary=binary)
+        expected_tags = {'key1': 'val1', 'key2': 'val2', 'key3': 'val3'}
+
+        self.assertEqual(expected_tags, tag_context.map)
+
+    def test_to_byte_array(self):
+        from opencensus.tags.tag_map import TagMap
+        from collections import OrderedDict
+
+        tag_map = OrderedDict(
+            [('key1', 'val1'), ('key2', 'val2'),
+             ('key3', 'val3'), ('key4', 'val4')])
+        tags = [tag_map]
+        tag_context = TagMap(tags=tags)
+        propagator = binary_serializer.BinarySerializer()
+        binary = propagator.to_byte_array(tag_context)
+
+        expected_binary = b'\x00\x00\x04key1\x04val1\x00\x04key2\x04val2\x00' \
+                          b'\x04key3\x04val3\x00\x04key4\x04val4'
+
+        self.assertEqual(binary, expected_binary)
+
+    def test__parse_tags_invalid_field_id(self):
+        from collections import OrderedDict
+
+        propagator = binary_serializer.BinarySerializer()
+        binary = bytearray(b'\x00\x00\x04key1\x04val1\x04key2\x04val2'
+                           b'\x00\x04key3\x04val3')
+        buffer = memoryview(binary)
+        tag_context = propagator._parse_tags(buffer)
+        expected_dict = OrderedDict(
+            [('key1', 'val1')])
+
+        self.assertEqual(frozenset(tag_context.map), frozenset(expected_dict))

tests/unit/tags/test_tag_map.py

@@ -21,7 +21,6 @@ class TestTagMap(unittest.TestCase):
 
     def test_constructor_defaults(self):
         tag_map = tag_map_module.TagMap()
-        self.assertEqual(tag_map.tags, {})
         self.assertEqual(tag_map.map, {})
 
     def test_constructor_explicit(self):