stream.py 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595
  1. # Copyright 2016 The LUCI Authors. All rights reserved.
  2. # Use of this source code is governed under the Apache License, Version 2.0
  3. # that can be found in the LICENSE file.
  4. import collections
  5. import contextlib
  6. import json
  7. import os
  8. import posixpath
  9. import socket
  10. import sys
  11. import threading
  12. import time
  13. from . import streamname, varint
  14. if sys.platform == "win32":
  15. from ctypes import GetLastError
  16. _PY2 = sys.version_info[0] == 2
  17. _MAPPING = collections.Mapping if _PY2 else collections.abc.Mapping
  18. _StreamParamsBase = collections.namedtuple(
  19. '_StreamParamsBase', ('name', 'type', 'content_type', 'tags'))
  20. # Magic number at the beginning of a Butler stream
  21. #
  22. # See "ProtocolFrameHeaderMagic" in:
  23. # <luci-go>/logdog/client/butlerlib/streamproto
  24. BUTLER_MAGIC = b'BTLR1\x1e'
  25. class StreamParams(_StreamParamsBase):
  26. """Defines the set of parameters to apply to a new stream."""
  27. # A text content stream.
  28. TEXT = 'text'
  29. # A binary content stream.
  30. BINARY = 'binary'
  31. # A datagram content stream.
  32. DATAGRAM = 'datagram'
  33. @classmethod
  34. def make(cls, **kwargs):
  35. """Returns (StreamParams): A new StreamParams instance with supplied values.
  36. Any parameter that isn't supplied will be set to None.
  37. Args:
  38. kwargs (dict): Named parameters to apply.
  39. """
  40. return cls(**{f: kwargs.get(f) for f in cls._fields})
  41. def validate(self):
  42. """Raises (ValueError): if the parameters are not valid."""
  43. streamname.validate_stream_name(self.name)
  44. if self.type not in (self.TEXT, self.BINARY, self.DATAGRAM):
  45. raise ValueError('Invalid type (%s)' % (self.type,))
  46. if self.tags is not None:
  47. if not isinstance(self.tags, _MAPPING):
  48. raise ValueError('Invalid tags type (%s)' % (self.tags,))
  49. for k, v in self.tags.items():
  50. streamname.validate_tag(k, v)
  51. def to_json(self):
  52. """Returns (str): The JSON representation of the StreamParams.
  53. Converts stream parameters to JSON for Butler consumption.
  54. Raises:
  55. ValueError: if these parameters are not valid.
  56. """
  57. self.validate()
  58. obj = {
  59. 'name': self.name,
  60. 'type': self.type,
  61. }
  62. def _maybe_add(key, value):
  63. if value is not None:
  64. obj[key] = value
  65. _maybe_add('contentType', self.content_type)
  66. _maybe_add('tags', self.tags)
  67. # Note that "dumps' will dump UTF-8 by default, which is what Butler wants.
  68. return json.dumps(obj, sort_keys=True, ensure_ascii=True, indent=None)
  69. class StreamProtocolRegistry(object):
  70. """Registry of streamserver URI protocols and their client classes.
  71. """
  72. def __init__(self):
  73. self._registry = {}
  74. def register_protocol(self, protocol, client_cls):
  75. assert issubclass(client_cls, StreamClient)
  76. if self._registry.get(protocol) is not None:
  77. raise KeyError('Duplicate protocol registered.')
  78. self._registry[protocol] = client_cls
  79. def create(self, uri, **kwargs):
  80. """Returns (StreamClient): A stream client for the specified URI.
  81. This uses the default StreamProtocolRegistry to instantiate a StreamClient
  82. for the specified URI.
  83. Args:
  84. uri (str): The streamserver URI.
  85. kwargs: keyword arguments to forward to the stream. See
  86. StreamClient.__init__.
  87. Raises:
  88. ValueError: if the supplied URI references an invalid or improperly
  89. configured streamserver.
  90. """
  91. uri = uri.split(':', 1)
  92. if len(uri) != 2:
  93. raise ValueError('Invalid stream server URI [%s]' % (uri,))
  94. protocol, value = uri
  95. client_cls = self._registry.get(protocol)
  96. if not client_cls:
  97. raise ValueError('Unknown stream client protocol (%s)' % (protocol,))
  98. return client_cls._create(value, **kwargs)
  99. # Default (global) registry.
  100. _default_registry = StreamProtocolRegistry()
  101. create = _default_registry.create
  102. class StreamClient(object):
  103. """Abstract base class for a streamserver client.
  104. """
  105. class _StreamBase(object):
  106. """ABC for StreamClient streams."""
  107. def __init__(self, stream_client, params):
  108. self._stream_client = stream_client
  109. self._params = params
  110. @property
  111. def params(self):
  112. """Returns (StreamParams): The stream parameters."""
  113. return self._params
  114. @property
  115. def path(self):
  116. """Returns (streamname.StreamPath): The stream path.
  117. Raises:
  118. ValueError: if the stream path is invalid, or if the stream prefix is
  119. not defined in the client.
  120. """
  121. return self._stream_client.get_stream_path(self._params.name)
  122. def get_viewer_url(self):
  123. """Returns (str): The viewer URL for this stream.
  124. Raises:
  125. KeyError: if information needed to construct the URL is missing.
  126. ValueError: if the stream prefix or name do not form a valid stream
  127. path.
  128. """
  129. return self._stream_client.get_viewer_url(self._params.name)
  130. class _BasicStream(_StreamBase):
  131. """Wraps a basic file descriptor, offering "write" and "close"."""
  132. def __init__(self, stream_client, params, fd):
  133. super(StreamClient._BasicStream, self).__init__(stream_client, params)
  134. self._fd = fd
  135. @property
  136. def fd(self):
  137. return self._fd
  138. def fileno(self):
  139. return self._fd.fileno()
  140. def write(self, data):
  141. return self._fd.write(data)
  142. def close(self):
  143. return self._fd.close()
  144. class _TextStream(_BasicStream):
  145. """Extends _BasicStream, ensuring data written is UTF-8 text."""
  146. def __init__(self, stream_client, params, fd):
  147. super(StreamClient._TextStream, self).__init__(stream_client, params, fd)
  148. self._fd = fd
  149. def write(self, data):
  150. if _PY2 and isinstance(data, str):
  151. # byte string is unfortunately accepted in py2 because of
  152. # undifferentiated usage of `str` and `unicode` but it should be
  153. # discontinued in py3. User should switch to binary stream instead
  154. # if there's a need to write bytes.
  155. return self._fd.write(data)
  156. elif _PY2 and isinstance(data, unicode):
  157. return self._fd.write(data.encode('utf-8'))
  158. elif not _PY2 and isinstance(data, str):
  159. return self._fd.write(data.encode('utf-8'))
  160. else:
  161. raise ValueError(
  162. 'expect str, got %r that is type %s' % (data, type(data),))
  163. class _DatagramStream(_StreamBase):
  164. """Wraps a stream object to write length-prefixed datagrams."""
  165. def __init__(self, stream_client, params, fd):
  166. super(StreamClient._DatagramStream, self).__init__(stream_client, params)
  167. self._fd = fd
  168. def send(self, data):
  169. varint.write_uvarint(self._fd, len(data))
  170. self._fd.write(data)
  171. def close(self):
  172. return self._fd.close()
  173. def __init__(self, project=None, prefix=None, coordinator_host=None,
  174. namespace=''):
  175. """Constructs a new base StreamClient instance.
  176. Args:
  177. project (str or None): If not None, the name of the log stream project.
  178. prefix (str or None): If not None, the log stream session prefix.
  179. coordinator_host (str or None): If not None, the name of the Coordinator
  180. host that this stream client is bound to. This will be used to
  181. construct viewer URLs for generated streams.
  182. namespace (str): The prefix to apply to all streams opened by this client.
  183. """
  184. self._project = project
  185. self._prefix = prefix
  186. self._coordinator_host = coordinator_host
  187. self._namespace = namespace
  188. self._name_lock = threading.Lock()
  189. self._names = set()
  190. @property
  191. def project(self):
  192. """Returns (str or None): The stream project, or None if not configured."""
  193. return self._project
  194. @property
  195. def prefix(self):
  196. """Returns (str or None): The stream prefix, or None if not configured."""
  197. return self._prefix
  198. @property
  199. def coordinator_host(self):
  200. """Returns (str or None): The coordinator host, or None if not configured.
  201. """
  202. return self._coordinator_host
  203. @property
  204. def namespace(self):
  205. """Returns (str): The namespace for all streams opened by this client.
  206. Empty if not configured.
  207. """
  208. return self._namespace
  209. def get_stream_path(self, name):
  210. """Returns (streamname.StreamPath): The stream path.
  211. Args:
  212. name (str): The name of the stream.
  213. Raises:
  214. KeyError: if information needed to construct the path is missing.
  215. ValueError: if the stream path is invalid, or if the stream prefix is
  216. not defined in the client.
  217. """
  218. if not self._prefix:
  219. raise KeyError('Stream prefix is not configured')
  220. return streamname.StreamPath.make(self._prefix, name)
  221. def get_viewer_url(self, name):
  222. """Returns (str): The LogDog viewer URL for the named stream.
  223. Args:
  224. name (str): The name of the stream. This can also be a query glob.
  225. Raises:
  226. KeyError: if information needed to construct the URL is missing.
  227. ValueError: if the stream prefix or name do not form a valid stream
  228. path.
  229. """
  230. if not self._coordinator_host:
  231. raise KeyError('Coordinator host is not configured')
  232. if not self._project:
  233. raise KeyError('Stream project is not configured')
  234. return streamname.get_logdog_viewer_url(
  235. self._coordinator_host,
  236. self._project,
  237. self.get_stream_path(name))
  238. def _register_new_stream(self, name):
  239. """Registers a new stream name.
  240. The Butler will internally reject any duplicate stream names. However, there
  241. isn't really feedback when this happens except a closed stream client. This
  242. is a client-side check to provide a more user-friendly experience in the
  243. event that a user attempts to register a duplicate stream name.
  244. Note that this is imperfect, as something else could register stream names
  245. with the same Butler instance and this library has no means of tracking.
  246. This is a best-effort experience, not a reliable check.
  247. Args:
  248. name (str): The name of the stream.
  249. Raises:
  250. ValueError if the stream name has already been registered.
  251. """
  252. with self._name_lock:
  253. if name in self._names:
  254. raise ValueError("Duplicate stream name [%s]" % (name,))
  255. self._names.add(name)
  256. @classmethod
  257. def _create(cls, value, **kwargs):
  258. """Returns (StreamClient): A new stream client instance.
  259. Validates the streamserver parameters and creates a new StreamClient
  260. instance that connects to them.
  261. Implementing classes must override this.
  262. """
  263. raise NotImplementedError()
  264. def _connect_raw(self):
  265. """Returns (file): A new file-like stream.
  266. Creates a new raw connection to the streamserver. This connection MUST not
  267. have any data written to it past initialization (if needed) when it has been
  268. returned.
  269. The file-like object must implement `write`, `fileno`, `flush`, and `close`.
  270. Implementing classes must override this.
  271. """
  272. raise NotImplementedError()
  273. def new_connection(self, params):
  274. """Returns (file): A new configured stream.
  275. The returned object implements (minimally) `write` and `close`.
  276. Creates a new LogDog stream with the specified parameters.
  277. Args:
  278. params (StreamParams): The parameters to use with the new connection.
  279. Raises:
  280. ValueError if the stream name has already been used, or if the parameters
  281. are not valid.
  282. """
  283. self._register_new_stream(params.name)
  284. params_bytes = params.to_json().encode('utf-8')
  285. fobj = self._connect_raw()
  286. fobj.write(BUTLER_MAGIC)
  287. varint.write_uvarint(fobj, len(params_bytes))
  288. fobj.write(params_bytes)
  289. return fobj
  290. @contextlib.contextmanager
  291. def text(self, name, **kwargs):
  292. """Context manager to create, use, and teardown a TEXT stream.
  293. This context manager creates a new butler TEXT stream with the specified
  294. parameters, yields it, and closes it on teardown.
  295. Args:
  296. name (str): the LogDog name of the stream.
  297. kwargs (dict): Log stream parameters. These may be any keyword arguments
  298. accepted by `open_text`.
  299. Returns (file): A file-like object to a Butler UTF-8 text stream supporting
  300. `write`.
  301. """
  302. fobj = None
  303. try:
  304. fobj = self.open_text(name, **kwargs)
  305. yield fobj
  306. finally:
  307. if fobj is not None:
  308. fobj.close()
  309. def open_text(self, name, content_type=None, tags=None):
  310. """Returns (file): A file-like object for a single text stream.
  311. This creates a new butler TEXT stream with the specified parameters.
  312. Args:
  313. name (str): the LogDog name of the stream.
  314. content_type (str): The optional content type of the stream. If None, a
  315. default content type will be chosen by the Butler.
  316. tags (dict): An optional key/value dictionary pair of LogDog stream tags.
  317. Returns (file): A file-like object to a Butler text stream. This object can
  318. have UTF-8 text content written to it with its `write` method, and must
  319. be closed when finished using its `close` method.
  320. """
  321. params = StreamParams.make(
  322. name=posixpath.join(self._namespace, name),
  323. type=StreamParams.TEXT,
  324. content_type=content_type,
  325. tags=tags)
  326. return self._TextStream(self, params, self.new_connection(params))
  327. @contextlib.contextmanager
  328. def binary(self, name, **kwargs):
  329. """Context manager to create, use, and teardown a BINARY stream.
  330. This context manager creates a new butler BINARY stream with the specified
  331. parameters, yields it, and closes it on teardown.
  332. Args:
  333. name (str): the LogDog name of the stream.
  334. kwargs (dict): Log stream parameters. These may be any keyword arguments
  335. accepted by `open_binary`.
  336. Returns (file): A file-like object to a Butler binary stream supporting
  337. `write`.
  338. """
  339. fobj = None
  340. try:
  341. fobj = self.open_binary(name, **kwargs)
  342. yield fobj
  343. finally:
  344. if fobj is not None:
  345. fobj.close()
  346. def open_binary(self, name, content_type=None, tags=None):
  347. """Returns (file): A file-like object for a single binary stream.
  348. This creates a new butler BINARY stream with the specified parameters.
  349. Args:
  350. name (str): the LogDog name of the stream.
  351. content_type (str): The optional content type of the stream. If None, a
  352. default content type will be chosen by the Butler.
  353. tags (dict): An optional key/value dictionary pair of LogDog stream tags.
  354. Returns (file): A file-like object to a Butler binary stream. This object
  355. can have UTF-8 content written to it with its `write` method, and must
  356. be closed when finished using its `close` method.
  357. """
  358. params = StreamParams.make(
  359. name=posixpath.join(self._namespace, name),
  360. type=StreamParams.BINARY,
  361. content_type=content_type,
  362. tags=tags)
  363. return self._BasicStream(self, params, self.new_connection(params))
  364. @contextlib.contextmanager
  365. def datagram(self, name, **kwargs):
  366. """Context manager to create, use, and teardown a DATAGRAM stream.
  367. This context manager creates a new butler DATAAGRAM stream with the
  368. specified parameters, yields it, and closes it on teardown.
  369. Args:
  370. name (str): the LogDog name of the stream.
  371. kwargs (dict): Log stream parameters. These may be any keyword arguments
  372. accepted by `open_datagram`.
  373. Returns (_DatagramStream): A datagram stream object. Datagrams can be
  374. written to it using its `send` method.
  375. """
  376. fobj = None
  377. try:
  378. fobj = self.open_datagram(name, **kwargs)
  379. yield fobj
  380. finally:
  381. if fobj is not None:
  382. fobj.close()
  383. def open_datagram(self, name, content_type=None, tags=None):
  384. """Creates a new butler DATAGRAM stream with the specified parameters.
  385. Args:
  386. name (str): the LogDog name of the stream.
  387. content_type (str): The optional content type of the stream. If None, a
  388. default content type will be chosen by the Butler.
  389. tags (dict): An optional key/value dictionary pair of LogDog stream tags.
  390. Returns (_DatagramStream): A datagram stream object. Datagrams can be
  391. written to it using its `send` method. This object must be closed when
  392. finished by using its `close` method.
  393. """
  394. params = StreamParams.make(
  395. name=posixpath.join(self._namespace, name),
  396. type=StreamParams.DATAGRAM,
  397. content_type=content_type,
  398. tags=tags)
  399. return self._DatagramStream(self, params, self.new_connection(params))
  400. class _NamedPipeStreamClient(StreamClient):
  401. """A StreamClient implementation that connects to a Windows named pipe.
  402. """
  403. def __init__(self, name, **kwargs):
  404. r"""Initializes a new Windows named pipe stream client.
  405. Args:
  406. name (str): The name of the Windows named pipe to use (e.g., "\\.\name")
  407. """
  408. super(_NamedPipeStreamClient, self).__init__(**kwargs)
  409. self._name = '\\\\.\\pipe\\' + name
  410. @classmethod
  411. def _create(cls, value, **kwargs):
  412. return cls(value, **kwargs)
  413. ERROR_PIPE_BUSY = 231
  414. def _connect_raw(self):
  415. # This is a similar procedure to the one in
  416. # https://github.com/microsoft/go-winio/blob/master/pipe.go (tryDialPipe)
  417. while True:
  418. try:
  419. return open(self._name, 'wb+', buffering=0)
  420. except (OSError, IOError):
  421. if GetLastError() != self.ERROR_PIPE_BUSY:
  422. raise
  423. time.sleep(0.001) # 1ms
  424. _default_registry.register_protocol('net.pipe', _NamedPipeStreamClient)
  425. class _UnixDomainSocketStreamClient(StreamClient):
  426. """A StreamClient implementation that uses a UNIX domain socket.
  427. """
  428. class SocketFile(object):
  429. """A write-only file-like object that writes to a UNIX socket."""
  430. def __init__(self, sock):
  431. self._sock = sock
  432. def fileno(self):
  433. return self._sock
  434. def write(self, data):
  435. self._sock.sendall(data)
  436. def flush(self):
  437. pass
  438. def close(self):
  439. self._sock.close()
  440. def __init__(self, path, **kwargs):
  441. """Initializes a new UNIX domain socket stream client.
  442. Args:
  443. path (str): The path to the named UNIX domain socket.
  444. """
  445. super(_UnixDomainSocketStreamClient, self).__init__(**kwargs)
  446. self._path = path
  447. @classmethod
  448. def _create(cls, value, **kwargs):
  449. if not os.path.exists(value):
  450. raise ValueError('UNIX domain socket [%s] does not exist.' % (value,))
  451. return cls(value, **kwargs)
  452. def _connect_raw(self):
  453. sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
  454. sock.connect(self._path)
  455. return self.SocketFile(sock)
  456. _default_registry.register_protocol('unix', _UnixDomainSocketStreamClient)