Add GAE Shutdown Handler

logging/google/cloud/logging/_environment_vars.py

@@ -0,0 +1,30 @@
+# Copyright 2017 Google Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Environment variables used by the package."""
+
+APPENGINE_FLEXIBLE_ENV_VM = 'GAE_APPENGINE_HOSTNAME'
+"""Environment variable set in App Engine when vm:true is set."""
+
+APPENGINE_FLEXIBLE_ENV_FLEX = 'GAE_INSTANCE'
+"""Environment variable set in App Engine when env:flex is set."""
+
+CONTAINER_ENGINE_ENV = 'KUBERNETES_SERVICE'
+"""Environment variable set in a Google Container Engine environment."""
+
+GAE_SERVICE = 'GAE_SERVICE'
+"""Environment variable in the App Engine Flex env that indicates version."""
+
+GAE_VERSION = 'GAE_VERSION'
+"""Environment variable in the App Engine Flex env that indicates version."""

logging/google/cloud/logging/_shutdown.py

@@ -0,0 +1,145 @@
+# Copyright 2017 Google Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Module for code related to shutdown logging."""
+
+import functools
+import os
+import signal
+import sys
+import traceback
+
+from google.cloud.logging._environment_vars import APPENGINE_FLEXIBLE_ENV_VM
+from google.cloud.logging._environment_vars import APPENGINE_FLEXIBLE_ENV_FLEX
+from google.cloud.logging._environment_vars import GAE_SERVICE
+from google.cloud.logging._environment_vars import GAE_VERSION
+from google.cloud import _helpers
+
+import six
+
+# Maximum size in bytes to send to Stackdriver Logging in one entry
+_MAX_PAYLOAD_SIZE = 100 * 1024
+
+_LOGGER_NAME_TMPL = 'projects/{}/logs/appengine.googleapis.com%2Fapp.shutdown'
+
+
+def _get_gae_instance():
+    """Returns the App Engine Flexible instance."""
+    return os.getenv(APPENGINE_FLEXIBLE_ENV_FLEX)
+
+
+def _get_gae_service():
+    """Returns the App Engine Flexible service."""
+    return os.getenv(GAE_SERVICE)
+
+
+def _get_gae_version():
+    """Returns the App Engine Flexible version."""
+    return os.getenv(GAE_VERSION)
+
+
+def _split_entry(payload):
+    """Splits payload into lists of maximum 100Kb.
+
+    Stackdriver Logging payloads are a maximum of 100Kb.
+    """
+    return [payload[i:i + _MAX_PAYLOAD_SIZE]
+            for i in six.moves.xrange(0, len(payload), _MAX_PAYLOAD_SIZE)]
+
+
+def _write_stacktrace_log(client, traces):
+    """Writes the trace logs to the appropriate GAE resource in Stackdriver.
+
+    :type client: :class:`google.cloud.logging.Client`
+    :param client: Stackdriver logging client.
+
+    :type traces: str
+    :param traces: String containing the stacktrace info to write to
+                   Stackdriver logging.
+    """
+    gae_version = _get_gae_version()
+    gae_service = _get_gae_service()
+    gae_instance = _get_gae_instance()
+
+    text_payload = '{}\nThread traces\n{}'.format(gae_instance, traces)
+
+    logger_name = _LOGGER_NAME_TMPL.format(client.project)
+
+    resource = {
+        'type': 'gae_app',
+        'labels': {
+            'project_id': client.project,
+            'version_id': gae_version,
+            'module_id': gae_service,
+        },
+    }
+
+    labels = {
+        'appengine.googleapis.com/version_id': gae_version,
+        'compute.googleapis.com/resource_type': 'instance',
+        'appengine.googleapis.com/instance_name': gae_instance,
+        'appengine.googleapis.com/module_id': gae_service,
+    }
+
+    split_payloads = _split_entry(_helpers._to_bytes(text_payload))
+    entries = [{'text_payload': payload} for payload in split_payloads]
+
+    client.logging_api.write_entries(
+        entries, logger_name=logger_name, resource=resource, labels=labels)
+
+
+def _is_on_appengine():
+    """Returns True if the environment is detected as App Engine flexible."""
+    return (os.getenv(APPENGINE_FLEXIBLE_ENV_VM) or os.getenv(
+        APPENGINE_FLEXIBLE_ENV_FLEX))
+
+
+def _report_stacktraces(
+        client, signal_, frame):  # pylint: disable=unused-argument
+    """Reports the stacktraces of all active threads to Stackdriver Logging.
+
+    :type client: `google.cloud.logging.Client`
+    :param client: Stackdriver logging client.
+
+    :type signal_: int
+    :param signal_: Signal number. Unused parameter since this function always
+                    expects SIGTERM, but the parameter is required to match
+                    the function signature for a signal handler.
+
+    :type frame: frame object
+    :param frame: The current stack frame. Unused parameter, this parameter
+                  is required to match the function signature for a signal
+                  handler.
+    """
+    traces = []
+    for thread_id, stack in sys._current_frames().items():
+        traces.append('\n# ThreadID: {}\n'.format(thread_id))
+        for filename, lineno, name, _ in traceback.extract_stack(stack):
+            traces.append('File: {}, line {}, in {}'.format(
+                filename, lineno, name))
+
+    _write_stacktrace_log(client, ''.join(traces))
+
+
+def setup_stacktrace_crash_report(client):
+    """Installs a SIGTERM handler to log stack traces to Stackdriver.
+
+    :type client: `google.cloud.logging.Client`
+    :param client: Stackdriver logging client.
+    """
+    if not _is_on_appengine():
+        raise RuntimeError('Shutdown reporting is only supported on App '
+                           'Engine flexible environment.')
+    signal.signal(
+        signal.SIGTERM, functools.partial(_report_stacktraces, client))

