Experiment with wrapping for gapic retries

Author: odeke-em

google/cloud/spanner_v1/_helpers.py

@@ -576,6 +576,7 @@ def _retry(
 
 
 def _check_rst_stream_error(exc):
+    print("\033[31mrst_", exc, "\033[00m")
     resumable_error = (
         any(
             resumable_message in exc.message
@@ -589,6 +590,11 @@ def _check_rst_stream_error(exc):
         raise
 
 
+def _check_unavailable(exc):
+    print("\033[31mcheck_unavailable", exc, "\033[00m")
+    raise
+
+
 def _metadata_with_leader_aware_routing(value, **kw):
     """Create RPC metadata containing a leader aware routing header
 
@@ -763,96 +769,164 @@ def __init__(self, original_callable: Callable):
 
 
 def inject_retry_header_control(api):
-    return
-    monkey_patch(type(api))
+    # monkey_patch(type(api))
+    # monkey_patch(api)
+    pass
 
-memoize_map = dict()
 
-def monkey_patch(obj):
-    return
+def monkey_patch(typ):
+    keys = dir(typ)
+    attempts = dict()
+    for key in keys:
+        if key.startswith("_"):
+            continue
 
-    """
-    klass = obj
-    attrs = dir(klass)
-    for attr_key in attrs:
-        if attr_key.startswith('_'):
+        if key != "batch_create_sessions":
             continue
 
-        attr_value = getattr(obj, attr_key)
-        if not callable(attr_value):
+        fn = getattr(typ, key)
+
+        signature = inspect.signature(fn)
+        if signature.parameters.get("metadata", None) is None:
             continue
 
-        signature = inspect.signature(attr_value)
-        print(attr_key, signature.parameters)
+        print("fn.__call__", inspect.getsource(fn))
 
-        call = attr_value
-        # Our goal is to replace the runtime pass through.
-        def wrapped(*args, **kwargs):
-            print(attr_key, 'called')
-            return call(*args, **kwargs)
+        def as_proxy(db, *args, **kwargs):
+            print("db_key", hex(id(db)))
+            print("as_proxy", args, kwargs)
+            metadata = kwargs.get("metadata", None)
+            if not metadata:
+                return fn(db, *args, **kwargs)
 
-        setattr(klass, attr_key, wrapped)
+            hash_key = hex(id(db)) + "." + hex(id(key))
+            attempts.setdefault(hash_key, 0)
+            attempts[hash_key] += 1
+            # 4. Find all the headers that match the target header key.
+            all_metadata = []
+            for mkey, value in metadata:
+                if mkey is not REQ_ID_HEADER_KEY:
+                    continue
 
-    return
-    """
+                splits = value.split(".")
+                # 5. Increment the original_attempt with that of our re-invocation count.
+                print("\033[34mkey", mkey, "\033[00m", splits)
+                hdr_attempt_plus_reinvocation = int(splits[-1]) + attempts[hash_key]
+                splits[-1] = str(hdr_attempt_plus_reinvocation)
+                value = ".".join(splits)
 
+                all_metadata.append((mkey, value))
+
+            kwargs["metadata"] = all_metadata
+            return fn(db, *args, **kwargs)
+
+        setattr(typ, key, as_proxy)
+
+
+def alt_foo():
+    memoize_map = dict()
     orig_get_attr = getattr(obj, "__getattribute__")
+    hex_orig = hex(id(orig_get_attr))
+    hex_patched = None
+
     def patched_getattribute(obj, key, *args, **kwargs):
-        if key.startswith('_'):
+        if key.startswith("_"):
             return orig_get_attr(obj, key, *args, **kwargs)
 
-        orig_value = orig_get_attr(obj, key, *args, **kwargs)
-        if not callable(orig_value):
-            return orig_value
+        if key != "batch_create_sessions":
+            return orig_get_attr(obj, key, *args, **kwargs)
 
         map_key = hex(id(key)) + hex(id(obj))
         memoized = memoize_map.get(map_key, None)
         if memoized:
-            print("memoized_hit", key, '\033[35m', inspect.getsource(orig_value), '\033[00m')
+            if False:
+                print(
+                    "memoized_hit",
+                    key,
+                    "\033[35m",
+                    inspect.getsource(orig_value),
+                    "\033[00m",
+                )
+            print("memoized_hit", key, "\033[35m", map_key, "\033[00m")
             return memoized
 
+        orig_value = orig_get_attr(obj, key, *args, **kwargs)
+        if not callable(orig_value):
+            return orig_value
+
         signature = inspect.signature(orig_value)
-        if signature.parameters.get('metadata', None) is None:
+        if signature.parameters.get("metadata", None) is None:
             return orig_value
 
-        print(key, '\033[34m', map_key, '\033[00m', signature, signature.parameters.get('metadata', None))
+        if False:
+            print(
+                key,
+                "\033[34m",
+                map_key,
+                "\033[00m",
+                signature,
+                signature.parameters.get("metadata", None),
+            )
+
+        if False:
+            stack = inspect.stack()
+            ends = stack[-50:-20]
+            for i, st in enumerate(ends):
+                print(i, st.filename, st.lineno)
+
+        print(
+            "\033[33mmonkey patching now\033[00m",
+            key,
+            "hex_orig",
+            hex_orig,
+            "hex_patched",
+            hex_patched,
+        )
         counters = dict(attempt=0)
+
         def patched_method(*aargs, **kkwargs):
-            counters['attempt'] += 1
-            metadata = kkwargs.get('metadata', None)
+            counters["attempt"] += 1
+            print("counters", counters)
+            metadata = kkwargs.get("metadata", None)
             if not metadata:
                 return orig_value(*aargs, **kkwargs)
 
             # 4. Find all the headers that match the target header key.
             all_metadata = []
             for mkey, value in metadata:
                 if mkey is REQ_ID_HEADER_KEY:
-                    attempt = counters['attempt']
+                    attempt = counters["attempt"]
                     if attempt > 1:
                         # 5. Increment the original_attempt with that of our re-invocation count.
                         splits = value.split(".")
-                        print('\033[34mkey', mkey, '\033[00m', splits)
-                        hdr_attempt_plus_reinvocation = (
-                            int(splits[-1]) + attempt
-                        )
+                        print("\033[34mkey", mkey, "\033[00m", splits)
+                        hdr_attempt_plus_reinvocation = int(splits[-1]) + attempt
                         splits[-1] = str(hdr_attempt_plus_reinvocation)
                         value = ".".join(splits)
 
                 all_metadata.append((mkey, value))
 
             kwargs["metadata"] = all_metadata
-            return orig_value(*aargs, **kkwargs)
+
+            try:
+                return orig_value(*aargs, **kkwargs)
+
+            except (InternalServerError, ServiceUnavailable) as exc:
+                print("caught this exception, incrementing", exc)
+                counters["attempt"] += 1
+                raise exc
 
         memoize_map[map_key] = patched_method
         return patched_method
 
-    setattr(obj, '__getattribute__', patched_getattribute)
+    hex_patched = hex(id(patched_getattribute))
+    setattr(obj, "__getattribute__", patched_getattribute)
 
 
 def foo(api):
     global patched
     global patched_mu
-    
+
     # For each method, add an _attempt value that'll then be
     # retrieved for each retry.
     # 1. Patch the __getattribute__ method to match items in our manifest.
@@ -878,20 +952,29 @@ def patched_getattribute(obj, key, *args, **kwargs):
         patched_key = hex(id(key)) + hex(id(obj))
         patched_mu.acquire()
         already_patched = patched.get(patched_key, None)
-        
+
         other_attempts = dict(attempts=0)
+
         # 3. Wrap the callable attribute and then capture its metadata keyed argument.
         def wrapped_attr(*args, **kwargs):
-            print("\033[31m", key, "attempt", other_attempts['attempts'], "\033[00m")
-            other_attempts['attempts'] += 1
+            print("\033[31m", key, "attempt", other_attempts["attempts"], "\033[00m")
+            other_attempts["attempts"] += 1
 
             metadata = kwargs.get("metadata", [])
             if not metadata:
                 # Increment the reinvocation count.
                 wrapped_attr._attempt += 1
                 return attr(*args, **kwargs)
 
-            print("\033[35mwrapped_attr", key, args, kwargs, 'attempt', wrapped_attr._attempt, "\033[00m")
+            print(
+                "\033[35mwrapped_attr",
+                key,
+                args,
+                kwargs,
+                "attempt",
+                wrapped_attr._attempt,
+                "\033[00m",
+            )
 
             # 4. Find all the headers that match the target header key.
             all_metadata = []
@@ -900,7 +983,7 @@ def wrapped_attr(*args, **kwargs):
                     if wrapped_attr._attempt > 0:
                         # 5. Increment the original_attempt with that of our re-invocation count.
                         splits = value.split(".")
-                        print('\033[34mkey', mkey, '\033[00m', splits)
+                        print("\033[34mkey", mkey, "\033[00m", splits)
                         hdr_attempt_plus_reinvocation = (
                             int(splits[-1]) + wrapped_attr._attempt
                         )
@@ -916,13 +999,13 @@ def wrapped_attr(*args, **kwargs):
 
         if already_patched:
             print("patched_key \033[32m", patched_key, key, "\033[00m", already_patched)
-            setattr(attr, 'patched', True)
+            setattr(attr, "patched", True)
             # Increment the reinvocation count.
             patched_mu.release()
             return already_patched
 
         patched[patched_key] = wrapped_attr
-        setattr(wrapped_attr, '_attempt', 0)
+        setattr(wrapped_attr, "_attempt", 0)
         patched_mu.release()
         return wrapped_attr
 

google/cloud/spanner_v1/batch.py

@@ -250,6 +250,7 @@ def commit(
             observability_options=observability_options,
             metadata=metadata,
         ), MetricsCapture():
+
             def wrapped_method(*args, **kwargs):
                 method = functools.partial(
                     api.commit,

google/cloud/spanner_v1/database.py

@@ -432,15 +432,6 @@ def logger(self):
 
     @property
     def spanner_api(self):
-        """Helper for session-related API calls."""
-        api = self.__generate_spanner_api()
-        if not api:
-            return api
-
-        monkey_patch(api)
-        return api
-
-    def __generate_spanner_api(self):
         """Helper for session-related API calls."""
         if self._spanner_api is None:
             client_info = self._instance._client._client_info