#!/usr/bin/env python
# Copyright 2014 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""
Estimates capacity needs for all builders of a given master.
"""

import argparse
import datetime
import json
import os
import subprocess
import sys
import tempfile
import urllib
import urllib2


BASE_DIR = os.path.dirname(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))))


def get_builds(mastername, buildername, days):
  results = []
  for day in days:
    cursor = ''
    while True:
      response = urllib2.urlopen(
          'https://chrome-build-extract.appspot.com/get_builds'
          '?master=%s&builder=%s&day=%s&cursor=%s' % (
              urllib.quote(mastername),
              urllib.quote(buildername),
              urllib.quote(str(day)),
              urllib.quote(cursor)))
      data = json.load(response)
      results.extend(data['builds'])
      cursor = data['cursor']
      if not cursor:
        break

  return results


def estimate_capacity(builds):
  hourly_buckets = {}
  daily_buckets = {}
  for build in builds:
    build_time = build['times'][1] - build['times'][0]

    # TODO(phajdan.jr): use timestamp of the build request being created
    # Using timestamp of build starting doesn't take pending queues
    # into account, and so is less accurate.
    timestamp = datetime.datetime.utcfromtimestamp(build['times'][0])

    hourly_bucket = timestamp.strftime('%Y-%m-%d-%H')
    hourly_buckets.setdefault(hourly_bucket, []).append(build_time)

    daily_bucket = timestamp.strftime('%Y-%m-%d')
    daily_buckets.setdefault(daily_bucket, []).append(build_time)

  def min_bots(buckets, resolution_s):
    result = 0
    for times in buckets.itervalues():
      total_time = sum(times)
      bots = total_time / resolution_s
      result = max(result, bots)
    return result

  return {
    'hourly_bots': min_bots(hourly_buckets, 3600),
    'daily_bots': min_bots(daily_buckets, 3600 * 24),
  }


def main(argv):
  parser = argparse.ArgumentParser()
  parser.add_argument('master')
  parser.add_argument('--days', type=int, default=14)

  args = parser.parse_args(argv)

  with tempfile.NamedTemporaryFile() as f:
    subprocess.check_call([
        os.path.join(BASE_DIR, 'scripts', 'tools', 'runit.py'),
        os.path.join(BASE_DIR, 'scripts', 'tools', 'dump_master_cfg.py'),
        'masters/%s' % args.master,
        f.name])
    master_config = json.load(f)

  builder_pools = []
  for builder in master_config['builders']:
    slave_set = set(builder['slavenames'])
    for pool in builder_pools:
      if pool['slave_set'] == slave_set:
        pool['builders'].append(builder['name'])
        break
    else:
      builder_pools.append({
        'slave_set': slave_set,
        'builders': [builder['name']],
      })

  days = []
  for i in range(args.days):
    days.append(datetime.date.today() - datetime.timedelta(days=(i + 1)))
  for index, pool in enumerate(builder_pools):
    print 'Pool #%d:' % (index + 1)
    pool_capacity = {
      'hourly_bots': 0.0,
      'daily_bots': 0.0,
    }
    # TODO(phajdan.jr): use multiprocessing pool to speed this up.
    for builder in pool['builders']:
      builds = get_builds(args.master.replace('master.', ''), builder, days)
      capacity = estimate_capacity(builds)
      for key in ('hourly_bots', 'daily_bots'):
        pool_capacity[key] += capacity[key]
      print '\t%-45s %5.1f %5.1f' % (
          builder,
          capacity['daily_bots'],
          capacity['hourly_bots'])
    print '\t%-45s %5.1f %5.1f %5.1f' % (
        'total',
        pool_capacity['daily_bots'],
        pool_capacity['hourly_bots'],
        len(pool['slave_set']))

  return 0


if __name__ == '__main__':
  sys.exit(main(sys.argv[1:]))

@@ -1211,190 +1211,272 @@
     
-# TODO(phajdan.jr): use timestamp of the build request being created%0A    # Using timestamp of build starting doesn't take pending queues%0A    # into account, and so is less accurate.%0A
+changes = build%5B'sourceStamp'%5D%5B'changes'%5D%0A    if changes:%0A      assert len(changes) == 1%0A      timestamp = datetime.datetime.utcfromtimestamp(changes%5B0%5D%5B'when'%5D)%0A    else:%0A      # Fallback for builds that don't have blamelist/source stamp,%0A      # e.g. win_pgo.%0A  
     

Improve accuracy of estimate_capacity.py