logging/google/cloud/logging/client.py

@@ -31,10 +31,14 @@
 
 from google.cloud.client import ClientWithProject
 from google.cloud.environment_vars import DISABLE_GRPC
+from google.cloud.logging._environment_vars import APPENGINE_FLEXIBLE_ENV_VM
+from google.cloud.logging._environment_vars import APPENGINE_FLEXIBLE_ENV_FLEX
+from google.cloud.logging._environment_vars import CONTAINER_ENGINE_ENV
 from google.cloud.logging._http import Connection
 from google.cloud.logging._http import _LoggingAPI as JSONLoggingAPI
 from google.cloud.logging._http import _MetricsAPI as JSONMetricsAPI
 from google.cloud.logging._http import _SinksAPI as JSONSinksAPI
+from google.cloud.logging._shutdown import setup_stacktrace_crash_report
 from google.cloud.logging.handlers import CloudLoggingHandler
 from google.cloud.logging.handlers import AppEngineHandler
 from google.cloud.logging.handlers import ContainerEngineHandler
@@ -49,15 +53,6 @@
 _DISABLE_GAX = os.getenv(DISABLE_GRPC, False)
 _USE_GAX = _HAVE_GAX and not _DISABLE_GAX
 
-_APPENGINE_FLEXIBLE_ENV_VM = 'GAE_APPENGINE_HOSTNAME'
-"""Environment variable set in App Engine when vm:true is set."""
-
-_APPENGINE_FLEXIBLE_ENV_FLEX = 'GAE_INSTANCE'
-"""Environment variable set in App Engine when env:flex is set."""
-
-_CONTAINER_ENGINE_ENV = 'KUBERNETES_SERVICE'
-"""Environment variable set in a Google Container Engine environment."""
-
 
 class Client(ClientWithProject):
     """Client to bundle configuration needed for API requests.
@@ -297,10 +292,10 @@ def get_default_handler(self):
         :rtype: :class:`logging.Handler`
         :returns: The default log handler based on the environment
         """
-        if (_APPENGINE_FLEXIBLE_ENV_VM in os.environ or
-                _APPENGINE_FLEXIBLE_ENV_FLEX in os.environ):
+        if (APPENGINE_FLEXIBLE_ENV_VM in os.environ or
+                APPENGINE_FLEXIBLE_ENV_FLEX in os.environ):
             return AppEngineHandler()
-        elif _CONTAINER_ENGINE_ENV in os.environ:
+        elif CONTAINER_ENGINE_ENV in os.environ:
             return ContainerEngineHandler()
         else:
             return CloudLoggingHandler(self)
