1# This Source Code Form is subject to the terms of the Mozilla Public
2# License, v. 2.0. If a copy of the MPL was not distributed with this
3# file, You can obtain one at http://mozilla.org/MPL/2.0/.
4
5from __future__ import absolute_import, print_function, unicode_literals
6
7import logging
8import json
9import os
10import urllib2
11
12from . import base
13from taskgraph.util.docker import (
14    docker_image,
15    generate_context_hash,
16    INDEX_PREFIX,
17)
18from taskgraph.util.templates import Templates
19
20logger = logging.getLogger(__name__)
21GECKO = os.path.realpath(os.path.join(__file__, '..', '..', '..', '..'))
22ARTIFACT_URL = 'https://queue.taskcluster.net/v1/task/{}/artifacts/{}'
23INDEX_URL = 'https://index.taskcluster.net/v1/task/{}'
24
25
26class DockerImageTask(base.Task):
27
28    def __init__(self, *args, **kwargs):
29        self.index_paths = kwargs.pop('index_paths')
30        super(DockerImageTask, self).__init__(*args, **kwargs)
31
32    def __eq__(self, other):
33        return super(DockerImageTask, self).__eq__(other) and \
34               self.index_paths == other.index_paths
35
36    @classmethod
37    def load_tasks(cls, kind, path, config, params, loaded_tasks):
38        parameters = {
39            'pushlog_id': params.get('pushlog_id', 0),
40            'pushdate': params['moz_build_date'],
41            'pushtime': params['moz_build_date'][8:],
42            'year': params['moz_build_date'][0:4],
43            'month': params['moz_build_date'][4:6],
44            'day': params['moz_build_date'][6:8],
45            'project': params['project'],
46            'docker_image': docker_image,
47            'base_repository': params['base_repository'] or params['head_repository'],
48            'head_repository': params['head_repository'],
49            'head_ref': params['head_ref'] or params['head_rev'],
50            'head_rev': params['head_rev'],
51            'owner': params['owner'],
52            'level': params['level'],
53            'source': '{repo}file/{rev}/taskcluster/ci/docker-image/image.yml'
54                      .format(repo=params['head_repository'], rev=params['head_rev']),
55            'index_image_prefix': INDEX_PREFIX,
56            'artifact_path': 'public/image.tar.zst',
57        }
58
59        tasks = []
60        templates = Templates(path)
61        for image_name, image_symbol in config['images'].iteritems():
62            context_path = os.path.join('testing', 'docker', image_name)
63            context_hash = generate_context_hash(GECKO, context_path, image_name)
64
65            image_parameters = dict(parameters)
66            image_parameters['image_name'] = image_name
67            image_parameters['context_hash'] = context_hash
68
69            image_task = templates.load('image.yml', image_parameters)
70            attributes = {'image_name': image_name}
71
72            # unique symbol for different docker image
73            if 'extra' in image_task['task']:
74                image_task['task']['extra']['treeherder']['symbol'] = image_symbol
75
76            # As an optimization, if the context hash exists for a high level, that image
77            # task ID will be used.  The reasoning behind this is that eventually everything ends
78            # up on level 3 at some point if most tasks use this as a common image
79            # for a given context hash, a worker within Taskcluster does not need to contain
80            # the same image per branch.
81            index_paths = ['{}.level-{}.{}.hash.{}'.format(
82                                INDEX_PREFIX, level, image_name, context_hash)
83                           for level in range(int(params['level']), 4)]
84
85            tasks.append(cls(kind, 'build-docker-image-' + image_name,
86                             task=image_task['task'], attributes=attributes,
87                             index_paths=index_paths))
88
89        return tasks
90
91    def get_dependencies(self, taskgraph):
92        return []
93
94    def optimize(self, params):
95        for index_path in self.index_paths:
96            try:
97                url = INDEX_URL.format(index_path)
98                existing_task = json.load(urllib2.urlopen(url))
99                # Only return the task ID if the artifact exists for the indexed
100                # task.  Otherwise, continue on looking at each of the branches.  Method
101                # continues trying other branches in case mozilla-central has an expired
102                # artifact, but 'project' might not. Only return no task ID if all
103                # branches have been tried
104                request = urllib2.Request(
105                    ARTIFACT_URL.format(existing_task['taskId'], 'public/image.tar.zst'))
106                request.get_method = lambda: 'HEAD'
107                urllib2.urlopen(request)
108
109                # HEAD success on the artifact is enough
110                return True, existing_task['taskId']
111            except urllib2.HTTPError:
112                pass
113
114        return False, None
115
116    @classmethod
117    def from_json(cls, task_dict):
118        # Generating index_paths for optimization
119        imgMeta = task_dict['task']['extra']['imageMeta']
120        image_name = imgMeta['imageName']
121        context_hash = imgMeta['contextHash']
122        index_paths = ['{}.level-{}.{}.hash.{}'.format(
123                            INDEX_PREFIX, level, image_name, context_hash)
124                       for level in range(int(imgMeta['level']), 4)]
125        docker_image_task = cls(kind='docker-image',
126                                label=task_dict['label'],
127                                attributes=task_dict['attributes'],
128                                task=task_dict['task'],
129                                index_paths=index_paths)
130        return docker_image_task
131