Skip to content

Commit 90127c1

Browse files
author
Peter Giacomo Lombardo
authored
AWS Lambda: Assure Server-Timing response header (#258)
1 parent f6cf9cc commit 90127c1

File tree

2 files changed

+78
-20
lines changed

2 files changed

+78
-20
lines changed

‎instana/instrumentation/aws/lambda_inst.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,18 @@ def lambda_handler_with_instana(wrapped, instance, args, kwargs):
2424
enrich_lambda_span(agent, scope.span, *args)
2525
try:
2626
result = wrapped(*args, **kwargs)
27+
28+
if isinstance(result, dict):
29+
server_timing_value = "intid;desc=%s" % scope.span.context.trace_id
30+
if 'headers' in result:
31+
result['headers']['Server-Timing'] = server_timing_value
32+
elif 'multiValueHeaders' in result:
33+
result['multiValueHeaders']['Server-Timing'] = [server_timing_value]
34+
else:
35+
# If both 'headers' and 'multiValueHeaders' aren't in result,
36+
# then default to setting single value 'headers'
37+
result['headers'] = dict()
38+
result['headers']['Server-Timing'] = server_timing_value
2739
except Exception as exc:
2840
if scope.span:
2941
scope.span.log_exception(exc)

‎tests/platforms/test_lambda.py‎

Lines changed: 66 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,11 @@ def __init__(self, **kwargs):
3131
# This is the target handler that will be instrumented for these tests
3232
def my_lambda_handler(event, context):
3333
# print("target_handler called")
34-
return "All Ok"
34+
return {
35+
'statusCode': 200,
36+
'headers': {'Content-Type': 'application/json'},
37+
'body': json.dumps({'site': 'pwpush.com', 'response': 204})
38+
}
3539

3640
# We only want to monkey patch the test handler once so do it here
3741
os.environ["LAMBDA_HANDLER"] = "tests.platforms.test_lambda.my_lambda_handler"
@@ -118,7 +122,7 @@ def test_has_extra_http_headers(self):
118122
def test_has_options(self):
119123
self.create_agent_and_setup_tracer()
120124
self.assertTrue(hasattr(self.agent, 'options'))
121-
self.assertTrue(type(self.agent.options) is AWSLambdaOptions)
125+
self.assertTrue(isinstance(self.agent.options, AWSLambdaOptions))
122126
assert(self.agent.options.endpoint_proxy == { })
123127

124128
def test_get_handler(self):
@@ -173,14 +177,17 @@ def test_custom_service_name(self):
173177
result = lambda_handler(event, self.context)
174178
os.environ.pop('INSTANA_SERVICE_NAME')
175179

176-
self.assertEqual('All Ok', result)
180+
assert isinstance(result, dict)
181+
assert 'headers' in result
182+
assert 'Server-Timing' in result['headers']
183+
177184
payload = self.agent.collector.prepare_payload()
178185

179186
self.assertTrue("metrics" in payload)
180187
self.assertTrue("spans" in payload)
181188
self.assertEqual(2, len(payload.keys()))
182189

183-
self.assertTrue(type(payload['metrics']['plugins']) is list)
190+
self.assertTrue(isinstance(payload['metrics']['plugins'], list))
184191
self.assertTrue(len(payload['metrics']['plugins']) == 1)
185192
plugin_data = payload['metrics']['plugins'][0]
186193

@@ -197,6 +204,9 @@ def test_custom_service_name(self):
197204
self.assertIsNotNone(span.ts)
198205
self.assertIsNotNone(span.d)
199206

207+
server_timing_value = "intid;desc=%s" % span.t
208+
assert result['headers']['Server-Timing'] == server_timing_value
209+
200210
self.assertEqual({'hl': True, 'cp': 'aws', 'e': 'arn:aws:lambda:us-east-2:12345:function:TestPython:1'},
201211
span.f)
202212

@@ -233,14 +243,17 @@ def test_api_gateway_trigger_tracing(self):
233243
# The original Lambda handler is set in os.environ["LAMBDA_HANDLER"]
234244
result = lambda_handler(event, self.context)
235245

236-
self.assertEqual('All Ok', result)
246+
assert isinstance(result, dict)
247+
assert 'headers' in result
248+
assert 'Server-Timing' in result['headers']
249+
237250
payload = self.agent.collector.prepare_payload()
238251

239252
self.assertTrue("metrics" in payload)
240253
self.assertTrue("spans" in payload)
241254
self.assertEqual(2, len(payload.keys()))
242255

243-
self.assertTrue(type(payload['metrics']['plugins']) is list)
256+
self.assertTrue(isinstance(payload['metrics']['plugins'], list))
244257
self.assertTrue(len(payload['metrics']['plugins']) == 1)
245258
plugin_data = payload['metrics']['plugins'][0]
246259

@@ -257,6 +270,9 @@ def test_api_gateway_trigger_tracing(self):
257270
self.assertIsNotNone(span.ts)
258271
self.assertIsNotNone(span.d)
259272

273+
server_timing_value = "intid;desc=%s" % span.t
274+
assert result['headers']['Server-Timing'] == server_timing_value
275+
260276
self.assertEqual({'hl': True, 'cp': 'aws', 'e': 'arn:aws:lambda:us-east-2:12345:function:TestPython:1'},
261277
span.f)
262278

@@ -292,14 +308,17 @@ def test_application_lb_trigger_tracing(self):
292308
# The original Lambda handler is set in os.environ["LAMBDA_HANDLER"]
293309
result = lambda_handler(event, self.context)
294310

295-
self.assertEqual('All Ok', result)
311+
assert isinstance(result, dict)
312+
assert 'headers' in result
313+
assert 'Server-Timing' in result['headers']
314+
296315
payload = self.agent.collector.prepare_payload()
297316

298317
self.assertTrue("metrics" in payload)
299318
self.assertTrue("spans" in payload)
300319
self.assertEqual(2, len(payload.keys()))
301320

302-
self.assertTrue(type(payload['metrics']['plugins']) is list)
321+
self.assertTrue(isinstance(payload['metrics']['plugins'], list))
303322
self.assertTrue(len(payload['metrics']['plugins']) == 1)
304323
plugin_data = payload['metrics']['plugins'][0]
305324

@@ -316,6 +335,9 @@ def test_application_lb_trigger_tracing(self):
316335
self.assertIsNotNone(span.ts)
317336
self.assertIsNotNone(span.d)
318337

338+
server_timing_value = "intid;desc=%s" % span.t
339+
assert result['headers']['Server-Timing'] == server_timing_value
340+
319341
self.assertEqual({'hl': True, 'cp': 'aws', 'e': 'arn:aws:lambda:us-east-2:12345:function:TestPython:1'},
320342
span.f)
321343

@@ -350,14 +372,17 @@ def test_cloudwatch_trigger_tracing(self):
350372
# The original Lambda handler is set in os.environ["LAMBDA_HANDLER"]
351373
result = lambda_handler(event, self.context)
352374

353-
self.assertEqual('All Ok', result)
375+
assert isinstance(result, dict)
376+
assert 'headers' in result
377+
assert 'Server-Timing' in result['headers']
378+
354379
payload = self.agent.collector.prepare_payload()
355380

356381
self.assertTrue("metrics" in payload)
357382
self.assertTrue("spans" in payload)
358383
self.assertEqual(2, len(payload.keys()))
359384

360-
self.assertTrue(type(payload['metrics']['plugins']) is list)
385+
self.assertTrue(isinstance(payload['metrics']['plugins'], list))
361386
self.assertTrue(len(payload['metrics']['plugins']) == 1)
362387
plugin_data = payload['metrics']['plugins'][0]
363388

@@ -374,6 +399,9 @@ def test_cloudwatch_trigger_tracing(self):
374399
self.assertIsNotNone(span.ts)
375400
self.assertIsNotNone(span.d)
376401

402+
server_timing_value = "intid;desc=%s" % span.t
403+
assert result['headers']['Server-Timing'] == server_timing_value
404+
377405
self.assertEqual({'hl': True, 'cp': 'aws', 'e': 'arn:aws:lambda:us-east-2:12345:function:TestPython:1'},
378406
span.f)
379407

@@ -392,7 +420,7 @@ def test_cloudwatch_trigger_tracing(self):
392420
self.assertEqual('aws:cloudwatch.events', span.data['lambda']['trigger'])
393421
self.assertEqual('cdc73f9d-aea9-11e3-9d5a-835b769c0d9c', span.data["lambda"]["cw"]["events"]["id"])
394422
self.assertEqual(False, span.data["lambda"]["cw"]["events"]["more"])
395-
self.assertTrue(type(span.data["lambda"]["cw"]["events"]["resources"]) is list)
423+
self.assertTrue(isinstance(span.data["lambda"]["cw"]["events"]["resources"], list))
396424
self.assertEqual(1, len(span.data["lambda"]["cw"]["events"]["resources"]))
397425
self.assertEqual('arn:aws:events:eu-west-1:123456789012:rule/ExampleRule',
398426
span.data["lambda"]["cw"]["events"]["resources"][0])
@@ -408,14 +436,17 @@ def test_cloudwatch_logs_trigger_tracing(self):
408436
# The original Lambda handler is set in os.environ["LAMBDA_HANDLER"]
409437
result = lambda_handler(event, self.context)
410438

411-
self.assertEqual('All Ok', result)
439+
assert isinstance(result, dict)
440+
assert 'headers' in result
441+
assert 'Server-Timing' in result['headers']
442+
412443
payload = self.agent.collector.prepare_payload()
413444

414445
self.assertTrue("metrics" in payload)
415446
self.assertTrue("spans" in payload)
416447
self.assertEqual(2, len(payload.keys()))
417448

418-
self.assertTrue(type(payload['metrics']['plugins']) is list)
449+
self.assertTrue(isinstance(payload['metrics']['plugins'], list))
419450
self.assertTrue(len(payload['metrics']['plugins']) == 1)
420451
plugin_data = payload['metrics']['plugins'][0]
421452

@@ -432,6 +463,9 @@ def test_cloudwatch_logs_trigger_tracing(self):
432463
self.assertIsNotNone(span.ts)
433464
self.assertIsNotNone(span.d)
434465

466+
server_timing_value = "intid;desc=%s" % span.t
467+
assert result['headers']['Server-Timing'] == server_timing_value
468+
435469
self.assertEqual({'hl': True, 'cp': 'aws', 'e': 'arn:aws:lambda:us-east-2:12345:function:TestPython:1'},
436470
span.f)
437471

@@ -452,7 +486,7 @@ def test_cloudwatch_logs_trigger_tracing(self):
452486
self.assertEqual('testLogGroup', span.data['lambda']['cw']['logs']['group'])
453487
self.assertEqual('testLogStream', span.data['lambda']['cw']['logs']['stream'])
454488
self.assertEqual(None, span.data['lambda']['cw']['logs']['more'])
455-
self.assertTrue(type(span.data['lambda']['cw']['logs']['events']) is list)
489+
self.assertTrue(isinstance(span.data['lambda']['cw']['logs']['events'], list))
456490
self.assertEqual(2, len(span.data['lambda']['cw']['logs']['events']))
457491
self.assertEqual('[ERROR] First test message', span.data['lambda']['cw']['logs']['events'][0])
458492
self.assertEqual('[ERROR] Second test message', span.data['lambda']['cw']['logs']['events'][1])
@@ -468,14 +502,17 @@ def test_s3_trigger_tracing(self):
468502
# The original Lambda handler is set in os.environ["LAMBDA_HANDLER"]
469503
result = lambda_handler(event, self.context)
470504

471-
self.assertEqual('All Ok', result)
505+
assert isinstance(result, dict)
506+
assert 'headers' in result
507+
assert 'Server-Timing' in result['headers']
508+
472509
payload = self.agent.collector.prepare_payload()
473510

474511
self.assertTrue("metrics" in payload)
475512
self.assertTrue("spans" in payload)
476513
self.assertEqual(2, len(payload.keys()))
477514

478-
self.assertTrue(type(payload['metrics']['plugins']) is list)
515+
self.assertTrue(isinstance(payload['metrics']['plugins'], list))
479516
self.assertTrue(len(payload['metrics']['plugins']) == 1)
480517
plugin_data = payload['metrics']['plugins'][0]
481518

@@ -492,6 +529,9 @@ def test_s3_trigger_tracing(self):
492529
self.assertIsNotNone(span.ts)
493530
self.assertIsNotNone(span.d)
494531

532+
server_timing_value = "intid;desc=%s" % span.t
533+
assert result['headers']['Server-Timing'] == server_timing_value
534+
495535
self.assertEqual({'hl': True, 'cp': 'aws', 'e': 'arn:aws:lambda:us-east-2:12345:function:TestPython:1'},
496536
span.f)
497537

@@ -508,7 +548,7 @@ def test_s3_trigger_tracing(self):
508548
self.assertIsNone(span.data['service'])
509549

510550
self.assertEqual('aws:s3', span.data['lambda']['trigger'])
511-
self.assertTrue(type(span.data["lambda"]["s3"]["events"]) is list)
551+
self.assertTrue(isinstance(span.data["lambda"]["s3"]["events"], list))
512552
events = span.data["lambda"]["s3"]["events"]
513553
self.assertEqual(1, len(events))
514554
event = events[0]
@@ -527,14 +567,17 @@ def test_sqs_trigger_tracing(self):
527567
# The original Lambda handler is set in os.environ["LAMBDA_HANDLER"]
528568
result = lambda_handler(event, self.context)
529569

530-
self.assertEqual('All Ok', result)
570+
assert isinstance(result, dict)
571+
assert 'headers' in result
572+
assert 'Server-Timing' in result['headers']
573+
531574
payload = self.agent.collector.prepare_payload()
532575

533576
self.assertTrue("metrics" in payload)
534577
self.assertTrue("spans" in payload)
535578
self.assertEqual(2, len(payload.keys()))
536579

537-
self.assertTrue(type(payload['metrics']['plugins']) is list)
580+
self.assertTrue(isinstance(payload['metrics']['plugins'], list))
538581
self.assertTrue(len(payload['metrics']['plugins']) == 1)
539582
plugin_data = payload['metrics']['plugins'][0]
540583

@@ -551,6 +594,9 @@ def test_sqs_trigger_tracing(self):
551594
self.assertIsNotNone(span.ts)
552595
self.assertIsNotNone(span.d)
553596

597+
server_timing_value = "intid;desc=%s" % span.t
598+
assert result['headers']['Server-Timing'] == server_timing_value
599+
554600
self.assertEqual({'hl': True, 'cp': 'aws', 'e': 'arn:aws:lambda:us-east-2:12345:function:TestPython:1'},
555601
span.f)
556602

@@ -567,7 +613,7 @@ def test_sqs_trigger_tracing(self):
567613
self.assertIsNone(span.data['service'])
568614

569615
self.assertEqual('aws:sqs', span.data['lambda']['trigger'])
570-
self.assertTrue(type(span.data["lambda"]["sqs"]["messages"]) is list)
616+
self.assertTrue(isinstance(span.data["lambda"]["sqs"]["messages"], list))
571617
messages = span.data["lambda"]["sqs"]["messages"]
572618
self.assertEqual(1, len(messages))
573619
message = messages[0]

0 commit comments

Comments
 (0)