@@ -327,3 +322,21 @@ def setup_logging(self, log_level=logging.INFO,
         handler = self.get_default_handler()
         setup_logging(handler, log_level=log_level,
                       excluded_loggers=excluded_loggers)
+
+    def enable_shutdown_logging(self, thread_dump=True):
+        """Enable shutdown report logging.
+
+        .. note::  This method is only supported on App Engine Flexible.
+
+        This method installs a SIGTERM handler that will report various
+        application metrics to Stackdriver Logging.
+
+        Currently the only supported option is thread stacktrace logging. This
+        will log all the stacktraces of all active threads.
+
+        :type thread_dump: bool
+        :param thread_dump: (Optional) When true, on SIGTERM the application
+                          will log the stacktrace of all running threads.
+        """
+        if thread_dump:
+            setup_stacktrace_crash_report(self)

logging/unit_tests/test__shutdown.py

@@ -0,0 +1,101 @@
+# Copyright 2017 Google Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import mock
+
+
+class Test_setup_stacktrace_crash_report(unittest.TestCase):
+
+    def _call_fut(self, request):
+        from google.cloud.logging._shutdown import (
+            setup_stacktrace_crash_report)
+        return setup_stacktrace_crash_report(request)
+
+    def test_outside_of_gae(self):
+        is_on_gae_patch = mock.patch(
+            'google.cloud.logging._shutdown._is_on_appengine',
+            return_value=False)
+        with is_on_gae_patch:
+            with self.assertRaises(RuntimeError):
+                self._call_fut(mock.sentinel.client)
+
+    def test_success(self):
+        import signal
+        from google.cloud.logging._environment_vars import (
+            APPENGINE_FLEXIBLE_ENV_VM)
+
+        signal_patch = mock.patch('google.cloud.logging._shutdown.signal')
+        environ_patch = mock.patch(
+            'os.environ',
+            new={APPENGINE_FLEXIBLE_ENV_VM: 'True'})
+
+        with environ_patch:
+            with signal_patch as signal_mock:
+                signal_mock.SIGTERM = signal.SIGTERM
+                self._call_fut(mock.sentinel.client)
+                self.assertTrue(signal_mock.signal.called)
+                positional = signal_mock.signal.call_args[0]
+                self.assertEqual(positional[0], signal.SIGTERM)
+                self.assertTrue(callable(type(positional[1])))
+
+
+class Test_write_stackrace_log(unittest.TestCase):
+
+    def _call_fut(self, client, traces):
+        from google.cloud.logging._shutdown import (_write_stacktrace_log)
+        return _write_stacktrace_log(client, traces)
+
+    def test_write_stracktrace_log(self):
+        from google.cloud.logging._gax import _LoggingAPI
+
+        mock_client = mock.Mock(spec=['project'])
+        mock_client.logging_api = mock.Mock(spec=_LoggingAPI)
+
+        trace = 'a simple stack trace'
+        environ_patch = mock.patch(
+            'os.environ', new={'GAE_INSTANCE': 'myversion'})
+        with environ_patch:
+            self._call_fut(mock_client, trace)
+            self.assertEqual(
+                len(mock_client.logging_api.write_entries.mock_calls), 1)
+            positional, _ = mock_client.logging_api.write_entries.call_args
+        expected_payload = 'myversion\nThread traces\n{}'.format(
+            trace).encode('utf-8')
+        self.assertEqual(positional[0], [{'text_payload': expected_payload}])
+
+
+class Test__report_stacktraces(unittest.TestCase):
+
+    def _call_fut(self, client, signal, frame):
+        from google.cloud.logging._shutdown import (
+            _report_stacktraces
+        )
+        return _report_stacktraces(client, signal, frame)
+
+    def test_success(self):
+        import re
+        patch = mock.patch(
+            'google.cloud.logging._shutdown._write_stacktrace_log')
+        with patch as write_log_mock:
+            self._call_fut(mock.sentinel.client, None, None)
+
+            traces = write_log_mock.call_args[0][1]
+
+        match = re.match(
+            '.*ThreadID: .*File:.*test__shutdown.*',
+            traces,
+            re.DOTALL)
+        self.assertIsNotNone(match)