From e1dc017df98f703e04916ee89858a85e1ee35811 Mon Sep 17 00:00:00 2001 From: Peter Giacomo Lombardo Date: Tue, 25 Aug 2020 11:30:14 +0200 Subject: [PATCH] AWS Lambda: Assure Server-Timing response header --- instana/instrumentation/aws/lambda_inst.py | 12 +++ tests/platforms/test_lambda.py | 86 +++++++++++++++++----- 2 files changed, 78 insertions(+), 20 deletions(-) diff --git a/instana/instrumentation/aws/lambda_inst.py b/instana/instrumentation/aws/lambda_inst.py index 502413ab..772e7214 100644 --- a/instana/instrumentation/aws/lambda_inst.py +++ b/instana/instrumentation/aws/lambda_inst.py @@ -24,6 +24,18 @@ def lambda_handler_with_instana(wrapped, instance, args, kwargs): enrich_lambda_span(agent, scope.span, *args) try: result = wrapped(*args, **kwargs) + + if isinstance(result, dict): + server_timing_value = "intid;desc=%s" % scope.span.context.trace_id + if 'headers' in result: + result['headers']['Server-Timing'] = server_timing_value + elif 'multiValueHeaders' in result: + result['multiValueHeaders']['Server-Timing'] = [server_timing_value] + else: + # If both 'headers' and 'multiValueHeaders' aren't in result, + # then default to setting single value 'headers' + result['headers'] = dict() + result['headers']['Server-Timing'] = server_timing_value except Exception as exc: if scope.span: scope.span.log_exception(exc) diff --git a/tests/platforms/test_lambda.py b/tests/platforms/test_lambda.py index 4aa68ce2..5861c928 100644 --- a/tests/platforms/test_lambda.py +++ b/tests/platforms/test_lambda.py @@ -31,7 +31,11 @@ def __init__(self, **kwargs): # This is the target handler that will be instrumented for these tests def my_lambda_handler(event, context): # print("target_handler called") - return "All Ok" + return { + 'statusCode': 200, + 'headers': {'Content-Type': 'application/json'}, + 'body': json.dumps({'site': 'pwpush.com', 'response': 204}) + } # We only want to monkey patch the test handler once so do it here os.environ["LAMBDA_HANDLER"] = "tests.platforms.test_lambda.my_lambda_handler" @@ -118,7 +122,7 @@ def test_has_extra_http_headers(self): def test_has_options(self): self.create_agent_and_setup_tracer() self.assertTrue(hasattr(self.agent, 'options')) - self.assertTrue(type(self.agent.options) is AWSLambdaOptions) + self.assertTrue(isinstance(self.agent.options, AWSLambdaOptions)) assert(self.agent.options.endpoint_proxy == { }) def test_get_handler(self): @@ -173,14 +177,17 @@ def test_custom_service_name(self): result = lambda_handler(event, self.context) os.environ.pop('INSTANA_SERVICE_NAME') - self.assertEqual('All Ok', result) + assert isinstance(result, dict) + assert 'headers' in result + assert 'Server-Timing' in result['headers'] + payload = self.agent.collector.prepare_payload() self.assertTrue("metrics" in payload) self.assertTrue("spans" in payload) self.assertEqual(2, len(payload.keys())) - self.assertTrue(type(payload['metrics']['plugins']) is list) + self.assertTrue(isinstance(payload['metrics']['plugins'], list)) self.assertTrue(len(payload['metrics']['plugins']) == 1) plugin_data = payload['metrics']['plugins'][0] @@ -197,6 +204,9 @@ def test_custom_service_name(self): self.assertIsNotNone(span.ts) self.assertIsNotNone(span.d) + server_timing_value = "intid;desc=%s" % span.t + assert result['headers']['Server-Timing'] == server_timing_value + self.assertEqual({'hl': True, 'cp': 'aws', 'e': 'arn:aws:lambda:us-east-2:12345:function:TestPython:1'}, span.f) @@ -233,14 +243,17 @@ def test_api_gateway_trigger_tracing(self): # The original Lambda handler is set in os.environ["LAMBDA_HANDLER"] result = lambda_handler(event, self.context) - self.assertEqual('All Ok', result) + assert isinstance(result, dict) + assert 'headers' in result + assert 'Server-Timing' in result['headers'] + payload = self.agent.collector.prepare_payload() self.assertTrue("metrics" in payload) self.assertTrue("spans" in payload) self.assertEqual(2, len(payload.keys())) - self.assertTrue(type(payload['metrics']['plugins']) is list) + self.assertTrue(isinstance(payload['metrics']['plugins'], list)) self.assertTrue(len(payload['metrics']['plugins']) == 1) plugin_data = payload['metrics']['plugins'][0] @@ -257,6 +270,9 @@ def test_api_gateway_trigger_tracing(self): self.assertIsNotNone(span.ts) self.assertIsNotNone(span.d) + server_timing_value = "intid;desc=%s" % span.t + assert result['headers']['Server-Timing'] == server_timing_value + self.assertEqual({'hl': True, 'cp': 'aws', 'e': 'arn:aws:lambda:us-east-2:12345:function:TestPython:1'}, span.f) @@ -292,14 +308,17 @@ def test_application_lb_trigger_tracing(self): # The original Lambda handler is set in os.environ["LAMBDA_HANDLER"] result = lambda_handler(event, self.context) - self.assertEqual('All Ok', result) + assert isinstance(result, dict) + assert 'headers' in result + assert 'Server-Timing' in result['headers'] + payload = self.agent.collector.prepare_payload() self.assertTrue("metrics" in payload) self.assertTrue("spans" in payload) self.assertEqual(2, len(payload.keys())) - self.assertTrue(type(payload['metrics']['plugins']) is list) + self.assertTrue(isinstance(payload['metrics']['plugins'], list)) self.assertTrue(len(payload['metrics']['plugins']) == 1) plugin_data = payload['metrics']['plugins'][0] @@ -316,6 +335,9 @@ def test_application_lb_trigger_tracing(self): self.assertIsNotNone(span.ts) self.assertIsNotNone(span.d) + server_timing_value = "intid;desc=%s" % span.t + assert result['headers']['Server-Timing'] == server_timing_value + self.assertEqual({'hl': True, 'cp': 'aws', 'e': 'arn:aws:lambda:us-east-2:12345:function:TestPython:1'}, span.f) @@ -350,14 +372,17 @@ def test_cloudwatch_trigger_tracing(self): # The original Lambda handler is set in os.environ["LAMBDA_HANDLER"] result = lambda_handler(event, self.context) - self.assertEqual('All Ok', result) + assert isinstance(result, dict) + assert 'headers' in result + assert 'Server-Timing' in result['headers'] + payload = self.agent.collector.prepare_payload() self.assertTrue("metrics" in payload) self.assertTrue("spans" in payload) self.assertEqual(2, len(payload.keys())) - self.assertTrue(type(payload['metrics']['plugins']) is list) + self.assertTrue(isinstance(payload['metrics']['plugins'], list)) self.assertTrue(len(payload['metrics']['plugins']) == 1) plugin_data = payload['metrics']['plugins'][0] @@ -374,6 +399,9 @@ def test_cloudwatch_trigger_tracing(self): self.assertIsNotNone(span.ts) self.assertIsNotNone(span.d) + server_timing_value = "intid;desc=%s" % span.t + assert result['headers']['Server-Timing'] == server_timing_value + self.assertEqual({'hl': True, 'cp': 'aws', 'e': 'arn:aws:lambda:us-east-2:12345:function:TestPython:1'}, span.f) @@ -392,7 +420,7 @@ def test_cloudwatch_trigger_tracing(self): self.assertEqual('aws:cloudwatch.events', span.data['lambda']['trigger']) self.assertEqual('cdc73f9d-aea9-11e3-9d5a-835b769c0d9c', span.data["lambda"]["cw"]["events"]["id"]) self.assertEqual(False, span.data["lambda"]["cw"]["events"]["more"]) - self.assertTrue(type(span.data["lambda"]["cw"]["events"]["resources"]) is list) + self.assertTrue(isinstance(span.data["lambda"]["cw"]["events"]["resources"], list)) self.assertEqual(1, len(span.data["lambda"]["cw"]["events"]["resources"])) self.assertEqual('arn:aws:events:eu-west-1:123456789012:rule/ExampleRule', span.data["lambda"]["cw"]["events"]["resources"][0]) @@ -408,14 +436,17 @@ def test_cloudwatch_logs_trigger_tracing(self): # The original Lambda handler is set in os.environ["LAMBDA_HANDLER"] result = lambda_handler(event, self.context) - self.assertEqual('All Ok', result) + assert isinstance(result, dict) + assert 'headers' in result + assert 'Server-Timing' in result['headers'] + payload = self.agent.collector.prepare_payload() self.assertTrue("metrics" in payload) self.assertTrue("spans" in payload) self.assertEqual(2, len(payload.keys())) - self.assertTrue(type(payload['metrics']['plugins']) is list) + self.assertTrue(isinstance(payload['metrics']['plugins'], list)) self.assertTrue(len(payload['metrics']['plugins']) == 1) plugin_data = payload['metrics']['plugins'][0] @@ -432,6 +463,9 @@ def test_cloudwatch_logs_trigger_tracing(self): self.assertIsNotNone(span.ts) self.assertIsNotNone(span.d) + server_timing_value = "intid;desc=%s" % span.t + assert result['headers']['Server-Timing'] == server_timing_value + self.assertEqual({'hl': True, 'cp': 'aws', 'e': 'arn:aws:lambda:us-east-2:12345:function:TestPython:1'}, span.f) @@ -452,7 +486,7 @@ def test_cloudwatch_logs_trigger_tracing(self): self.assertEqual('testLogGroup', span.data['lambda']['cw']['logs']['group']) self.assertEqual('testLogStream', span.data['lambda']['cw']['logs']['stream']) self.assertEqual(None, span.data['lambda']['cw']['logs']['more']) - self.assertTrue(type(span.data['lambda']['cw']['logs']['events']) is list) + self.assertTrue(isinstance(span.data['lambda']['cw']['logs']['events'], list)) self.assertEqual(2, len(span.data['lambda']['cw']['logs']['events'])) self.assertEqual('[ERROR] First test message', span.data['lambda']['cw']['logs']['events'][0]) self.assertEqual('[ERROR] Second test message', span.data['lambda']['cw']['logs']['events'][1]) @@ -468,14 +502,17 @@ def test_s3_trigger_tracing(self): # The original Lambda handler is set in os.environ["LAMBDA_HANDLER"] result = lambda_handler(event, self.context) - self.assertEqual('All Ok', result) + assert isinstance(result, dict) + assert 'headers' in result + assert 'Server-Timing' in result['headers'] + payload = self.agent.collector.prepare_payload() self.assertTrue("metrics" in payload) self.assertTrue("spans" in payload) self.assertEqual(2, len(payload.keys())) - self.assertTrue(type(payload['metrics']['plugins']) is list) + self.assertTrue(isinstance(payload['metrics']['plugins'], list)) self.assertTrue(len(payload['metrics']['plugins']) == 1) plugin_data = payload['metrics']['plugins'][0] @@ -492,6 +529,9 @@ def test_s3_trigger_tracing(self): self.assertIsNotNone(span.ts) self.assertIsNotNone(span.d) + server_timing_value = "intid;desc=%s" % span.t + assert result['headers']['Server-Timing'] == server_timing_value + self.assertEqual({'hl': True, 'cp': 'aws', 'e': 'arn:aws:lambda:us-east-2:12345:function:TestPython:1'}, span.f) @@ -508,7 +548,7 @@ def test_s3_trigger_tracing(self): self.assertIsNone(span.data['service']) self.assertEqual('aws:s3', span.data['lambda']['trigger']) - self.assertTrue(type(span.data["lambda"]["s3"]["events"]) is list) + self.assertTrue(isinstance(span.data["lambda"]["s3"]["events"], list)) events = span.data["lambda"]["s3"]["events"] self.assertEqual(1, len(events)) event = events[0] @@ -527,14 +567,17 @@ def test_sqs_trigger_tracing(self): # The original Lambda handler is set in os.environ["LAMBDA_HANDLER"] result = lambda_handler(event, self.context) - self.assertEqual('All Ok', result) + assert isinstance(result, dict) + assert 'headers' in result + assert 'Server-Timing' in result['headers'] + payload = self.agent.collector.prepare_payload() self.assertTrue("metrics" in payload) self.assertTrue("spans" in payload) self.assertEqual(2, len(payload.keys())) - self.assertTrue(type(payload['metrics']['plugins']) is list) + self.assertTrue(isinstance(payload['metrics']['plugins'], list)) self.assertTrue(len(payload['metrics']['plugins']) == 1) plugin_data = payload['metrics']['plugins'][0] @@ -551,6 +594,9 @@ def test_sqs_trigger_tracing(self): self.assertIsNotNone(span.ts) self.assertIsNotNone(span.d) + server_timing_value = "intid;desc=%s" % span.t + assert result['headers']['Server-Timing'] == server_timing_value + self.assertEqual({'hl': True, 'cp': 'aws', 'e': 'arn:aws:lambda:us-east-2:12345:function:TestPython:1'}, span.f) @@ -567,7 +613,7 @@ def test_sqs_trigger_tracing(self): self.assertIsNone(span.data['service']) self.assertEqual('aws:sqs', span.data['lambda']['trigger']) - self.assertTrue(type(span.data["lambda"]["sqs"]["messages"]) is list) + self.assertTrue(isinstance(span.data["lambda"]["sqs"]["messages"], list)) messages = span.data["lambda"]["sqs"]["messages"] self.assertEqual(1, len(messages)) message = messages[0]