Supports breakpoint canary in GCP

The breakpoint canary feature is controlled using the new 'breakpoint_enable_canary' and 'breakpoint_allow_canary_override' parameters when calling enable() or the same-name flags when run as a module.

PiperOrigin-RevId: 315558797
Change-Id: I13c7a2077055ecfb6e8c28c726a8315b60ce505e

Author: Louis Ye

README.md

@@ -271,3 +271,9 @@ Default Credentials](https://cloud.google.com/docs/authentication/production)
 which are automatically available on machines hosted on GCP, or can be set via
 `gcloud auth application-default login` or the `GOOGLE_APPLICATION_CREDENTIALS`
 environment variable.
+
+`breakpoint_enable_canary`: Whether to enable the
+[breakpoint canary feature](https://cloud.google.com/debugger/docs/using/snapshots#with_canarying).
+It expects a boolean value (`True`/`False`) or a string, with `'True'`
+interpreted as `True` and any other string interpreted as `False`). If not
+provided, the breakpoint canarying will not be enabled.

src/googleclouddebugger/__init__.py

@@ -72,6 +72,9 @@ def _StartDebugger():
       _flags.get('project_id'),
       _flags.get('project_number'),
       _flags.get('service_account_json_file'))
+  _hub_client.SetupCanaryMode(
+      _flags.get('breakpoint_enable_canary'),
+      _flags.get('breakpoint_allow_canary_override'))
   _hub_client.InitializeDebuggeeLabels(_flags)
   _hub_client.Start()
 

src/googleclouddebugger/gcp_hub_client.py

@@ -75,6 +75,15 @@
 # of 40 seconds.
 _HTTP_TIMEOUT_SECONDS = 100
 
+# The map from the values of flags (breakpoint_enable_canary,
+# breakpoint_allow_canary_override) to canary mode.
+_CANARY_MODE_MAP = {
+    (True, True): 'CANARY_MODE_DEFAULT_ENABLED',
+    (True, False): 'CANARY_MODE_ALWAYS_ENABLED',
+    (False, True): 'CANARY_MODE_DEFAULT_DISABLED',
+    (False, False): 'CANARY_MODE_ALWAYS_DISABLED',
+}
+
 
 class NoProjectIdError(Exception):
   """Used to indicate the project id cannot be determined."""
@@ -101,6 +110,8 @@ def __init__(self):
     self._debuggee_labels = {}
     self._service_account_auth = False
     self._debuggee_id = None
+    self._agent_id = None
+    self._canary_mode = None
     self._wait_token = 'init'
     self._breakpoints = []
     self._main_thread = None
@@ -218,6 +229,21 @@ def SetupAuth(self,
     self._project_id = project_id
     self._project_number = project_number or project_id
 
+  def SetupCanaryMode(self, breakpoint_enable_canary,
+                      breakpoint_allow_canary_override):
+    """Sets up canaryMode for the debuggee according to input parameters.
+
+    Args:
+      breakpoint_enable_canary: str or bool, whether to enable breakpoint
+          canary. Any string except 'True' is interpreted as False.
+      breakpoint_allow_canary_override: str or bool, whether to allow the
+          individually set breakpoint to override the canary behavior. Any
+          string except 'True' is interpreted as False.
+    """
+    enable_canary = breakpoint_enable_canary in ('True', True)
+    allow_canary_override = breakpoint_allow_canary_override in ('True', True)
+    self._canary_mode = _CANARY_MODE_MAP[enable_canary, allow_canary_override]
+
   def Start(self):
     """Starts the worker thread."""
     self._shutdown = False
@@ -327,8 +353,11 @@ def _RegisterDebuggee(self, service):
         self._project_number = project_number or self._project_number
 
         self._debuggee_id = response['debuggee']['id']
-        native.LogInfo('Debuggee registered successfully, ID: %s' % (
-            self._debuggee_id))
+        self._agent_id = response['agentId']
+        native.LogInfo(
+            'Debuggee registered successfully, ID: %s, agent ID: %s, '
+            'canary mode: %s' % (self._debuggee_id, self._agent_id,
+                                 response['debuggee'].get('canaryMode')))
         self.register_backoff.Succeeded()
         return (False, 0)  # Proceed immediately to list active breakpoints.
       except BaseException:
@@ -355,7 +384,9 @@ def _ListActiveBreakpoints(self, service):
     """
     try:
       response = service.debuggees().breakpoints().list(
-          debuggeeId=self._debuggee_id, waitToken=self._wait_token,
+          debuggeeId=self._debuggee_id,
+          agentId=self._agent_id,
+          waitToken=self._wait_token,
           successOnTimeout=True).execute()
       if not response.get('waitExpired'):
         self._wait_token = response.get('nextWaitToken')
@@ -469,6 +500,7 @@ def _GetDebuggee(self):
         'description': self._GetDebuggeeDescription(),
         'labels': self._debuggee_labels,
         'agentVersion': agent_version,
+        'canaryMode': self._canary_mode,
     }
 
     source_context = self._ReadAppJsonFile('source-context.json')