Skip to content
This repository has been archived by the owner on Apr 26, 2024. It is now read-only.

Commit

Permalink
Batch of easier annotations
Browse files Browse the repository at this point in the history
  • Loading branch information
David Robertson committed May 27, 2022
1 parent e6d193a commit dda0cc8
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 22 deletions.
6 changes: 4 additions & 2 deletions synapse/config/tracer.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Any, Set
from typing import Any, List, Set

from synapse.types import JsonDict
from synapse.util.check_dependencies import DependencyException, check_requirements
Expand Down Expand Up @@ -49,7 +49,9 @@ def read_config(self, config: JsonDict, **kwargs: Any) -> None:

# The tracer is enabled so sanitize the config

self.opentracer_whitelist = opentracing_config.get("homeserver_whitelist", [])
self.opentracer_whitelist: List[str] = opentracing_config.get(
"homeserver_whitelist", []
)
if not isinstance(self.opentracer_whitelist, list):
raise ConfigError("Tracer homeserver_whitelist config is malformed")

Expand Down
46 changes: 26 additions & 20 deletions synapse/logging/opentracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -170,14 +170,18 @@ def set_fates(clotho, lachesis, atropos, father="Zues", mother="Themis"):
from functools import wraps
from typing import (
TYPE_CHECKING,
Any,
Callable,
Collection,
Dict,
Generator,
Iterable,
List,
Optional,
Pattern,
Type,
TypeVar,
Union,
)

import attr
Expand Down Expand Up @@ -374,7 +378,7 @@ def ensure_active_span_inner_2(*args, **kwargs):
# Setup


def init_tracer(hs: "HomeServer"):
def init_tracer(hs: "HomeServer") -> None:
"""Set the whitelists and initialise the JaegerClient tracer"""
global opentracing
if not hs.config.tracing.opentracer_enabled:
Expand Down Expand Up @@ -416,11 +420,11 @@ def init_tracer(hs: "HomeServer"):


@only_if_tracing
def set_homeserver_whitelist(homeserver_whitelist):
def set_homeserver_whitelist(homeserver_whitelist: Iterable[str]) -> None:
"""Sets the homeserver whitelist
Args:
homeserver_whitelist (Iterable[str]): regex of whitelisted homeservers
homeserver_whitelist: regexes specifying whitelisted homeservers
"""
global _homeserver_whitelist
if homeserver_whitelist:
Expand All @@ -431,15 +435,15 @@ def set_homeserver_whitelist(homeserver_whitelist):


@only_if_tracing
def whitelisted_homeserver(destination):
def whitelisted_homeserver(destination: str) -> bool:
"""Checks if a destination matches the whitelist
Args:
destination (str)
destination
"""

if _homeserver_whitelist:
return _homeserver_whitelist.match(destination)
return _homeserver_whitelist.match(destination) is not None
return False


Expand Down Expand Up @@ -586,27 +590,27 @@ def start_active_span_from_edu(

# Opentracing setters for tags, logs, etc
@only_if_tracing
def active_span():
def active_span() -> Optional["opentracing.Span"]:
"""Get the currently active span, if any"""
return opentracing.tracer.active_span


@ensure_active_span("set a tag")
def set_tag(key, value):
def set_tag(key: str, value: Union[str, bool, int, float]) -> None:
"""Sets a tag on the active span"""
assert opentracing.tracer.active_span is not None
opentracing.tracer.active_span.set_tag(key, value)


@ensure_active_span("log")
def log_kv(key_values, timestamp=None):
def log_kv(key_values: Dict[str, Any], timestamp=None) -> None:
"""Log to the active span"""
assert opentracing.tracer.active_span is not None
opentracing.tracer.active_span.log_kv(key_values, timestamp)


@ensure_active_span("set the traces operation name")
def set_operation_name(operation_name):
def set_operation_name(operation_name: str) -> None:
"""Sets the operation name of the active span"""
assert opentracing.tracer.active_span is not None
opentracing.tracer.active_span.set_operation_name(operation_name)
Expand Down Expand Up @@ -758,21 +762,21 @@ def span_context_from_request(request: Request) -> "Optional[opentracing.SpanCon


@only_if_tracing
def span_context_from_string(carrier):
def span_context_from_string(carrier: str) -> Optional["opentracing.SpanContext"]:
"""
Returns:
The active span context decoded from a string.
"""
carrier = json_decoder.decode(carrier)
return opentracing.tracer.extract(opentracing.Format.TEXT_MAP, carrier)
payload: Dict[str, str] = json_decoder.decode(carrier)
return opentracing.tracer.extract(opentracing.Format.TEXT_MAP, payload)


@only_if_tracing
def extract_text_map(carrier):
def extract_text_map(carrier: Dict[str, str]) -> Optional["opentracing.SpanContext"]:
"""
Wrapper method for opentracing's tracer.extract for TEXT_MAP.
Args:
carrier (dict): a dict possibly containing a span context.
carrier: a dict possibly containing a span context.
Returns:
The active span context extracted from carrier.
Expand Down Expand Up @@ -851,7 +855,7 @@ def err_back(result):
return decorator


def tag_args(func):
def tag_args(func: Callable[P, R]) -> Callable[P, R]:
"""
Tags all of the args to the active span.
"""
Expand All @@ -860,19 +864,21 @@ def tag_args(func):
return func

@wraps(func)
def _tag_args_inner(*args, **kwargs):
def _tag_args_inner(*args: P.args, **kwargs: P.kwargs) -> R:
argspec = inspect.getfullargspec(func)
for i, arg in enumerate(argspec.args[1:]):
set_tag("ARG_" + arg, args[i])
set_tag("args", args[len(argspec.args) :])
set_tag("ARG_" + arg, args[i]) # type: ignore[index]
set_tag("args", args[len(argspec.args) :]) # type: ignore[index]
set_tag("kwargs", kwargs)
return func(*args, **kwargs)

return _tag_args_inner


@contextlib.contextmanager
def trace_servlet(request: "SynapseRequest", extract_context: bool = False):
def trace_servlet(
request: "SynapseRequest", extract_context: bool = False
) -> Generator[None, None, None]:
"""Returns a context manager which traces a request. It starts a span
with some servlet specific tags such as the request metrics name and
request information.
Expand Down

0 comments on commit dda0cc8

Please sign in to comment.