src/bindings/python/bt2/bt2/trace_collection_message_iterator.py

   1 # The MIT License (MIT)
   2 #
   3 # Copyright (c) 2017 Philippe Proulx <pproulx@efficios.com>
   4 #
   5 # Permission is hereby granted, free of charge, to any person obtaining a copy
   6 # of this software and associated documentation files (the "Software"), to deal
   7 # in the Software without restriction, including without limitation the rights
   8 # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
   9 # copies of the Software, and to permit persons to whom the Software is
  10 # furnished to do so, subject to the following conditions:
  11 #
  12 # The above copyright notice and this permission notice shall be included in
  13 # all copies or substantial portions of the Software.
  14 #
  15 # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  16 # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  17 # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  18 # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  19 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  20 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
  21 # THE SOFTWARE.
  22
  23 from bt2 import utils
  24 import bt2
  25 import itertools
  26 from bt2 import message_iterator as bt2_message_iterator
  27 from bt2 import logging as bt2_logging
  28 from bt2 import port as bt2_port
  29 import datetime
  30 from collections import namedtuple
  31 import numbers
  32
  33
  34 # a pair of component and ComponentSpec
  35 _ComponentAndSpec = namedtuple('_ComponentAndSpec', ['comp', 'spec'])
  36
  37
  38 class ComponentSpec:
  39     def __init__(
  40         self,
  41         plugin_name,
  42         class_name,
  43         params=None,
  44         logging_level=bt2_logging.LoggingLevel.NONE,
  45     ):
  46         utils._check_str(plugin_name)
  47         utils._check_str(class_name)
  48         utils._check_log_level(logging_level)
  49         self._plugin_name = plugin_name
  50         self._class_name = class_name
  51         self._logging_level = logging_level
  52
  53         if type(params) is str:
  54             self._params = bt2.create_value({'inputs': [params]})
  55         else:
  56             self._params = bt2.create_value(params)
  57
  58     @property
  59     def plugin_name(self):
  60         return self._plugin_name
  61
  62     @property
  63     def class_name(self):
  64         return self._class_name
  65
  66     @property
  67     def logging_level(self):
  68         return self._logging_level
  69
  70     @property
  71     def params(self):
  72         return self._params
  73
  74
  75 # datetime.datetime or integral to nanoseconds
  76 def _get_ns(obj):
  77     if obj is None:
  78         return
  79
  80     if isinstance(obj, numbers.Real):
  81         # consider that it's already in seconds
  82         s = obj
  83     elif isinstance(obj, datetime.datetime):
  84         # s -> ns
  85         s = obj.timestamp()
  86     else:
  87         raise TypeError(
  88             '"{}" is not an integral number or a datetime.datetime object'.format(obj)
  89         )
  90
  91     return int(s * 1e9)
  92
  93
  94 class _CompClsType:
  95     SOURCE = 0
  96     FILTER = 1
  97
  98
  99 class TraceCollectionMessageIterator(bt2_message_iterator._MessageIterator):
 100     def __init__(
 101         self,
 102         source_component_specs,
 103         filter_component_specs=None,
 104         stream_intersection_mode=False,
 105         begin=None,
 106         end=None,
 107     ):
 108         utils._check_bool(stream_intersection_mode)
 109         self._stream_intersection_mode = stream_intersection_mode
 110         self._begin_ns = _get_ns(begin)
 111         self._end_ns = _get_ns(end)
 112
 113         if type(source_component_specs) is ComponentSpec:
 114             source_component_specs = [source_component_specs]
 115
 116         if type(filter_component_specs) is ComponentSpec:
 117             filter_component_specs = [filter_component_specs]
 118         elif filter_component_specs is None:
 119             filter_component_specs = []
 120
 121         self._src_comp_specs = source_component_specs
 122         self._flt_comp_specs = filter_component_specs
 123         self._next_suffix = 1
 124         self._connect_ports = False
 125
 126         # lists of _ComponentAndSpec
 127         self._src_comps_and_specs = []
 128         self._flt_comps_and_specs = []
 129
 130         self._validate_component_specs(source_component_specs)
 131         self._validate_component_specs(filter_component_specs)
 132         self._build_graph()
 133
 134     def _validate_component_specs(self, comp_specs):
 135         for comp_spec in comp_specs:
 136             if type(comp_spec) is not ComponentSpec:
 137                 raise TypeError(
 138                     '"{}" object is not a ComponentSpec'.format(type(comp_spec))
 139                 )
 140
 141     def __next__(self):
 142         return next(self._msg_iter)
 143
 144     def _create_stream_intersection_trimmer(self, component, port):
 145         # find the original parameters specified by the user to create
 146         # this port's component to get the `inputs` parameter
 147         for src_comp_and_spec in self._src_comps_and_specs:
 148             if component == src_comp_and_spec.comp:
 149                 break
 150
 151         try:
 152             inputs = src_comp_and_spec.spec.params['inputs']
 153         except Exception as e:
 154             raise ValueError(
 155                 'all source components must be created with an "inputs" parameter in stream intersection mode'
 156             ) from e
 157
 158         params = {'inputs': inputs}
 159
 160         # query the port's component for the `babeltrace.trace-info`
 161         # object which contains the stream intersection range for each
 162         # exposed trace
 163         query_exec = bt2.QueryExecutor(
 164             src_comp_and_spec.comp.cls, 'babeltrace.trace-info', params
 165         )
 166         trace_info_res = query_exec.query()
 167         begin = None
 168         end = None
 169
 170         # find the trace info for this port's trace
 171         try:
 172             for trace_info in trace_info_res:
 173                 for stream in trace_info['streams']:
 174                     if stream['port-name'] == port.name:
 175                         range_ns = trace_info['intersection-range-ns']
 176                         begin = range_ns['begin']
 177                         end = range_ns['end']
 178                         break
 179         except Exception:
 180             pass
 181
 182         if begin is None or end is None:
 183             raise RuntimeError(
 184                 'cannot find stream intersection range for port "{}"'.format(port.name)
 185             )
 186
 187         name = 'trimmer-{}-{}'.format(src_comp_and_spec.comp.name, port.name)
 188         return self._create_trimmer(begin, end, name)
 189
 190     def _create_muxer(self):
 191         plugin = bt2.find_plugin('utils')
 192
 193         if plugin is None:
 194             raise RuntimeError('cannot find "utils" plugin (needed for the muxer)')
 195
 196         if 'muxer' not in plugin.filter_component_classes:
 197             raise RuntimeError(
 198                 'cannot find "muxer" filter component class in "utils" plugin'
 199             )
 200
 201         comp_cls = plugin.filter_component_classes['muxer']
 202         return self._graph.add_component(comp_cls, 'muxer')
 203
 204     def _create_trimmer(self, begin_ns, end_ns, name):
 205         plugin = bt2.find_plugin('utils')
 206
 207         if plugin is None:
 208             raise RuntimeError('cannot find "utils" plugin (needed for the trimmer)')
 209
 210         if 'trimmer' not in plugin.filter_component_classes:
 211             raise RuntimeError(
 212                 'cannot find "trimmer" filter component class in "utils" plugin'
 213             )
 214
 215         params = {}
 216
 217         def ns_to_string(ns):
 218             s_part = ns // 1000000000
 219             ns_part = ns % 1000000000
 220             return '{}.{:09d}'.format(s_part, ns_part)
 221
 222         if begin_ns is not None:
 223             params['begin'] = ns_to_string(begin_ns)
 224
 225         if end_ns is not None:
 226             params['end'] = ns_to_string(end_ns)
 227
 228         comp_cls = plugin.filter_component_classes['trimmer']
 229         return self._graph.add_component(comp_cls, name, params)
 230
 231     def _get_unique_comp_name(self, comp_spec):
 232         name = '{}-{}'.format(comp_spec.plugin_name, comp_spec.class_name)
 233         comps_and_specs = itertools.chain(
 234             self._src_comps_and_specs, self._flt_comps_and_specs
 235         )
 236
 237         if name in [comp_and_spec.comp.name for comp_and_spec in comps_and_specs]:
 238             name += '-{}'.format(self._next_suffix)
 239             self._next_suffix += 1
 240
 241         return name
 242
 243     def _create_comp(self, comp_spec, comp_cls_type):
 244         plugin = bt2.find_plugin(comp_spec.plugin_name)
 245
 246         if plugin is None:
 247             raise ValueError('no such plugin: {}'.format(comp_spec.plugin_name))
 248
 249         if comp_cls_type == _CompClsType.SOURCE:
 250             comp_classes = plugin.source_component_classes
 251         else:
 252             comp_classes = plugin.filter_component_classes
 253
 254         if comp_spec.class_name not in comp_classes:
 255             cc_type = 'source' if comp_cls_type == _CompClsType.SOURCE else 'filter'
 256             raise ValueError(
 257                 'no such {} component class in "{}" plugin: {}'.format(
 258                     cc_type, comp_spec.plugin_name, comp_spec.class_name
 259                 )
 260             )
 261
 262         comp_cls = comp_classes[comp_spec.class_name]
 263         name = self._get_unique_comp_name(comp_spec)
 264         comp = self._graph.add_component(
 265             comp_cls, name, comp_spec.params, comp_spec.logging_level
 266         )
 267         return comp
 268
 269     def _get_free_muxer_input_port(self):
 270         for port in self._muxer_comp.input_ports.values():
 271             if not port.is_connected:
 272                 return port
 273
 274     def _connect_src_comp_port(self, component, port):
 275         # if this trace collection iterator is in stream intersection
 276         # mode, we need this connection:
 277         #
 278         #     port -> trimmer -> muxer
 279         #
 280         # otherwise, simply:
 281         #
 282         #     port -> muxer
 283         if self._stream_intersection_mode:
 284             trimmer_comp = self._create_stream_intersection_trimmer(component, port)
 285             self._graph.connect_ports(port, trimmer_comp.input_ports['in'])
 286             port_to_muxer = trimmer_comp.output_ports['out']
 287         else:
 288             port_to_muxer = port
 289
 290         self._graph.connect_ports(port_to_muxer, self._get_free_muxer_input_port())
 291
 292     def _graph_port_added(self, component, port):
 293         if not self._connect_ports:
 294             return
 295
 296         if type(port) is bt2_port._InputPort:
 297             return
 298
 299         if component not in [comp.comp for comp in self._src_comps_and_specs]:
 300             # do not care about non-source components (muxer, trimmer, etc.)
 301             return
 302
 303         self._connect_src_comp_port(component, port)
 304
 305     def _build_graph(self):
 306         self._graph = bt2.Graph()
 307         self._graph.add_port_added_listener(self._graph_port_added)
 308         self._muxer_comp = self._create_muxer()
 309
 310         if self._begin_ns is not None or self._end_ns is not None:
 311             trimmer_comp = self._create_trimmer(self._begin_ns, self._end_ns, 'trimmer')
 312             self._graph.connect_ports(
 313                 self._muxer_comp.output_ports['out'], trimmer_comp.input_ports['in']
 314             )
 315             msg_iter_port = trimmer_comp.output_ports['out']
 316         else:
 317             msg_iter_port = self._muxer_comp.output_ports['out']
 318
 319         # create extra filter components (chained)
 320         for comp_spec in self._flt_comp_specs:
 321             comp = self._create_comp(comp_spec, _CompClsType.FILTER)
 322             self._flt_comps_and_specs.append(_ComponentAndSpec(comp, comp_spec))
 323
 324         # connect the extra filter chain
 325         for comp_and_spec in self._flt_comps_and_specs:
 326             in_port = list(comp_and_spec.comp.input_ports.values())[0]
 327             out_port = list(comp_and_spec.comp.output_ports.values())[0]
 328             self._graph.connect_ports(msg_iter_port, in_port)
 329             msg_iter_port = out_port
 330
 331         # Here we create the components, self._graph_port_added() is
 332         # called when they add ports, but the callback returns early
 333         # because self._connect_ports is False. This is because the
 334         # self._graph_port_added() could not find the associated source
 335         # component specification in self._src_comps_and_specs because
 336         # it does not exist yet (it needs the created component to
 337         # exist).
 338         for comp_spec in self._src_comp_specs:
 339             comp = self._create_comp(comp_spec, _CompClsType.SOURCE)
 340             self._src_comps_and_specs.append(_ComponentAndSpec(comp, comp_spec))
 341
 342         # Now we connect the ports which exist at this point. We allow
 343         # self._graph_port_added() to automatically connect _new_ ports.
 344         self._connect_ports = True
 345
 346         for comp_and_spec in self._src_comps_and_specs:
 347             # Keep a separate list because comp_and_spec.output_ports
 348             # could change during the connection of one of its ports.
 349             # Any new port is handled by self._graph_port_added().
 350             out_ports = [port for port in comp_and_spec.comp.output_ports.values()]
 351
 352             for out_port in out_ports:
 353                 if out_port.is_connected:
 354                     continue
 355
 356                 self._connect_src_comp_port(comp_and_spec.comp, out_port)
 357
 358         # create this trace collection iterator's message iterator
 359         self._msg_iter = self._graph.create_output_port_message_iterator(msg_iter_port)