| # Copyright (C) 2022 The Android Open Source Project |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| |
| import dataclasses as dc |
| import enum |
| from typing import BinaryIO, Dict, Generator, List, Type, Union |
| from typing import Generic, Tuple, TypeVar, get_type_hints |
| |
| from perfetto.trace_uri_resolver import util |
| |
| TraceUri = str |
| TraceGenerator = Generator[bytes, None, None] |
| TraceContent = Union[BinaryIO, TraceGenerator] |
| _T = TypeVar('_T') |
| |
| |
| @dc.dataclass |
| class ConstraintClass(Generic[_T]): |
| |
| class Op(enum.Enum): |
| EQ = '=' |
| NE = '!=' |
| LE = '<=' |
| GE = '>=' |
| GT = '>' |
| LT = '<' |
| |
| def __str__(self): |
| return self.value |
| |
| value: _T |
| op: Op = Op.EQ |
| |
| |
| Constraint = Union[_T, ConstraintClass[_T]] |
| ConstraintWithList = Union[Constraint[_T], Constraint[List[_T]]] |
| |
| |
| class TraceUriResolver: |
| """"Resolves a trace URI (e.g. 'ants:trace_id=1234') into a list of traces. |
| |
| This class can be subclassed to provide a pluggable mechanism for looking |
| up traces using URI strings. |
| |
| For example: |
| class CustomTraceResolver(TraceUriResolver): |
| PREFIX = 'custom' |
| |
| def __init__(self, build_branch: List[str] = None, id: str = None): |
| self.build_branch = build_branch |
| self.id = id |
| self.db = init_db() |
| |
| def resolve(self): |
| traces = self.db.lookup( |
| id=self.id, build_branch=self.build_branch)['path'] |
| return [ |
| TraceUriResolver.Result( |
| trace=t['path'], |
| args={'iteration': t['iteration'], 'device': t['device']} |
| ) |
| for t in traces |
| ] |
| |
| Trace resolvers can be passed to trace processor directly: |
| with TraceProcessor(CustomTraceResolver(id='abcdefg')) as tp: |
| tp.query('select * from slice') |
| |
| Alternatively, a trace addesses can be passed: |
| config = TraceProcessorConfig( |
| resolver_registry=ResolverRegistry(resolvers=[CustomTraceResolver]) |
| ) |
| with TraceProcessor('custom:id=abcdefg', config=config) as tp: |
| tp.query('select * from slice') |
| """ |
| |
| # Subclasses should set PREFIX to match the trace address prefix they |
| # want to handle. |
| PREFIX: str = None |
| |
| @dc.dataclass |
| class Result: |
| # TraceUri is present here because it allows recursive lookups (i.e. |
| # a resolver which returns a path to a trace). |
| trace: Union[TraceUri, TraceContent] |
| |
| # metadata allows additional key-value pairs to be provided which are |
| # associated for trace. For example, test names and iteration numbers |
| # could be provivded for traces originating from lab tests. |
| metadata: Dict[str, str] |
| |
| def __init__(self, |
| trace: Union[TraceUri, TraceContent], |
| metadata: Dict[str, str] = dict()): |
| self.trace = trace |
| self.metadata = metadata |
| |
| def resolve(self) -> List['TraceUriResolver.Result']: |
| """Resolves a list of traces. |
| |
| Subclasses should implement this method and resolve the parameters |
| specified in the constructor to a list of traces. |
| """ |
| raise Exception('resolve is unimplemented for this resolver') |
| |
| @classmethod |
| def from_trace_uri(cls: Type['TraceUriResolver'], |
| uri: TraceUri) -> 'TraceUriResolver': |
| """Creates a resolver from a URI. |
| |
| URIs have the form: |
| android_ci:day=2021-01-01;devices=blueline,crosshatch;key>=value |
| |
| This is converted to a dictionary of the form: |
| {'day': '2021-01-01', 'id': ['blueline', 'crosshatch'], |
| 'key': ConstraintClass('value', Op.GE)} |
| |
| and passed as kwargs to the constructor of the trace resolver (see class |
| documentation for info). |
| |
| Generally, sublcasses should not override this method as the standard |
| trace address format should work for most usecases. Instead, simply |
| define your constructor with the parameters you expect to see in the |
| trace address. |
| """ |
| return cls(**_args_dict_from_uri(uri, get_type_hints(cls.__init__))) |
| |
| |
| def _read_op(arg_str: str, op_start_ind: int) -> ConstraintClass.Op: |
| """Parse operator from string. |
| |
| Given string and an expected start index for operator it returns Op object or |
| raises error if operator was not found. |
| |
| For example: |
| _read_op('a>4', 1) returns Op.GE |
| _read_op('a>4', 0) raises ValueError |
| _read_op('a>4', 3) raises ValueError |
| """ |
| first = arg_str[op_start_ind] if op_start_ind < len(arg_str) else None |
| second = arg_str[op_start_ind + |
| 1] if op_start_ind + 1 < len(arg_str) else None |
| Op = ConstraintClass.Op |
| if first == '>': |
| return Op.GE if second == '=' else Op.GT |
| elif first == '<': |
| return Op.LE if second == '=' else Op.LT |
| elif first == '!' and second == '=': |
| return Op.NE |
| elif first == '=': |
| return Op.EQ |
| raise ValueError('Could not find valid operator in uri arg_str: ' + arg_str) |
| |
| |
| def _parse_arg(arg_str: str) -> Tuple[str, ConstraintClass.Op, str]: |
| """Parse argument string and return a tuple (key, operator, value). |
| |
| Given a string like 'branch_num>=4000', it returns a tuple ('branch_num', |
| Op.GE,'4000'). Raises ValueError exceptions in case ill formed arg_str is |
| passed like '>30', 'key>', 'key', 'key--31' |
| """ |
| op_start_ind = 0 |
| for ind, c in enumerate(arg_str): |
| if not c.isalnum() and c != '_': |
| op_start_ind = ind |
| break |
| if op_start_ind == 0: |
| raise ValueError('Could not find valid key in arg_str: ' + arg_str) |
| key = arg_str[:op_start_ind] |
| op = _read_op(arg_str, op_start_ind) |
| value = arg_str[op_start_ind + len(str(op)):] |
| if not value: |
| raise ValueError('Empty value in trace uri arg_str: ' + arg_str) |
| return (key, op, value) |
| |
| |
| def _args_dict_from_uri(uri: str, |
| type_hints) -> Dict[str, ConstraintWithList[str]]: |
| """Creates an the args dictionary from a trace URI. |
| |
| URIs have the form: |
| android_ci:day=2021-01-01;devices=blueline,crosshatch;key>=value |
| |
| This is converted to a dictionary of the form: |
| {'day': '2021-01-01', 'id': ['blueline', 'crosshatch'], |
| 'key': ConstraintClass('value', Op.GE)} |
| """ |
| _, args_str = util.parse_trace_uri(uri) |
| if not args_str: |
| return {} |
| |
| args_lst = args_str.split(';') |
| args_dict = dict() |
| for arg in args_lst: |
| (key, op, value) = _parse_arg(arg) |
| lst = value.split(',') |
| if len(lst) > 1: |
| args_dict[key] = lst |
| else: |
| args_dict[key] = value |
| |
| if key not in type_hints: |
| if op != ConstraintClass.Op.EQ: |
| raise ValueError(f'{key} only supports "=" operator') |
| continue |
| have_constraint = False |
| type_hint = type_hints[key] |
| type_args = type_hint.__args__ if hasattr(type_hint, '__args__') else () |
| for type_arg in type_args: |
| type_origin = type_arg.__origin__ if hasattr(type_arg, |
| '__origin__') else None |
| if type_origin is ConstraintClass: |
| have_constraint = True |
| break |
| if not have_constraint and op != ConstraintClass.Op.EQ: |
| raise ValueError('Operator other than "=" passed to argument which ' |
| 'does not have constraint type: ' + arg) |
| if have_constraint: |
| args_dict[key] = ConstraintClass(args_dict[key], op) |
| return args_dict |