buildkite/gcloud_utils.py

#!/usr/bin/env python3
#
# Copyright 2018 The Bazel Authors. All rights reserved.
#
# 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 gcloud
import json
import subprocess
import time
import threading
import re


def wait_for_instance(instance_name, zone, status):
    while True:
        result = gcloud.describe_instance(instance_name, zone=zone, format='json')
        current_status = json.loads(result.stdout)['status']
        if current_status == status:
            gcloud.debug('wait_for_instance: {}/{} arrived at status {}'.format(zone, instance_name, status))
            break
        else:
            gcloud.debug('wait_for_instance: Waiting for {}/{} to go from status {} to status {}'.format(
                zone, instance_name, current_status, status))
        time.sleep(5)


def prettify_logs(instance_name, log, with_prefix=True):
    for line in log.splitlines():
        # Skip empty lines.
        if not line:
            continue

        # Filter for log lines printed by our startup script, ignore the rest.
        # Then drop the common prefix to make the output easier to read.
        # For unknown platforms, we just take every line unmodified.
        if 'ubuntu' in instance_name:
            match = re.match(r'.*INFO startup-script: (.*)', line)
            if not match:
                continue
            line = match.group(1)
        elif 'windows' in instance_name:
            match = re.match(r'.*windows-startup-script-ps1: (.*)', line)
            if not match:
                continue
            line = match.group(1)

        if with_prefix:
            yield "{}: {}".format(instance_name, line)
        else:
            yield line


def print_pretty_logs(instance_name, log):
    lines = ('\n'.join(prettify_logs(instance_name, log))).strip()
    if lines:
        with gcloud.PRINT_LOCK:
            print(lines)


def tail_serial_console(instance_name, zone, start=None, until=None):
    next_start = start if start else '0'
    while True:
        try:
            result = gcloud.get_serial_port_output(instance_name, zone=zone, start=next_start)
        except subprocess.CalledProcessError as e:
            if 'Could not fetch serial port output: TIMEOUT' in e.stderr:
                gcloud.debug('tail_serial_console: Retrying after TIMEOUT')
                continue
            gcloud.debug('tail_serial_console: Done, because got exception: {}'.format(e))
            if e.stdout:
                gcloud.debug('stdout: ' + e.stdout)
            if e.stderr:
                gcloud.debug('stderr: ' + e.stderr)
            break
        print_pretty_logs(instance_name, result.stdout)
        next_start = re.search(r'--start=(\d*)', result.stderr).group(1)
        if until and until in result.stdout:
            gcloud.debug('tail_serial_console: Done, because found string "{}"'.format(until))
            break
        time.sleep(5)
    return next_start