Add operation name frame in stacks with boto api calls

When using API calls from boto we often do not see the operation names
in the stack traces even though there is a dedicate function name. This
is because of the generic way that api calls are implemented in boto.
Unfortunately this is not convenient for profiling as we do not see in
flamegraphs which calls are using resources.

This change adds a specific check when we build the stacks to detect api
calls inside boto and if possible add a new frame with the operation
function name.

i.e. instead of
```
(...)
botocore.client:CloudWatch:_make_api_call
botocore.client:CloudWatch:_api_call
user_module:send_metric
```

we should end up with
```
(...)
botocore.client:CloudWatch:_make_api_call
botocore.client:CloudWatch:_api_call
botocore.client:CloudWatch:put_metric_data
user_module:send_metric
```

This new frame should help when investigating on flamegraphs but also
with automatic analysis.

Author: PapaPedro

codeguru_profiler_agent/sampling_utils.py

@@ -4,9 +4,10 @@
 import linecache
 import threading
 import traceback
-
+import re
 from codeguru_profiler_agent.model.frame import Frame
 
+BOTO_CLIENT_PATH = re.compile("[/\\\\]botocore[/\\\\]client.py$")
 TRUNCATED_FRAME = Frame(name="<Truncated>")
 
 TIME_SLEEP_FRAME = Frame(name="<Sleep>")
@@ -64,15 +65,35 @@ def _extract_stack(stack, max_depth):
     """
     result = []
     for raw_frame, line_no in stack:
+        _maybe_add_boto_operation_name(raw_frame, result)
         co = raw_frame.f_code
         result.append(
             Frame(name=co.co_name, class_name=_extract_class(raw_frame.f_locals), line_no=line_no,
                   file_path=co.co_filename)
         )
-    if len(stack) < max_depth:
+    if len(result) < max_depth:
         last_frame, last_frame_line_no = stack[-1]
         _maybe_append_synthetic_frame(result, last_frame, last_frame_line_no)
-    return result
+    return result[:max_depth]
+
+
+def _maybe_add_boto_operation_name(raw_frame, result):
+    """
+    boto is dealing with API calls in a very generic way so by default the sampling
+    would only show that we are making an api call without having the actual operation name.
+    This function checks if this frame is botocore.client.py:_api_call and if it is, it adds
+    a frame with the actual operation name.
+    :param raw_frame: the raw frame
+    """
+    if (raw_frame.f_code.co_name == '_api_call'
+            and BOTO_CLIENT_PATH.search(raw_frame.f_code.co_filename) is not None
+            and raw_frame.f_locals and 'py_operation_name' in raw_frame.f_locals.keys()
+            and raw_frame.f_locals.get('py_operation_name')):
+        result.append(
+            Frame(name=raw_frame.f_locals.get('py_operation_name'),
+                  class_name=_extract_class(raw_frame.f_locals),
+                  file_path=raw_frame.f_code.co_filename)
+        )
 
 
 def _maybe_append_synthetic_frame(result, frame, line_no):

test/help_utils.py

@@ -1,6 +1,8 @@
 import threading
 import time
 from queue import Queue
+import boto3
+from botocore.client import Config
 
 INTEGRATION_TEST_ACCOUNT_ID = "519630429520"
 MY_PROFILING_GROUP_NAME_FOR_INTEG_TESTS = "MyProfilingGroupForIntegrationTests"
@@ -48,6 +50,23 @@ def dummy_parent_method(self):
         """
         self.dummy_method()
 
