123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595 |
- # Copyright 2016 The LUCI Authors. All rights reserved.
- # Use of this source code is governed under the Apache License, Version 2.0
- # that can be found in the LICENSE file.
- import collections
- import contextlib
- import json
- import os
- import posixpath
- import socket
- import sys
- import threading
- import time
- from . import streamname, varint
- if sys.platform == "win32":
- from ctypes import GetLastError
- _PY2 = sys.version_info[0] == 2
- _MAPPING = collections.Mapping if _PY2 else collections.abc.Mapping
- _StreamParamsBase = collections.namedtuple(
- '_StreamParamsBase', ('name', 'type', 'content_type', 'tags'))
- # Magic number at the beginning of a Butler stream
- #
- # See "ProtocolFrameHeaderMagic" in:
- # <luci-go>/logdog/client/butlerlib/streamproto
- BUTLER_MAGIC = b'BTLR1\x1e'
- class StreamParams(_StreamParamsBase):
- """Defines the set of parameters to apply to a new stream."""
- # A text content stream.
- TEXT = 'text'
- # A binary content stream.
- BINARY = 'binary'
- # A datagram content stream.
- DATAGRAM = 'datagram'
- @classmethod
- def make(cls, **kwargs):
- """Returns (StreamParams): A new StreamParams instance with supplied values.
- Any parameter that isn't supplied will be set to None.
- Args:
- kwargs (dict): Named parameters to apply.
- """
- return cls(**{f: kwargs.get(f) for f in cls._fields})
- def validate(self):
- """Raises (ValueError): if the parameters are not valid."""
- streamname.validate_stream_name(self.name)
- if self.type not in (self.TEXT, self.BINARY, self.DATAGRAM):
- raise ValueError('Invalid type (%s)' % (self.type,))
- if self.tags is not None:
- if not isinstance(self.tags, _MAPPING):
- raise ValueError('Invalid tags type (%s)' % (self.tags,))
- for k, v in self.tags.items():
- streamname.validate_tag(k, v)
- def to_json(self):
- """Returns (str): The JSON representation of the StreamParams.
- Converts stream parameters to JSON for Butler consumption.
- Raises:
- ValueError: if these parameters are not valid.
- """
- self.validate()
- obj = {
- 'name': self.name,
- 'type': self.type,
- }
- def _maybe_add(key, value):
- if value is not None:
- obj[key] = value
- _maybe_add('contentType', self.content_type)
- _maybe_add('tags', self.tags)
- # Note that "dumps' will dump UTF-8 by default, which is what Butler wants.
- return json.dumps(obj, sort_keys=True, ensure_ascii=True, indent=None)
- class StreamProtocolRegistry(object):
- """Registry of streamserver URI protocols and their client classes.
- """
- def __init__(self):
- self._registry = {}
- def register_protocol(self, protocol, client_cls):
- assert issubclass(client_cls, StreamClient)
- if self._registry.get(protocol) is not None:
- raise KeyError('Duplicate protocol registered.')
- self._registry[protocol] = client_cls
- def create(self, uri, **kwargs):
- """Returns (StreamClient): A stream client for the specified URI.
- This uses the default StreamProtocolRegistry to instantiate a StreamClient
- for the specified URI.
- Args:
- uri (str): The streamserver URI.
- kwargs: keyword arguments to forward to the stream. See
- StreamClient.__init__.
- Raises:
- ValueError: if the supplied URI references an invalid or improperly
- configured streamserver.
- """
- uri = uri.split(':', 1)
- if len(uri) != 2:
- raise ValueError('Invalid stream server URI [%s]' % (uri,))
- protocol, value = uri
- client_cls = self._registry.get(protocol)
- if not client_cls:
- raise ValueError('Unknown stream client protocol (%s)' % (protocol,))
- return client_cls._create(value, **kwargs)
- # Default (global) registry.
- _default_registry = StreamProtocolRegistry()
- create = _default_registry.create
- class StreamClient(object):
- """Abstract base class for a streamserver client.
- """
- class _StreamBase(object):
- """ABC for StreamClient streams."""
- def __init__(self, stream_client, params):
- self._stream_client = stream_client
- self._params = params
- @property
- def params(self):
- """Returns (StreamParams): The stream parameters."""
- return self._params
- @property
- def path(self):
- """Returns (streamname.StreamPath): The stream path.
- Raises:
- ValueError: if the stream path is invalid, or if the stream prefix is
- not defined in the client.
- """
- return self._stream_client.get_stream_path(self._params.name)
- def get_viewer_url(self):
- """Returns (str): The viewer URL for this stream.
- Raises:
- KeyError: if information needed to construct the URL is missing.
- ValueError: if the stream prefix or name do not form a valid stream
- path.
- """
- return self._stream_client.get_viewer_url(self._params.name)
- class _BasicStream(_StreamBase):
- """Wraps a basic file descriptor, offering "write" and "close"."""
- def __init__(self, stream_client, params, fd):
- super(StreamClient._BasicStream, self).__init__(stream_client, params)
- self._fd = fd
- @property
- def fd(self):
- return self._fd
- def fileno(self):
- return self._fd.fileno()
- def write(self, data):
- return self._fd.write(data)
- def close(self):
- return self._fd.close()
- class _TextStream(_BasicStream):
- """Extends _BasicStream, ensuring data written is UTF-8 text."""
- def __init__(self, stream_client, params, fd):
- super(StreamClient._TextStream, self).__init__(stream_client, params, fd)
- self._fd = fd
- def write(self, data):
- if _PY2 and isinstance(data, str):
- # byte string is unfortunately accepted in py2 because of
- # undifferentiated usage of `str` and `unicode` but it should be
- # discontinued in py3. User should switch to binary stream instead
- # if there's a need to write bytes.
- return self._fd.write(data)
- elif _PY2 and isinstance(data, unicode):
- return self._fd.write(data.encode('utf-8'))
- elif not _PY2 and isinstance(data, str):
- return self._fd.write(data.encode('utf-8'))
- else:
- raise ValueError(
- 'expect str, got %r that is type %s' % (data, type(data),))
- class _DatagramStream(_StreamBase):
- """Wraps a stream object to write length-prefixed datagrams."""
- def __init__(self, stream_client, params, fd):
- super(StreamClient._DatagramStream, self).__init__(stream_client, params)
- self._fd = fd
- def send(self, data):
- varint.write_uvarint(self._fd, len(data))
- self._fd.write(data)
- def close(self):
- return self._fd.close()
- def __init__(self, project=None, prefix=None, coordinator_host=None,
- namespace=''):
- """Constructs a new base StreamClient instance.
- Args:
- project (str or None): If not None, the name of the log stream project.
- prefix (str or None): If not None, the log stream session prefix.
- coordinator_host (str or None): If not None, the name of the Coordinator
- host that this stream client is bound to. This will be used to
- construct viewer URLs for generated streams.
- namespace (str): The prefix to apply to all streams opened by this client.
- """
- self._project = project
- self._prefix = prefix
- self._coordinator_host = coordinator_host
- self._namespace = namespace
- self._name_lock = threading.Lock()
- self._names = set()
- @property
- def project(self):
- """Returns (str or None): The stream project, or None if not configured."""
- return self._project
- @property
- def prefix(self):
- """Returns (str or None): The stream prefix, or None if not configured."""
- return self._prefix
- @property
- def coordinator_host(self):
- """Returns (str or None): The coordinator host, or None if not configured.
- """
- return self._coordinator_host
- @property
- def namespace(self):
- """Returns (str): The namespace for all streams opened by this client.
- Empty if not configured.
- """
- return self._namespace
- def get_stream_path(self, name):
- """Returns (streamname.StreamPath): The stream path.
- Args:
- name (str): The name of the stream.
- Raises:
- KeyError: if information needed to construct the path is missing.
- ValueError: if the stream path is invalid, or if the stream prefix is
- not defined in the client.
- """
- if not self._prefix:
- raise KeyError('Stream prefix is not configured')
- return streamname.StreamPath.make(self._prefix, name)
- def get_viewer_url(self, name):
- """Returns (str): The LogDog viewer URL for the named stream.
- Args:
- name (str): The name of the stream. This can also be a query glob.
- Raises:
- KeyError: if information needed to construct the URL is missing.
- ValueError: if the stream prefix or name do not form a valid stream
- path.
- """
- if not self._coordinator_host:
- raise KeyError('Coordinator host is not configured')
- if not self._project:
- raise KeyError('Stream project is not configured')
- return streamname.get_logdog_viewer_url(
- self._coordinator_host,
- self._project,
- self.get_stream_path(name))
- def _register_new_stream(self, name):
- """Registers a new stream name.
- The Butler will internally reject any duplicate stream names. However, there
- isn't really feedback when this happens except a closed stream client. This
- is a client-side check to provide a more user-friendly experience in the
- event that a user attempts to register a duplicate stream name.
- Note that this is imperfect, as something else could register stream names
- with the same Butler instance and this library has no means of tracking.
- This is a best-effort experience, not a reliable check.
- Args:
- name (str): The name of the stream.
- Raises:
- ValueError if the stream name has already been registered.
- """
- with self._name_lock:
- if name in self._names:
- raise ValueError("Duplicate stream name [%s]" % (name,))
- self._names.add(name)
- @classmethod
- def _create(cls, value, **kwargs):
- """Returns (StreamClient): A new stream client instance.
- Validates the streamserver parameters and creates a new StreamClient
- instance that connects to them.
- Implementing classes must override this.
- """
- raise NotImplementedError()
- def _connect_raw(self):
- """Returns (file): A new file-like stream.
- Creates a new raw connection to the streamserver. This connection MUST not
- have any data written to it past initialization (if needed) when it has been
- returned.
- The file-like object must implement `write`, `fileno`, `flush`, and `close`.
- Implementing classes must override this.
- """
- raise NotImplementedError()
- def new_connection(self, params):
- """Returns (file): A new configured stream.
- The returned object implements (minimally) `write` and `close`.
- Creates a new LogDog stream with the specified parameters.
- Args:
- params (StreamParams): The parameters to use with the new connection.
- Raises:
- ValueError if the stream name has already been used, or if the parameters
- are not valid.
- """
- self._register_new_stream(params.name)
- params_bytes = params.to_json().encode('utf-8')
- fobj = self._connect_raw()
- fobj.write(BUTLER_MAGIC)
- varint.write_uvarint(fobj, len(params_bytes))
- fobj.write(params_bytes)
- return fobj
- @contextlib.contextmanager
- def text(self, name, **kwargs):
- """Context manager to create, use, and teardown a TEXT stream.
- This context manager creates a new butler TEXT stream with the specified
- parameters, yields it, and closes it on teardown.
- Args:
- name (str): the LogDog name of the stream.
- kwargs (dict): Log stream parameters. These may be any keyword arguments
- accepted by `open_text`.
- Returns (file): A file-like object to a Butler UTF-8 text stream supporting
- `write`.
- """
- fobj = None
- try:
- fobj = self.open_text(name, **kwargs)
- yield fobj
- finally:
- if fobj is not None:
- fobj.close()
- def open_text(self, name, content_type=None, tags=None):
- """Returns (file): A file-like object for a single text stream.
- This creates a new butler TEXT stream with the specified parameters.
- Args:
- name (str): the LogDog name of the stream.
- content_type (str): The optional content type of the stream. If None, a
- default content type will be chosen by the Butler.
- tags (dict): An optional key/value dictionary pair of LogDog stream tags.
- Returns (file): A file-like object to a Butler text stream. This object can
- have UTF-8 text content written to it with its `write` method, and must
- be closed when finished using its `close` method.
- """
- params = StreamParams.make(
- name=posixpath.join(self._namespace, name),
- type=StreamParams.TEXT,
- content_type=content_type,
- tags=tags)
- return self._TextStream(self, params, self.new_connection(params))
- @contextlib.contextmanager
- def binary(self, name, **kwargs):
- """Context manager to create, use, and teardown a BINARY stream.
- This context manager creates a new butler BINARY stream with the specified
- parameters, yields it, and closes it on teardown.
- Args:
- name (str): the LogDog name of the stream.
- kwargs (dict): Log stream parameters. These may be any keyword arguments
- accepted by `open_binary`.
- Returns (file): A file-like object to a Butler binary stream supporting
- `write`.
- """
- fobj = None
- try:
- fobj = self.open_binary(name, **kwargs)
- yield fobj
- finally:
- if fobj is not None:
- fobj.close()
- def open_binary(self, name, content_type=None, tags=None):
- """Returns (file): A file-like object for a single binary stream.
- This creates a new butler BINARY stream with the specified parameters.
- Args:
- name (str): the LogDog name of the stream.
- content_type (str): The optional content type of the stream. If None, a
- default content type will be chosen by the Butler.
- tags (dict): An optional key/value dictionary pair of LogDog stream tags.
- Returns (file): A file-like object to a Butler binary stream. This object
- can have UTF-8 content written to it with its `write` method, and must
- be closed when finished using its `close` method.
- """
- params = StreamParams.make(
- name=posixpath.join(self._namespace, name),
- type=StreamParams.BINARY,
- content_type=content_type,
- tags=tags)
- return self._BasicStream(self, params, self.new_connection(params))
- @contextlib.contextmanager
- def datagram(self, name, **kwargs):
- """Context manager to create, use, and teardown a DATAGRAM stream.
- This context manager creates a new butler DATAAGRAM stream with the
- specified parameters, yields it, and closes it on teardown.
- Args:
- name (str): the LogDog name of the stream.
- kwargs (dict): Log stream parameters. These may be any keyword arguments
- accepted by `open_datagram`.
- Returns (_DatagramStream): A datagram stream object. Datagrams can be
- written to it using its `send` method.
- """
- fobj = None
- try:
- fobj = self.open_datagram(name, **kwargs)
- yield fobj
- finally:
- if fobj is not None:
- fobj.close()
- def open_datagram(self, name, content_type=None, tags=None):
- """Creates a new butler DATAGRAM stream with the specified parameters.
- Args:
- name (str): the LogDog name of the stream.
- content_type (str): The optional content type of the stream. If None, a
- default content type will be chosen by the Butler.
- tags (dict): An optional key/value dictionary pair of LogDog stream tags.
- Returns (_DatagramStream): A datagram stream object. Datagrams can be
- written to it using its `send` method. This object must be closed when
- finished by using its `close` method.
- """
- params = StreamParams.make(
- name=posixpath.join(self._namespace, name),
- type=StreamParams.DATAGRAM,
- content_type=content_type,
- tags=tags)
- return self._DatagramStream(self, params, self.new_connection(params))
- class _NamedPipeStreamClient(StreamClient):
- """A StreamClient implementation that connects to a Windows named pipe.
- """
- def __init__(self, name, **kwargs):
- r"""Initializes a new Windows named pipe stream client.
- Args:
- name (str): The name of the Windows named pipe to use (e.g., "\\.\name")
- """
- super(_NamedPipeStreamClient, self).__init__(**kwargs)
- self._name = '\\\\.\\pipe\\' + name
- @classmethod
- def _create(cls, value, **kwargs):
- return cls(value, **kwargs)
- ERROR_PIPE_BUSY = 231
- def _connect_raw(self):
- # This is a similar procedure to the one in
- # https://github.com/microsoft/go-winio/blob/master/pipe.go (tryDialPipe)
- while True:
- try:
- return open(self._name, 'wb+', buffering=0)
- except (OSError, IOError):
- if GetLastError() != self.ERROR_PIPE_BUSY:
- raise
- time.sleep(0.001) # 1ms
- _default_registry.register_protocol('net.pipe', _NamedPipeStreamClient)
- class _UnixDomainSocketStreamClient(StreamClient):
- """A StreamClient implementation that uses a UNIX domain socket.
- """
- class SocketFile(object):
- """A write-only file-like object that writes to a UNIX socket."""
- def __init__(self, sock):
- self._sock = sock
- def fileno(self):
- return self._sock
- def write(self, data):
- self._sock.sendall(data)
- def flush(self):
- pass
- def close(self):
- self._sock.close()
- def __init__(self, path, **kwargs):
- """Initializes a new UNIX domain socket stream client.
- Args:
- path (str): The path to the named UNIX domain socket.
- """
- super(_UnixDomainSocketStreamClient, self).__init__(**kwargs)
- self._path = path
- @classmethod
- def _create(cls, value, **kwargs):
- if not os.path.exists(value):
- raise ValueError('UNIX domain socket [%s] does not exist.' % (value,))
- return cls(value, **kwargs)
- def _connect_raw(self):
- sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
- sock.connect(self._path)
- return self.SocketFile(sock)
- _default_registry.register_protocol('unix', _UnixDomainSocketStreamClient)
|