xref: /aosp_15_r20/tools/asuite/atest/rollout_control.py (revision c2e18aaa1096c836b086f94603d04f4eb9cf37f5) 
1#!/usr/bin/env python3
2# Copyright 2024, The Android Open Source Project
3#
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
16"""Rollout control for Atest features."""
17
18import functools
19import getpass
20import hashlib
21import importlib.resources
22import logging
23import os
24from atest import atest_enum
25from atest import atest_utils
26from atest.metrics import metrics
27
28
29@functools.cache
30def _get_project_owners() -> list[str]:
31  """Returns the owners of the feature."""
32  owners = []
33  try:
34    with importlib.resources.as_file(
35        importlib.resources.files('atest').joinpath('OWNERS')
36    ) as version_file_path:
37      owners.extend(version_file_path.read_text(encoding='utf-8').splitlines())
38  except (ModuleNotFoundError, FileNotFoundError) as e:
39    logging.error(e)
40  try:
41    with importlib.resources.as_file(
42        importlib.resources.files('atest').joinpath('OWNERS_ADTE_TEAM')
43    ) as version_file_path:
44      owners.extend(version_file_path.read_text(encoding='utf-8').splitlines())
45  except (ModuleNotFoundError, FileNotFoundError) as e:
46    logging.error(e)
47  return [line.split('@')[0] for line in owners if '@google.com' in line]
48
49
50class RolloutControlledFeature:
51  """Base class for Atest features under rollout control."""
52
53  def __init__(
54      self,
55      name: str,
56      rollout_percentage: float,
57      env_control_flag: str,
58      feature_id: int = None,
59      owners: list[str] | None = None,
60      print_message: str | None = None,
61  ):
62    """Initializes the object.
63
64    Args:
65        name: The name of the feature.
66        rollout_percentage: The percentage of users to enable the feature for.
67          The value should be in [0, 100].
68        env_control_flag: The environment variable name to override the feature
69          enablement. When set, 'true' or '1' means enable, other values means
70          disable.
71        feature_id: The ID of the feature that is controlled by rollout control
72          for metric collection purpose. Must be a positive integer.
73        owners: The owners of the feature. If not provided, the owners of the
74          feature will be read from OWNERS file.
75        print_message: The message to print to the console when the feature is
76          enabled for the user.
77    """
78    if rollout_percentage < 0 or rollout_percentage > 100:
79      raise ValueError(
80          'Rollout percentage must be in [0, 100]. Got %s instead.'
81          % rollout_percentage
82      )
83    if feature_id is not None and feature_id <= 0:
84      raise ValueError(
85          'Feature ID must be a positive integer. Got %s instead.' % feature_id
86      )
87    if owners is None:
88      owners = _get_project_owners()
89    self._name = name
90    self._rollout_percentage = rollout_percentage
91    self._env_control_flag = env_control_flag
92    self._feature_id = feature_id
93    self._owners = owners
94    self._print_message = print_message
95
96  def _check_env_control_flag(self) -> bool | None:
97    """Checks the environment variable to override the feature enablement.
98
99    Returns:
100        True if the feature is enabled, False if disabled, None if not set.
101    """
102    if self._env_control_flag not in os.environ:
103      return None
104    return os.environ[self._env_control_flag] in ('TRUE', 'True', 'true', '1')
105
106  def _is_enabled_for_user(self, username: str | None) -> bool:
107    """Checks whether the feature is enabled for the user.
108
109    Args:
110        username: The username to check the feature enablement for. If not
111          provided, the current user's username will be used.
112
113    Returns:
114        True if the feature is enabled for the user, False otherwise.
115    """
116    if self._rollout_percentage == 100:
117      return True
118
119    if username is None:
120      username = getpass.getuser()
121
122    if not username:
123      logging.debug(
124          'Unable to determine the username. Disabling the feature %s.',
125          self._name,
126      )
127      return False
128
129    if username in self._owners:
130      return True
131
132    hash_object = hashlib.sha256()
133    hash_object.update((username + ' ' + self._name).encode('utf-8'))
134    return int(hash_object.hexdigest(), 16) % 100 < self._rollout_percentage
135
136  @functools.cache
137  def is_enabled(self, username: str | None = None) -> bool:
138    """Checks whether the current feature is enabled for the user.
139
140    Args:
141        username: The username to check the feature enablement for. If not
142          provided, the current user's username will be used.
143
144    Returns:
145        True if the feature is enabled for the user, False otherwise.
146    """
147    override_flag_value = self._check_env_control_flag()
148    if override_flag_value is not None:
149      logging.debug(
150          'Feature %s is %s by env variable %s.',
151          self._name,
152          'enabled' if override_flag_value else 'disabled',
153          self._env_control_flag,
154      )
155      if self._feature_id:
156        metrics.LocalDetectEvent(
157            detect_type=atest_enum.DetectType.ROLLOUT_CONTROLLED_FEATURE_ID_OVERRIDE,
158            result=self._feature_id
159            if override_flag_value
160            else -self._feature_id,
161        )
162      return override_flag_value
163
164    is_enabled = self._is_enabled_for_user(username)
165
166    logging.debug(
167        'Feature %s is %s for user %s.',
168        self._name,
169        'enabled' if is_enabled else 'disabled',
170        username,
171    )
172
173    if self._feature_id:
174      metrics.LocalDetectEvent(
175          detect_type=atest_enum.DetectType.ROLLOUT_CONTROLLED_FEATURE_ID,
176          result=self._feature_id if is_enabled else -self._feature_id,
177      )
178
179    if is_enabled and self._print_message:
180      print(atest_utils.mark_magenta(self._print_message))
181
182    return is_enabled
183
184
185deprecate_bazel_mode = RolloutControlledFeature(
186    name='Deprecate Bazel Mode',
187    rollout_percentage=60,
188    env_control_flag='DEPRECATE_BAZEL_MODE',
189    feature_id=1,
190)
191
192rolling_tf_subprocess_output = RolloutControlledFeature(
193    name='Rolling TradeFed subprocess output',
194    rollout_percentage=100,
195    env_control_flag='ROLLING_TF_SUBPROCESS_OUTPUT',
196    feature_id=2,
197    print_message=(
198        'You are one of the first users receiving the "Rolling subprocess'
199        ' output" feature. If you are happy with it, please +1 on'
200        ' http://b/380460196.'
201    ),
202)
203
204tf_preparer_incremental_setup = RolloutControlledFeature(
205    name='TradeFed preparer incremental setup',
206    rollout_percentage=0,
207    env_control_flag='TF_PREPARER_INCREMENTAL_SETUP',
208    feature_id=3,
209    print_message=(
210        'You are one of the first users selected to receive the "Incremental'
211        ' setup for TradeFed preparers" feature. If you are happy with it,'
212        ' please +1 on http://b/381900378. If you experienced any issues,'
213        ' please comment on the same bug.'
214    ),
215)
216