fixes

Author: michael-zhao459

datadog_lambda/dsm.py

@@ -70,12 +70,12 @@ def _dsm_set_sns_context(event):
             if not sns_data:
                 return ""
             arn = sns_data.get("TopicArn", "")
-            # Trace data needed for sns payload size calculation
+            context_json = _get_dsm_context_from_lambda(record)
             trace_data = base64.b64encode(
-                json.dumps(_get_dsm_context_from_lambda(record)).encode("utf-8")
+                json.dumps(context_json).encode("utf-8")
             ).decode("utf-8")
             payload_size = calculate_sns_payload_size(record, trace_data)
-            _dsm_set_context_helper(record, "sns", arn, payload_size)
+            _dsm_set_context_helper("sns", arn, payload_size, context_json)
         except Exception as e:
             logger.error(format_err_with_traceback(e))