1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
3# Copyright (c) 2021 Huawei Device Co., Ltd.
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16import shutil
17import os
18import hashlib
19import json
20import http.client as client
21from . import build_utils
22
23
24class Storage():
25    def __init__(self):
26        pass
27
28    @classmethod
29    def retrieve_object(cls, cache_artifact, obj):
30        possible_dir_cache_artifact = '{}.directory'.format(cache_artifact)
31
32        if os.path.exists(cache_artifact):
33            os.makedirs(os.path.dirname(obj), exist_ok=True)
34            shutil.copyfile(cache_artifact, obj)
35            os.utime(cache_artifact)
36            if pycache_debug_enable:
37                print('Retrieve {} from cache'.format(obj))
38        elif os.path.exists(possible_dir_cache_artifact):
39            # Extract zip archive if it's cache artifact for directory.
40            os.makedirs(obj, exist_ok=True)
41            build_utils.extract_all(possible_dir_cache_artifact,
42                                   obj,
43                                   no_clobber=False)
44            os.utime(possible_dir_cache_artifact)
45            if pycache_debug_enable:
46                print('Extract {} from cache'.format(obj))
47        else:
48            if pycache_debug_enable:
49                print('Failed to retrieve {} from cache'.format(obj))
50            return 0
51        return 1
52
53    @classmethod
54    def add_object(cls, cache_artifact, obj):
55        cache_dir = os.path.dirname(cache_artifact)
56        os.makedirs(cache_dir, exist_ok=True)
57
58        if not os.path.exists(obj):
59            return
60        # If path is directory, store an zip archive.
61        if os.path.isdir(obj):
62            dir_cache_artifact = '{}.directory'.format(cache_artifact)
63            build_utils.zip_dir(dir_cache_artifact, obj)
64            if pycache_debug_enable:
65                print("archive {} to {}".format(obj, dir_cache_artifact))
66        else:
67            shutil.copyfile(obj, cache_artifact)
68            if pycache_debug_enable:
69                print("copying {} to {}".format(obj, cache_artifact))
70
71
72class PyCache():
73    def __init__(self, cache_dir=None):
74        cache_dir = os.environ.get('PYCACHE_DIR')
75        if cache_dir:
76            self.pycache_dir = cache_dir
77        else:
78            raise Exception('Error: failed to get PYCACHE_DIR')
79        self.storage = Storage()
80
81    @classmethod
82    def cache_key(cls, path):
83        sha256 = hashlib.sha256()
84        sha256.update(path.encode())
85        return sha256.hexdigest()
86
87    def retrieve(self, output_paths, prefix=''):
88        for path in output_paths:
89            _, cache_artifact = self.descend_directory('{}{}'.format(
90                prefix, path))
91            result = self.storage.retrieve_object(cache_artifact, path)
92            if not result:
93                return result
94
95        try:
96            self.report_cache_stat('cache_hit')
97        except:  # noqa: E722 pylint: disable=bare-except
98            pass
99        return 1
100
101    def save(self, output_paths, prefix=''):
102        for path in output_paths:
103            _, cache_artifact = self.descend_directory('{}{}'.format(
104                prefix, path))
105            self.storage.add_object(cache_artifact, path)
106
107    def report_cache_stat(self, hit_or_miss):
108        pyd_server, pyd_port = self.get_pyd()
109        conn = client.HTTPConnection(pyd_server, pyd_port)
110        conn.request(hit_or_miss, '/')
111        conn.close()
112
113    def get_pyd(self):
114        daemon_config_file = '{}/.config'.format(self.pycache_dir)
115        if not os.path.exists(daemon_config_file):
116            raise Exception('Warning: no pycache daemon process exists.')
117        with open(daemon_config_file, 'r') as jsonfile:
118            data = json.load(jsonfile)
119            return data.get('host'), data.get('port')
120
121    def descend_directory(self, path):
122        digest = self.cache_key(path)
123        cache_dir = os.path.join(self.pycache_dir, digest[:2])
124        return cache_dir, os.path.join(cache_dir, digest[2:])
125
126    # Manifest file to record inputs/outputs/commands.
127    def get_manifest_path(self, path):
128        manifest_dir, manifest_file = self.descend_directory(path)
129        os.makedirs(manifest_dir, exist_ok=True)
130        return manifest_file
131
132
133pycache_enabled = (os.environ.get('PYCACHE_DIR') is not None)
134pycache_debug_enable = int(os.environ.get('PRINT_BUILD_EXPLANATIONS', 0))
135if pycache_enabled:
136    pycache = PyCache()
137else:
138    pycache = None  # pylint: disable=invalid-name
139