no_cron_children_tracing.py

class NoCronChildrenTracer(trace_sdk.Tracer):
    '''
    Custom Tracer that breaks the parent relationship for cron tasks.

    This is so we can use a parent based sampler without excluding or including a whole
    cron run. Otherwise we would see one whole run a day with a 1/100 sampling rate. If
    the cron isn't the parent, we can sample its children by trace ID, and preserve the
    whole trace with descendant spans.

    This will also have the benfit of making the UI easier to use, as query results
    won't all link to the top level cron trace.
    '''

    # Would be nice to say "any task started in a cron" instead of explicit names, but
    # apply_async of the same task happens before run, and can only see one parent span
    # above, so can't tell if started in a cron at this point. Could maybe do something
    # with Baggage later.
    _CRON_CHILD_SPAN_NAMES = {
        'run/task_name_1',
        'run/task_name_2',
        'run/task_name_3',
    }

    def start_span(  # type: ignore
        self,
        name: str,
        context: Optional[context_api.Context] = None,
        kind: trace_api.SpanKind = trace_api.SpanKind.INTERNAL,
        attributes: types.Attributes = None,
        links: Sequence[trace_api.Link] = (),
        start_time: Optional[int] = None,
        record_exception: bool = True,
        set_status_on_exception: bool = True,
    ) -> trace_api.Span:
        if name in self._CRON_CHILD_SPAN_NAMES:
            old_parent_context = get_current_span(context=context).get_span_context()
            # Making current span INVALID_SPAN forces the next span to be a new trace.
            context = set_span_in_context(INVALID_SPAN, context=context)
            # Add a link so we can still find old parent in UI.
            links = [*links, trace_api.Link(old_parent_context)]
        return super().start_span(
            name,
            context=context,
            kind=kind,
            attributes=attributes,
            links=links,
            start_time=start_time,
            record_exception=record_exception,
            set_status_on_exception=set_status_on_exception,
        )


class NoCronChildrenTracerProvider(trace_sdk.TracerProvider):
    '''Provides NoCronChildrenTracerProvider instead of a regular Tracer.'''

    # This method is copied from the super class, just with the returned class changed.
    def get_tracer(
        self,
        instrumenting_module_name: str,
        instrumenting_library_version: Optional[str] = None,
        schema_url: Optional[str] = None,
    ) -> 'trace_api.Tracer':
        if not instrumenting_module_name:
            instrumenting_module_name = ''
            trace_sdk.logger.error('get_tracer called with missing module name.')
        if instrumenting_library_version is None:
            instrumenting_library_version = ''
        return NoCronChildrenTracer(
            self.sampler,
            self.resource,
            self._active_span_processor,
            self.id_generator,  # type: ignore
            InstrumentationInfo(
                instrumenting_module_name,
                instrumenting_library_version,
                schema_url,
            ),
            self._span_limits,
        )