github.com/abayer/test-infra@v0.0.5/queue_health/poll/poller.py (about)

     1  #!/usr/bin/env python
     2  
     3  # Copyright 2016 The Kubernetes Authors.
     4  #
     5  # Licensed under the Apache License, Version 2.0 (the "License");
     6  # you may not use this file except in compliance with the License.
     7  # You may obtain a copy of the License at
     8  #
     9  #     http://www.apache.org/licenses/LICENSE-2.0
    10  #
    11  # Unless required by applicable law or agreed to in writing, software
    12  # distributed under the License is distributed on an "AS IS" BASIS,
    13  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  # See the License for the specific language governing permissions and
    15  # limitations under the License.
    16  
    17  import cStringIO
    18  import datetime
    19  import pprint
    20  import subprocess
    21  import sys
    22  import time
    23  import traceback
    24  
    25  import requests
    26  
    27  def get_submit_queue_json(path):
    28      for count in range(3):
    29          uri = 'https://submit-queue.k8s.io/%s' % path
    30          print >>sys.stderr, 'GET %s' % uri
    31          resp = requests.get(uri, allow_redirects=True)
    32          if resp.ok:
    33              break
    34          time.sleep(2**count)
    35      resp.raise_for_status()
    36      return resp.json()
    37  
    38  
    39  def is_blocked():
    40      json = get_submit_queue_json('health')
    41      return json['MergePossibleNow'] != True
    42  
    43  
    44  def get_stats():
    45      stats = get_submit_queue_json('sq-stats')
    46      return stats['Initialized'] is True, stats['MergesSinceRestart']
    47  
    48  
    49  def poll():
    50      prs = get_submit_queue_json('prs')
    51      e2e = get_submit_queue_json('github-e2e-queue')
    52      online, merge_count = get_stats()
    53      return (
    54          online,  # Is mergebot initialized?
    55          len(prs['PRStatus']),  # number of open PRs
    56          len(e2e['E2EQueue']),  # number of items in the e2e queue
    57          len(e2e['E2ERunning']),  # Worthless: number of keys in this dict.
    58          is_blocked(),  # Whether we can merge
    59          merge_count,  # Number of merges the bot has done
    60      )
    61  
    62  
    63  def load_stats(uri):
    64      while True:
    65          try:
    66              return subprocess.check_output(['gsutil', '-q', 'cat', uri])
    67          except subprocess.CalledProcessError:
    68              traceback.print_exc()
    69              time.sleep(5)
    70  
    71  
    72  def save_stats(uri, buf):
    73      proc = subprocess.Popen(
    74          # TODO(fejta): add -Z if this gets resolved:
    75          # https://github.com/GoogleCloudPlatform/gsutil/issues/364
    76          ['gsutil', '-q', '-h', 'Content-Type:text/plain',
    77           'cp', '-a', 'public-read', '-', uri],
    78          stdin=subprocess.PIPE)
    79      proc.communicate(buf.getvalue())
    80      code = proc.wait()
    81      if code:
    82          print >>sys.stderr, 'Failed to copy stats to %s: %d' % (uri, code)
    83  
    84  
    85  def poll_forever(uri, service_account=None):
    86      if service_account:
    87          print >>sys.stderr, 'Activating service account using: %s' % service_account
    88          subprocess.check_call(
    89              ['gcloud', 'auth', 'activate-service-account', '--key-file=%s' % service_account])
    90      print >>sys.stderr, 'Loading historical stats from %s...' % uri
    91      buf = cStringIO.StringIO()
    92      buf.write(load_stats(uri))
    93      secs = 60
    94  
    95      while True:
    96          try:
    97              print >>sys.stderr, 'Waiting %ds...' % secs
    98              time.sleep(secs)
    99              now = datetime.datetime.now()
   100              print >>sys.stderr, 'Polling current status...'
   101              online, prs, queue, running, blocked, merge_count = False, 0, 0, 0, False, 0
   102              try:
   103                  online, prs, queue, running, blocked, merge_count = poll()
   104              except KeyboardInterrupt:
   105                  raise
   106              except (KeyError, IOError):
   107                  traceback.print_exc()
   108                  continue
   109  
   110              data = '{} {} {} {} {} {} {}\n'.format(
   111                  now, online, prs, queue, running, blocked, merge_count)
   112              print >>sys.stderr, 'Appending to history: %s' % data
   113              buf.write(data)
   114  
   115              print >>sys.stderr, 'Saving historical stats to %s...' % uri
   116              save_stats(uri, buf)
   117          except KeyboardInterrupt:
   118              break
   119  
   120  
   121  if __name__ == '__main__':
   122      # log all arguments.
   123      PP = pprint.PrettyPrinter(stream=sys.stderr)
   124      PP.pprint(sys.argv)
   125  
   126      poll_forever(*sys.argv[1:])  # pylint: disable=no-value-for-parameter