Fix trace client memory leak (googleapis#6856)

c24t · tseaver · commit c0abaa4d6860 · 2018-12-07T13:05:48.000-05:00
diff --git a/trace/google/cloud/trace/client.py b/trace/google/cloud/trace/client.py
@@ -51,7 +51,8 @@ def trace_api(self):
         https://cloud.google.com/trace/docs/reference/v2/rpc/google.devtools.
         cloudtrace.v2
         """
-        self._trace_api = make_trace_api(self)
+        if self._trace_api is None:
+            self._trace_api = make_trace_api(self)
         return self._trace_api
 
     def batch_write_spans(self, name, spans, retry=None, timeout=None):
diff --git a/trace/google/cloud/trace/v1/client.py b/trace/google/cloud/trace/v1/client.py
@@ -53,7 +53,8 @@ def trace_api(self):
         https://cloud.google.com/trace/docs/reference/v1/rpc/google.devtools.
         cloudtrace.v1
         """
-        self._trace_api = make_trace_api(self)
+        if self._trace_api is None:
+            self._trace_api = make_trace_api(self)
         return self._trace_api
 
     def patch_traces(self, traces, project_id=None):
diff --git a/trace/tests/unit/v1/test_client_v1.py b/trace/tests/unit/v1/test_client_v1.py
@@ -60,6 +60,12 @@ def make_api(client_obj):
         self.assertIs(api, api_obj)
         self.assertEqual(clients, [client])
 
+    def test_trace_api_existing(self):
+        """Check that the client caches _trace_api."""
+        client = self._make_one(project=self.project, credentials=_make_credentials())
+        client._trace_api = mock.sentinel.trace_api
+        self.assertIs(client.trace_api, mock.sentinel.trace_api)
+
     def test_patch_traces_default(self):
         from google.cloud.trace.v1._gapic import _TraceAPI
 
diff --git a/trace/tests/unit/v2/test_client_v2.py b/trace/tests/unit/v2/test_client_v2.py
@@ -60,6 +60,12 @@ def make_api(client_obj):
         self.assertIs(api, api_obj)
         self.assertEqual(clients, [client])
 
+    def test_trace_api_existing(self):
+        """Check that the client caches _trace_api."""
+        client = self._make_one(project=self.project, credentials=_make_credentials())
+        client._trace_api = mock.sentinel.trace_api
+        self.assertIs(client.trace_api, mock.sentinel.trace_api)
+
     def test_batch_write_spans(self):
         from google.cloud.trace._gapic import _TraceAPI
 
