From 831f248877bb3896c9279a69776a5d2e4596f168 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ferenc=20G=C3=A9czi?= <ferenc.geczi@ibm.com>
Date: Mon, 8 Jan 2024 12:00:00 +0000
Subject: [PATCH] fix(lambda): Record the statusCode and errors
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Ferenc Géczi <ferenc.geczi@ibm.com>
---
 lib/instana/instrumentation/aws_sdk_lambda.rb | 15 ++++++-
 test/instrumentation/aws_test.rb              | 44 +++++++++++++++++++
 2 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/lib/instana/instrumentation/aws_sdk_lambda.rb b/lib/instana/instrumentation/aws_sdk_lambda.rb
index 8217183f..f684b263 100644
--- a/lib/instana/instrumentation/aws_sdk_lambda.rb
+++ b/lib/instana/instrumentation/aws_sdk_lambda.rb
@@ -25,7 +25,20 @@ def call(context)
             type: context.params[:invocation_type]
           }.reject { |_, v| v.nil? }
 
-          ::Instana.tracer.trace(:"aws.lambda.invoke", {aws: {lambda: {invoke: tags}}}) { @handler.call(context) }
+          span = ::Instana.tracer.log_entry(:"aws.lambda.invoke", {aws: {lambda: {invoke: tags}}})
+          begin
+            response = @handler.call(context)
+            if response.respond_to? :status_code
+              span.set_tags(:http => {:status => response.status_code })
+
+            end
+            response
+          rescue => e
+            ::Instana.tracer.log_error(e)
+            raise
+          ensure
+            ::Instana.tracer.log_exit(:"aws.lambda.invoke")
+          end
         end
       end
 
diff --git a/test/instrumentation/aws_test.rb b/test/instrumentation/aws_test.rb
index cc16bf9d..74ca4976 100644
--- a/test/instrumentation/aws_test.rb
+++ b/test/instrumentation/aws_test.rb
@@ -193,5 +193,49 @@ def test_lambda
     assert_equal :"aws.lambda.invoke", lambda_span[:n]
     assert_equal 'Test', lambda_span[:data][:aws][:lambda][:invoke][:function]
     assert_equal 'Event', lambda_span[:data][:aws][:lambda][:invoke][:type]
+    assert_equal 200, lambda_span[:data][:http][:status]
+    assert_nil lambda_span[:ec]
+    assert_nil lambda_span[:stack]
+  end
+
+  def test_lambda_with_500_status
+    stub_request(:post, "https://lambda.local.amazonaws.com/2015-03-31/functions/Test/invocations")
+      .with(
+        body: "data",
+        headers: {
+          'X-Amz-Client-Context' => /.+/
+        }
+      )
+      .to_return(status: 500, body: '{"message": "Internal Server Error" }', headers: {})
+
+    lambda = Aws::Lambda::Client.new(
+      endpoint: 'https://lambda.local.amazonaws.com',
+      region: 'local',
+      access_key_id: "test",
+      secret_access_key: "test"
+    )
+
+    assert_raises(RuntimeError) do
+      Instana::Tracer.start_or_continue_trace(:lambda_test, {}) do
+        lambda.invoke(
+          function_name: 'Test',
+          invocation_type: 'Event',
+          payload: 'data'
+        )
+      end
+    end
+
+    spans = ::Instana.processor.queued_spans
+    lambda_span, _entry_span, *rest = spans
+
+    assert rest.empty?
+
+    assert_equal :"aws.lambda.invoke", lambda_span[:n]
+    assert_equal 'Test', lambda_span[:data][:aws][:lambda][:invoke][:function]
+    assert_equal 'Event', lambda_span[:data][:aws][:lambda][:invoke][:type]
+    refute_nil lambda_span[:ec]
+    assert_equal 1, lambda_span[:ec]
+    refute_nil lambda_span[:stack]
+    assert_equal 30, lambda_span[:stack].length # default limit is 30 in span.add_stack
   end
 end