recipe_modules/default_timeout/api.py

# Copyright 2022 The Pigweed Authors
#
# 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
#
#     https://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.
"""Ensure the given call completes before swarming kills the build."""

import contextlib
import dataclasses
from typing import Generator

from recipe_engine import recipe_api
from RECIPE_MODULES.fuchsia.utils import nice_duration


class TimeoutApi(recipe_api.RecipeApi):
    """Ensure the given call completes before swarming kills the build."""

    @contextlib.contextmanager
    def __call__(self, buffer_sec: float = 120) -> Generator[None, None, None]:
        current_time_sec: float = self.m.time.time()

        # Amount of time elapsed in the run.
        elapsed_time_sec: float = (
            current_time_sec - self.m.buildbucket.build.start_time.seconds
        )

        # Amount of time before build times out.
        time_remaining_sec: float = (
            self.m.buildbucket.build.execution_timeout.seconds
            - elapsed_time_sec
        )

        # Give a buffer before build times out and kill this step then. This
        # should give enough time to read any logfiles and maybe upload to
        # GCS before the build times out.
        timeout_sec: float = time_remaining_sec - buffer_sec

        # Most tests don't appropriately set up timing, and this could be
        # invoked at the very end of a build where it's about to timeout, so
        # make sure the timeout is always at least 10 seconds in the future.
        if timeout_sec < 10.0:
            timeout_sec = 10.0

        with self.m.step.nest(f'timeout {nice_duration(timeout_sec)}'):
            pass

        try:
            with self.m.time.timeout(timeout_sec):
                yield

        finally:
            pass