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 from bt2 import component as bt2_component
  30 import datetime
  31 from collections import namedtuple
  32 import numbers
  33
  34
  35 # a pair of component and ComponentSpec
  36 _ComponentAndSpec = namedtuple('_ComponentAndSpec', ['comp', 'spec'])
  37
  38
  39 class ComponentSpec:
  40     def __init__(
  41         self,
  42         plugin_name,
  43         class_name,
  44         params=None,
  45         obj=None,
  46         logging_level=bt2_logging.LoggingLevel.NONE,
  47     ):
  48         utils._check_str(plugin_name)
  49         utils._check_str(class_name)
  50         utils._check_log_level(logging_level)
  51         self._plugin_name = plugin_name
  52         self._class_name = class_name
  53         self._logging_level = logging_level
  54         self._obj = obj
  55
  56         if type(params) is str:
  57             self._params = bt2.create_value({'inputs': [params]})
  58         else:
  59             self._params = bt2.create_value(params)
  60
  61     @property
  62     def plugin_name(self):
  63         return self._plugin_name
  64
  65     @property
  66     def class_name(self):
  67         return self._class_name
  68
  69     @property
  70     def logging_level(self):
  71         return self._logging_level
  72
  73     @property
  74     def params(self):
  75         return self._params
  76
  77     @property
  78     def obj(self):
  79         return self._obj
  80
  81
  82 # datetime.datetime or integral to nanoseconds
  83 def _get_ns(obj):
  84     if obj is None:
  85         return
  86
  87     if isinstance(obj, numbers.Real):
  88         # consider that it's already in seconds
  89         s = obj
  90     elif isinstance(obj, datetime.datetime):
  91         # s -> ns
  92         s = obj.timestamp()
  93     else:
  94         raise TypeError(
  95             '"{}" is not an integral number or a datetime.datetime object'.format(obj)
  96         )
  97
  98     return int(s * 1e9)
  99
 100
 101 class _CompClsType:
 102     SOURCE = 0
 103     FILTER = 1
 104
 105
 106 class _TraceCollectionMessageIteratorProxySink(bt2_component._UserSinkComponent):
 107     def __init__(self, params, msg_list):
 108         assert type(msg_list) is list
 109         self._msg_list = msg_list
 110         self._add_input_port('in')
 111
 112     def _user_graph_is_configured(self):
 113         self._msg_iter = self._create_input_port_message_iterator(
 114             self._input_ports['in']
 115         )
 116
 117     def _user_consume(self):
 118         assert self._msg_list[0] is None
 119         self._msg_list[0] = next(self._msg_iter)
 120
 121
 122 class TraceCollectionMessageIterator(bt2_message_iterator._MessageIterator):
 123     def __init__(
 124         self,
 125         source_component_specs,
 126         filter_component_specs=None,
 127         stream_intersection_mode=False,
 128         begin=None,
 129         end=None,
 130     ):
 131         utils._check_bool(stream_intersection_mode)
 132         self._stream_intersection_mode = stream_intersection_mode
 133         self._begin_ns = _get_ns(begin)
 134         self._end_ns = _get_ns(end)
 135         self._msg_list = [None]
 136
 137         if type(source_component_specs) is ComponentSpec:
 138             source_component_specs = [source_component_specs]
 139
 140         if type(filter_component_specs) is ComponentSpec:
 141             filter_component_specs = [filter_component_specs]
 142         elif filter_component_specs is None:
 143             filter_component_specs = []
 144
 145         self._src_comp_specs = source_component_specs
 146         self._flt_comp_specs = filter_component_specs
 147         self._next_suffix = 1
 148         self._connect_ports = False
 149
 150         # lists of _ComponentAndSpec
 151         self._src_comps_and_specs = []
 152         self._flt_comps_and_specs = []
 153
 154         self._validate_component_specs(source_component_specs)
 155         self._validate_component_specs(filter_component_specs)
 156         self._build_graph()
 157
 158     def _validate_component_specs(self, comp_specs):
 159         for comp_spec in comp_specs:
 160             if type(comp_spec) is not ComponentSpec:
 161                 raise TypeError(
 162                     '"{}" object is not a ComponentSpec'.format(type(comp_spec))
 163                 )
 164
 165     def __next__(self):
 166         assert self._msg_list[0] is None
 167         self._graph.run_once()
 168         msg = self._msg_list[0]
 169         assert msg is not None
 170         self._msg_list[0] = None
 171         return msg
 172
 173     def _create_stream_intersection_trimmer(self, component, port):
 174         # find the original parameters specified by the user to create
 175         # this port's component to get the `inputs` parameter
 176         for src_comp_and_spec in self._src_comps_and_specs:
 177             if component == src_comp_and_spec.comp:
 178                 break
 179
 180         try:
 181             inputs = src_comp_and_spec.spec.params['inputs']
 182         except Exception as e:
 183             raise ValueError(
 184                 'all source components must be created with an "inputs" parameter in stream intersection mode'
 185             ) from e
 186
 187         params = {'inputs': inputs}
 188
 189         # query the port's component for the `babeltrace.trace-info`
 190         # object which contains the stream intersection range for each
 191         # exposed trace
 192         query_exec = bt2.QueryExecutor(
 193             src_comp_and_spec.comp.cls, 'babeltrace.trace-info', params
 194         )
 195         trace_info_res = query_exec.query()
 196         begin = None
 197         end = None
 198
 199         # find the trace info for this port's trace
 200         try:
 201             for trace_info in trace_info_res:
 202                 for stream in trace_info['streams']:
 203                     if stream['port-name'] == port.name:
 204                         range_ns = trace_info['intersection-range-ns']
 205                         begin = range_ns['begin']
 206                         end = range_ns['end']
 207                         break
 208         except Exception:
 209             pass
 210
 211         if begin is None or end is None:
 212             raise RuntimeError(
 213                 'cannot find stream intersection range for port "{}"'.format(port.name)
 214             )
 215
 216         name = 'trimmer-{}-{}'.format(src_comp_and_spec.comp.name, port.name)
 217         return self._create_trimmer(begin, end, name)
 218
 219     def _create_muxer(self):
 220         plugin = bt2.find_plugin('utils')
 221
 222         if plugin is None:
 223             raise RuntimeError('cannot find "utils" plugin (needed for the muxer)')
 224
 225         if 'muxer' not in plugin.filter_component_classes:
 226             raise RuntimeError(
 227                 'cannot find "muxer" filter component class in "utils" plugin'
 228             )
 229
 230         comp_cls = plugin.filter_component_classes['muxer']
 231         return self._graph.add_component(comp_cls, 'muxer')
 232
 233     def _create_trimmer(self, begin_ns, end_ns, name):
 234         plugin = bt2.find_plugin('utils')
 235
 236         if plugin is None:
 237             raise RuntimeError('cannot find "utils" plugin (needed for the trimmer)')
 238
 239         if 'trimmer' not in plugin.filter_component_classes:
 240             raise RuntimeError(
 241                 'cannot find "trimmer" filter component class in "utils" plugin'
 242             )
 243
 244         params = {}
 245
 246         def ns_to_string(ns):
 247             s_part = ns // 1000000000
 248             ns_part = ns % 1000000000
 249             return '{}.{:09d}'.format(s_part, ns_part)
 250
 251         if begin_ns is not None:
 252             params['begin'] = ns_to_string(begin_ns)
 253
 254         if end_ns is not None:
 255             params['end'] = ns_to_string(end_ns)
 256
 257         comp_cls = plugin.filter_component_classes['trimmer']
 258         return self._graph.add_component(comp_cls, name, params)
 259
 260     def _get_unique_comp_name(self, comp_spec):
 261         name = '{}-{}'.format(comp_spec.plugin_name, comp_spec.class_name)
 262         comps_and_specs = itertools.chain(
 263             self._src_comps_and_specs, self._flt_comps_and_specs
 264         )
 265
 266         if name in [comp_and_spec.comp.name for comp_and_spec in comps_and_specs]:
 267             name += '-{}'.format(self._next_suffix)
 268             self._next_suffix += 1
 269
 270         return name
 271
 272     def _create_comp(self, comp_spec, comp_cls_type):
 273         plugin = bt2.find_plugin(comp_spec.plugin_name)
 274
 275         if plugin is None:
 276             raise ValueError('no such plugin: {}'.format(comp_spec.plugin_name))
 277
 278         if comp_cls_type == _CompClsType.SOURCE:
 279             comp_classes = plugin.source_component_classes
 280         else:
 281             comp_classes = plugin.filter_component_classes
 282
 283         if comp_spec.class_name not in comp_classes:
 284             cc_type = 'source' if comp_cls_type == _CompClsType.SOURCE else 'filter'
 285             raise ValueError(
 286                 'no such {} component class in "{}" plugin: {}'.format(
 287                     cc_type, comp_spec.plugin_name, comp_spec.class_name
 288                 )
 289             )
 290
 291         comp_cls = comp_classes[comp_spec.class_name]
 292         name = self._get_unique_comp_name(comp_spec)
 293         comp = self._graph.add_component(
 294             comp_cls, name, comp_spec.params, comp_spec.obj, comp_spec.logging_level
 295         )
 296         return comp
 297
 298     def _get_free_muxer_input_port(self):
 299         for port in self._muxer_comp.input_ports.values():
 300             if not port.is_connected:
 301                 return port
 302
 303     def _connect_src_comp_port(self, component, port):
 304         # if this trace collection iterator is in stream intersection
 305         # mode, we need this connection:
 306         #
 307         #     port -> trimmer -> muxer
 308         #
 309         # otherwise, simply:
 310         #
 311         #     port -> muxer
 312         if self._stream_intersection_mode:
 313             trimmer_comp = self._create_stream_intersection_trimmer(component, port)
 314             self._graph.connect_ports(port, trimmer_comp.input_ports['in'])
 315             port_to_muxer = trimmer_comp.output_ports['out']
 316         else:
 317             port_to_muxer = port
 318
 319         self._graph.connect_ports(port_to_muxer, self._get_free_muxer_input_port())
 320
 321     def _graph_port_added(self, component, port):
 322         if not self._connect_ports:
 323             return
 324
 325         if type(port) is bt2_port._InputPort:
 326             return
 327
 328         if component not in [comp.comp for comp in self._src_comps_and_specs]:
 329             # do not care about non-source components (muxer, trimmer, etc.)
 330             return
 331
 332         self._connect_src_comp_port(component, port)
 333
 334     def _build_graph(self):
 335         self._graph = bt2.Graph()
 336         self._graph.add_port_added_listener(self._graph_port_added)
 337         self._muxer_comp = self._create_muxer()
 338
 339         if self._begin_ns is not None or self._end_ns is not None:
 340             trimmer_comp = self._create_trimmer(self._begin_ns, self._end_ns, 'trimmer')
 341             self._graph.connect_ports(
 342                 self._muxer_comp.output_ports['out'], trimmer_comp.input_ports['in']
 343             )
 344             last_flt_out_port = trimmer_comp.output_ports['out']
 345         else:
 346             last_flt_out_port = self._muxer_comp.output_ports['out']
 347
 348         # create extra filter components (chained)
 349         for comp_spec in self._flt_comp_specs:
 350             comp = self._create_comp(comp_spec, _CompClsType.FILTER)
 351             self._flt_comps_and_specs.append(_ComponentAndSpec(comp, comp_spec))
 352
 353         # connect the extra filter chain
 354         for comp_and_spec in self._flt_comps_and_specs:
 355             in_port = list(comp_and_spec.comp.input_ports.values())[0]
 356             out_port = list(comp_and_spec.comp.output_ports.values())[0]
 357             self._graph.connect_ports(last_flt_out_port, in_port)
 358             last_flt_out_port = out_port
 359
 360         # Here we create the components, self._graph_port_added() is
 361         # called when they add ports, but the callback returns early
 362         # because self._connect_ports is False. This is because the
 363         # self._graph_port_added() could not find the associated source
 364         # component specification in self._src_comps_and_specs because
 365         # it does not exist yet (it needs the created component to
 366         # exist).
 367         for comp_spec in self._src_comp_specs:
 368             comp = self._create_comp(comp_spec, _CompClsType.SOURCE)
 369             self._src_comps_and_specs.append(_ComponentAndSpec(comp, comp_spec))
 370
 371         # Now we connect the ports which exist at this point. We allow
 372         # self._graph_port_added() to automatically connect _new_ ports.
 373         self._connect_ports = True
 374
 375         for comp_and_spec in self._src_comps_and_specs:
 376             # Keep a separate list because comp_and_spec.output_ports
 377             # could change during the connection of one of its ports.
 378             # Any new port is handled by self._graph_port_added().
 379             out_ports = [port for port in comp_and_spec.comp.output_ports.values()]
 380
 381             for out_port in out_ports:
 382                 if out_port.is_connected:
 383                     continue
 384
 385                 self._connect_src_comp_port(comp_and_spec.comp, out_port)
 386
 387         # Add the proxy sink, passing our message list to share consumed
 388         # messages with this trace collection message iterator.
 389         sink = self._graph.add_component(
 390             _TraceCollectionMessageIteratorProxySink, 'proxy-sink', obj=self._msg_list
 391         )
 392         sink_in_port = sink.input_ports['in']
 393
 394         # connect last filter to proxy sink
 395         self._graph.connect_ports(last_flt_out_port, sink_in_port)