+    def make_boto_api_call(self, boto_client):
+        try:
+            boto_client.put_metric_data(Namespace="any_namespace", MetricData=[])
+        except Exception as e:
+            print("This should be a ConnectTimeoutError", e)
+
+    def new_thread_sending_boto_api_call(self, timeout_seconds=1, thread_name="test-boto-api-call"):
+        no_retry_config = Config(connect_timeout=timeout_seconds, retries={'max_attempts': 0})
+        # we do not want boto to look for real credentials so provide fake ones
+        session = boto3.Session(region_name="us-east-1", aws_access_key_id="fake_id", aws_secret_access_key="fake_key")
+        # we set a fake endpoint in the client because we do not want to make a real call
+        # this is only so we can have a thread inside an api call trying to make a connection
+        # long enough for us to take a sample
+        no_target_client = session.client('cloudwatch', endpoint_url='https://notExisting.com/', config=no_retry_config)
+        self.boto_thread = threading.Thread(
+            name=thread_name, target=self.make_boto_api_call, daemon=True, kwargs={"boto_client": no_target_client})
+        self.boto_thread.start()
 
 def wait_for(condition, timeout_seconds=1.0, poll_interval_seconds=0.01):
     """

test/unit/test_sampling_utils.py

@@ -3,11 +3,20 @@
 import sys
 
 from test import help_utils
+from collections import namedtuple
 
 from codeguru_profiler_agent.sampling_utils import get_stacks
 
 DEFAULT_TRUNCATED_FRAME_NAME = "<Truncated>"
 
+test_code = namedtuple('code', ['co_filename', 'co_name'])
+test_frame = namedtuple('frame', ['f_code', 'f_locals'])
+test_tb = namedtuple('tb', ['tb_frame', 'tb_lineno'])
+
+
+def make_frame(path, method, line_nbr, f_locals={}):
+    return test_tb(test_frame(test_code(path, method), f_locals), line_nbr)
+
 
 def is_frame_in_stacks(stacks, target_frame):
     for stack in stacks:
@@ -88,3 +97,52 @@ def test_it_does_not_include_zombie_threads(self):
 
                 assert not is_frame_in_stacks(
                     stacks, "dummy_parent_method")
+
+        def test_it_adds_operation_name_frame_for_boto(self):
+            raw_stack = [
+                make_frame('path/to/foo.py', 'foo', 3),
+                make_frame('site-packages/botocore/client.py', '_api_call', 3, {'py_operation_name': 'boto_api_call'}),
+                make_frame('path/to/bar.py', 'bar', 3)
+            ]
+            with mock.patch(
+                    "traceback.walk_stack",
+                    side_effect=
+                    lambda end_frame: raw_stack
+            ):
+                stacks = get_stacks(
+                    threads_to_sample=sys._current_frames().items(),
+                    excluded_threads=set(),
+                    max_depth=100)
+                assert len(stacks[0]) == 4
+                assert is_frame_in_stacks(stacks, "boto_api_call")
+
+        def test_adding_boto_frame_does_not_exceed_maximum_depth(self):
+            raw_stack = [
+                make_frame('site-packages/botocore/client.py', '_api_call', 34, {'py_operation_name': 'boto_api_call'}),
+                make_frame('path/to/foo.py', 'bar', 12),
+
+            ]
+            for i in range(100):
+                raw_stack.insert(0, make_frame('path/to/foo.py', 'bar' + str(i), 1))
+            with mock.patch(
+                    "traceback.walk_stack",
+                    side_effect=
+                    lambda end_frame: raw_stack
+            ):
+                stacks = get_stacks(
+                    threads_to_sample=sys._current_frames().items(),
+                    excluded_threads=set(),
+                    max_depth=100)
+                assert len(stacks[0]) == 100
+                assert is_frame_in_stacks(stacks, "boto_api_call")
+
+        def test_it_adds_operation_name_frame_for_real_boto_call(self):
+            # Run a thread that will try to do a boto3 api call for 1 second then fail with a log
+            # the function will call put_metric_data on a cloudwatch client
+            # so get_stack should capture it.
+            self.helper.new_thread_sending_boto_api_call(timeout_seconds=1)
+            stacks = get_stacks(
+                threads_to_sample=sys._current_frames().items(),
+                excluded_threads=set(),
+                max_depth=100)
+            assert is_frame_in_stacks(stacks, "put_metric_data")