fix(v05): reset string table index map after rollbacks [backport 1.20] (#7879)

Backport 3945eb6 from #7874 to 1.20.

Resolves: #7662

This PR resolves an encoding issue when `DD_TRACE_API_VERSION` is set to
`v0.5`

## Background

The python tracer always attempts to encode a trace. While encoding a
trace if the payload/buffer size limit is reached the trace buffer and
string table are rolled back to a previous state. However the rollback
only resets the index in the string table and the trace buffer. It does
not removing stale string->index pairs from `StringTable._table`. This
causes problems when:
- buffer limit is reached while encoding a trace -> then trace buffer
and string table are rolled back -> then a smaller trace containing
strings that map to stale/rolled back string table index is successfully
encoded.

This scenario is very rare and only occurs on a very small fraction of
python traces.

The scenario described above can mis-encode traces in two different ways
resulting two types of decoding errors:
- out of bounds access: The encoded trace references string table
indices that have not been serialized. In python decoders this will
raise an IndexError.
- string swapping/overwriting: Two different strings map to the same
string table index. When decoding the trace both strings will be mapped
to the same value. This is the behavior we saw in #incident-22455

## Fix 

Remove stale _id->string pairs from the `StringTable._table` after
rolling back

## Checklist

- [x] Change(s) are motivated and described in the PR description.
- [x] Testing strategy is described if automated tests are not included
in the PR.
- [x] Risk is outlined (performance impact, potential for breakage,
maintainability, etc).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] [Library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
are followed. If no release note is required, add label
`changelog/no-changelog`.
- [x] Documentation is included (in-code, generated user docs, [public
corp docs](https://github.com/DataDog/documentation/)).
- [x] Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist

- [x] Title is accurate.
- [x] No unnecessary changes are introduced.
- [x] Description motivates each change.
- [x] Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes unless absolutely necessary.
- [x] Testing strategy adequately addresses listed risk(s).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] Release note makes sense to a user of the library.
- [x] Reviewer has explicitly acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment.
- [x] Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)
- [x] If this PR touches code that signs or publishes builds or
packages, or handles credentials of any kind, I've requested a review
from `@DataDog/security-design-and-guidance`.
- [x] This PR doesn't touch any of that.

---------

Co-authored-by: Munir Abdinur <[email protected]>

Author: github-actions[bot]

ddtrace/internal/_encoding.pyx

@@ -286,13 +286,26 @@ cdef class MsgpackStringTable(StringTable):
             self.pk.length = self._sp_len
             self._next_id = self._sp_id
 
+        # After rolling back the string table next_id we must remove all stale string -> _id pairs
+        # This will resolve two classes of encoding errors:
+        #  - multiple strings referencing the same string table index. In this scenario
+        #    two different strings in the encoded trace could be serialized with the same _id. In this scenario
+        #    two different strings could reference one string in the encoded trace (string swapping).
+        # - when the string table references an index of the string table that is not serialized. The encoded
+        #    trace can not be decoded without accessing an invalid index. In this scenario the agent will
+        #    return a 400 status code.
+        self._table = {s: idx for s, idx in self._table.items() if idx < self._next_id}
+
     cdef get_bytes(self):
         cdef int ret
-        cdef stdint.uint32_t table_size = self._next_id
-        cdef int offset = MSGPACK_STRING_TABLE_LENGTH_PREFIX_SIZE - array_prefix_size(table_size)
-        cdef int old_pos = self.pk.length
-
+        cdef stdint.uint32_t table_size
+        cdef int offset
+        cdef int old_pos
         with self._lock:
+            table_size = self._next_id
+            offset = MSGPACK_STRING_TABLE_LENGTH_PREFIX_SIZE - array_prefix_size(table_size)
+            old_pos = self.pk.length
+
             # Update table size prefix
             self.pk.length = offset
             ret = msgpack_pack_array(&self.pk, table_size)
@@ -305,7 +318,7 @@ cdef class MsgpackStringTable(StringTable):
                 return None
             self.pk.length = old_pos
 
-        return PyBytes_FromStringAndSize(self.pk.buf + offset, self.pk.length - offset)
+            return PyBytes_FromStringAndSize(self.pk.buf + offset, self.pk.length - offset)
 
     @property
     def size(self):

releasenotes/notes/trace-api-v05-encoding-reset-string-table-17ee19c026b22af9.yaml

@@ -0,0 +1,6 @@
+---
+fixes:
+  - |
+    tracing: Resolves trace encoding errors raised when ``DD_TRACE_API_VERSION`` is set to ``v0.5`` and a BufferFull
+    Exception is raised by the TraceWriter. This fix ensures span fields are not overwritten and reduces the frequency
+    of 4XX errors in the trace agent.

tests/tracer/test_encoders.py

@@ -302,6 +302,36 @@ def allencodings(f):
     return pytest.mark.parametrize("encoding", MSGPACK_ENCODERS.keys())(f)
 
 
+def test_msgpack_encoding_after_an_exception_was_raised():
+    """Ensure that the encoder's state is consistent after an Exception is raised during encoding"""
+    # Encode a trace after a rollback/BufferFull occurs exception
+    rolledback_encoder = MsgpackEncoderV05(1 << 12, 1 << 12)
+    trace = gen_trace(nspans=1, ntags=100, nmetrics=100, key_size=10, value_size=10)
+    rand_string = rands(size=20, chars=string.ascii_letters)
+    # trace only has one span
+    trace[0].set_tag_str("some_tag", rand_string)
+    try:
+        # Encode a trace that will trigger a rollback/BufferItemTooLarge exception
+        # BufferFull is not raised since only one span is being encoded
+        rolledback_encoder.put(trace)
+    except BufferItemTooLarge:
+        pass
+    else:
+        pytest.fail("Encoding the trace did not overflow the trace buffer. We should increase the size of the span.")
+    # Successfully encode a small trace
+    small_trace = gen_trace(nspans=1, ntags=0, nmetrics=0)
+    # Add a tag to the small trace that was previously encoded in the encoder's StringTable
+    small_trace[0].set_tag_str("previously_encoded_string", rand_string)
+    rolledback_encoder.put(small_trace)
+
+    # Encode a trace without triggering a rollback/BufferFull exception
+    ref_encoder = MsgpackEncoderV05(1 << 20, 1 << 20)
+    ref_encoder.put(small_trace)
+
+    # Ensure the two encoders have the same state
+    assert rolledback_encoder.encode() == ref_encoder.encode()
+
+
 @allencodings
 def test_custom_msgpack_encode(encoding):
     encoder = MSGPACK_ENCODERS[encoding](1 << 20, 1 << 20)