#
# Copyright (C) 2009-2016 CEA/DAM
# Copyright (C) 2016-2017 Stephane Thiell <sthiell@stanford.edu>
#
# This file is part of ClusterShell.
#
# ClusterShell is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# ClusterShell is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with ClusterShell; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
"""
EngineClient
ClusterShell engine's client interface.
An engine client is similar to a process, you can start/stop it, read data from
it and write data to it. Multiple data channels are supported (eg. stdin, stdout
and stderr, or even more...)
"""
import errno
import logging
import os
try:
import queue
except ImportError:
# Python 2 compatibility
import Queue as queue
import threading
from ClusterShell.Defaults import DEFAULTS
from ClusterShell.Worker.fastsubprocess import Popen, PIPE, STDOUT, \
set_nonblock_flag
from ClusterShell.Engine.Engine import EngineBaseTimer, E_READ, E_WRITE
LOGGER = logging.getLogger(__name__)
class EngineClientException(Exception):
"""Generic EngineClient exception."""
class EngineClientEOF(EngineClientException):
"""EOF from client."""
class EngineClientError(EngineClientException):
"""Base EngineClient error exception."""
class EngineClientNotSupportedError(EngineClientError):
"""Operation not supported by EngineClient."""
class EngineClientStream(object):
"""EngineClient I/O stream object.
Internal object used by EngineClient to manage its Engine-registered I/O
streams. Each EngineClientStream is bound to a file object (file
descriptor). It can be either an input, an output or a bidirectional
stream (not used for now)."""
def __init__(self, name, sfile=None, evmask=0):
"""Initialize an EngineClientStream object.
@param name: Name of stream.
@param sfile: File object or file descriptor.
@param evmask: Config I/O event bitmask.
"""
self.name = name
self.fd = None
self.rbuf = bytes()
self.wbuf = bytes()
self.eof = False
self.evmask = evmask
self.events = 0
self.new_events = 0
self.retain = False
self.closefd = False
self.set_file(sfile)
def set_file(self, sfile, evmask=0, retain=True, closefd=True):
"""
Set the stream file and event mask for this object.
sfile should be a file object or a file descriptor.
Event mask can be either E_READ, E_WRITE or both.
Currently does NOT retain file object.
"""
try:
# file descriptor
self.fd = sfile.fileno()
except AttributeError:
self.fd = sfile
# Set I/O event mask
self.evmask = evmask
# Set retain flag
self.retain = retain
# Set closefd flag
self.closefd = closefd
def __repr__(self):
return "<%s at 0x%s (name=%s fd=%s rbuflen=%d wbuflen=%d eof=%d " \
"evmask=0x%x)>" % (self.__class__.__name__, id(self), self.name,
self.fd, len(self.rbuf), len(self.wbuf), self.eof, self.evmask)
def close(self):
"""Close stream."""
if self.closefd and self.fd is not None:
os.close(self.fd)
def readable(self):
"""Return whether the stream is setup as readable."""
return self.evmask & E_READ
def writable(self):
"""Return whether the stream is setup as writable."""
return self.evmask & E_WRITE
class EngineClientStreamDict(dict):
"""EngineClient's named stream dictionary."""
def set_stream(self, sname, sfile=None, evmask=0, retain=True,
closefd=True):
"""Set stream based on file object or file descriptor.
This method can be used to add a stream or update its
parameters.
"""
engfile = dict.setdefault(self, sname, EngineClientStream(sname))
engfile.set_file(sfile, evmask, retain, closefd)
return engfile
def set_reader(self, sname, sfile=None, retain=True, closefd=True):
"""Set readable stream based on file object or file descriptor."""
self.set_stream(sname, sfile, E_READ, retain, closefd)
def set_writer(self, sname, sfile=None, retain=True, closefd=True):
"""Set writable stream based on file object or file descriptor."""
self.set_stream(sname, sfile, E_WRITE, retain, closefd)
def destroy(self, key):
"""Close file object and remove it from this pool."""
self[key].close()
dict.pop(self, key)
def __delitem__(self, key):
self.destroy(key)
def clear(self):
"""Clear File Pool"""
for stream in self.values():
stream.close()
dict.clear(self)
def active_readers(self):
"""Get an iterator on readable streams (with fd set)."""
return (s for s in self.readers() if s.fd is not None)
def readers(self):
"""Get an iterator on all streams setup as readable."""
return (s for s in list(self.values()) if s.evmask & E_READ)
def active_writers(self):
"""Get an iterator on writable streams (with fd set)."""
return (s for s in self.writers() if s.fd is not None)
def writers(self):
"""Get an iterator on all streams setup as writable."""
return (s for s in list(self.values()) if s.evmask & E_WRITE)
def retained(self):
"""Check whether this set of streams is retained.
Note on retain: an active stream with retain=True keeps the
engine client alive. When only streams with retain=False
remain, the engine client terminates.
Return:
True -- when at least one stream is retained
False -- when no retainable stream remain
"""
for stream in self.values():
if stream.fd is not None and stream.retain:
return True
return False
class EngineClient(EngineBaseTimer):
"""
Abstract class EngineClient.
"""
def __init__(self, worker, key, stderr, timeout, autoclose):
"""EngineClient initializer.
Should be called from derived classes.
Arguments:
worker -- parent worker instance
key -- client key used by MsgTree (eg. node name)
stderr -- boolean set if stderr is on a separate stream
timeout -- client execution timeout value (float)
autoclose -- boolean set to indicate whether this engine
client should be aborted as soon as all other
non-autoclosing clients have finished.
"""
EngineBaseTimer.__init__(self, timeout, -1, autoclose)
self._reg_epoch = 0 # registration generation number
# read-only public
self.registered = False # registered on engine or not
self.delayable = True # subject to fanout limit
self.worker = worker
if key is None:
key = id(worker)
self.key = key
# boolean indicating whether stderr is on a separate fd
self._stderr = stderr
# streams associated with this client
self.streams = EngineClientStreamDict()
def __repr__(self):
# added repr(self.key)
return '<%s.%s instance at 0x%x key %r>' % (self.__module__,
self.__class__.__name__,
id(self), self.key)
def _fire(self):
"""
Fire timeout timer.
"""
if self._engine:
self._engine.remove(self, abort=True, did_timeout=True)
def _start(self):
"""
Starts client and returns client instance as a convenience.
Derived classes must implement.
"""
raise NotImplementedError("Derived classes must implement.")
def _close(self, abort, timeout):
"""
Close client. Called by the engine after client has been unregistered.
This method should handle both termination types (normal or aborted)
and should set timeout status accordingly.
Derived classes should implement.
"""
for sname in list(self.streams):
# for engine-induced timeout with abort, do not actually abort
# the stream so we can flush buffers (original behavior)
self._close_stream(sname, abort and not timeout)
self.invalidate() # set self._engine to None
def _close_stream(self, sname, abort=False):
"""
Close specific stream by name (internal, called by engine). This method
is the regular way to close a stream flushing read buffers accordingly,
unless we are closing the stream with the abort flag set to True.
"""
# flush_read() is useful but may generate user events (ev_read) that
# could lead to worker abort and then ev_close. Be careful there.
# If we are doing a close with abort, skip the flush (this is newer).
if not abort:
self._flush_read(sname)
if sname in self.streams:
del self.streams[sname]
def _set_reading(self, sname):
"""
Set reading state.
"""
self._engine.set_reading(self, sname)
def _set_writing(self, sname):
"""
Set writing state.
"""
self._engine.set_writing(self, sname)
def _read(self, sname, size=65536):
"""
Read data from process.
"""
result = os.read(self.streams[sname].fd, size)
if len(result) == 0:
raise EngineClientEOF()
self._set_reading(sname)
return result
def _flush_read(self, sname):
"""Called when stream is closing to flush read buffers."""
pass # derived classes may implement
def _handle_read(self, sname):
"""
Handle a read notification. Called by the engine as the result of an
event indicating that a read is available.
"""
raise NotImplementedError("Derived classes must implement.")
def _handle_write(self, sname):
"""
Handle a write notification. Called by the engine as the result of an
event indicating that a write can be performed now.
"""
wfile = self.streams[sname]
if not wfile.wbuf and wfile.eof:
# remove stream from engine (not directly)
if self._engine:
self._engine.remove_stream(self, wfile)
elif len(wfile.wbuf) > 0:
try:
wcnt = os.write(wfile.fd, wfile.wbuf)
except OSError as exc:
if exc.errno == errno.EAGAIN:
# _handle_write() is not only called by the engine but also
# by _write(), so this is legit: we just try again later
self._set_writing(sname)
return
if exc.errno == errno.EPIPE:
# broken pipe: log warning message and do NOT retry
LOGGER.warning('%r: %s', self, exc)
# remove stream now; select engine has no POLLERR
if self._engine:
self._engine.remove_stream(self, wfile)
return
raise
if wcnt > 0:
# dequeue written buffer
wfile.wbuf = wfile.wbuf[wcnt:]
# check for possible ending
if wfile.eof and not wfile.wbuf:
self.worker._on_written(self.key, wcnt, sname)
# remove stream from engine (not directly)
if self._engine:
self._engine.remove_stream(self, wfile)
else:
self._set_writing(sname)
self.worker._on_written(self.key, wcnt, sname)
def _exec_nonblock(self, commandlist, shell=False, env=None):
"""
Utility method to launch a command with stdin/stdout file
descriptors configured in non-blocking mode.
"""
full_env = None
if env:
full_env = os.environ.copy()
full_env.update(env)
if self._stderr:
stderr_setup = PIPE
else:
stderr_setup = STDOUT
# Launch process in non-blocking mode
proc = Popen(commandlist, bufsize=0, stdin=PIPE, stdout=PIPE,
stderr=stderr_setup, shell=shell, env=full_env)
if self._stderr:
self.streams.set_stream(self.worker.SNAME_STDERR, proc.stderr,
E_READ)
self.streams.set_stream(self.worker.SNAME_STDOUT, proc.stdout, E_READ)
self.streams.set_stream(self.worker.SNAME_STDIN, proc.stdin, E_WRITE,
retain=False)
return proc
def _readlines(self, sname):
"""Utility method to read client lines."""
# read a chunk of data, may raise eof
readbuf = self._read(sname)
assert len(readbuf) > 0, "assertion failed: len(readbuf) > 0"
# Current version implements line-buffered reads. If needed, we could
# easily provide direct, non-buffered, data reads in the future.
rfile = self.streams[sname]
buf = rfile.rbuf + readbuf
lines = buf.splitlines(True)
rfile.rbuf = bytes()
for line in lines:
if line.endswith(b'\n'):
if line.endswith(b'\r\n'):
yield line[:-2] # trim CRLF
else:
# trim LF
yield line[:-1] # trim LF
else:
# keep partial line in buffer
rfile.rbuf = line
# breaking here
def _write(self, sname, buf):
"""Add some data to be written to the client."""
if sname not in self.streams:
LOGGER.debug("dropping write to unknown or closed stream %s on "
"client %s", sname, self.key)
return
wfile = self.streams[sname]
if self._engine and wfile.fd:
wfile.wbuf += buf
# give it a try now (will set writing flag anyhow)
self._handle_write(sname)
else:
# bufferize until pipe is ready
wfile.wbuf += buf
def _set_write_eof(self, sname):
"""Set EOF on specific writable stream."""
if sname not in self.streams:
LOGGER.debug("stream %s was already closed on client %s, skipping",
sname, self.key)
return
wfile = self.streams[sname]
wfile.eof = True
if self._engine and wfile.fd and not wfile.wbuf:
# sendq empty, remove stream now
self._engine.remove_stream(self, wfile)
def abort(self):
"""Abort processing any action by this client.
Safe to call on an already closing or aborting client.
"""
engine = self._engine
if engine:
self.invalidate() # set self._engine to None
engine.remove(self, abort=True)
[docs]
class EnginePort(EngineClient):
"""
An EnginePort is an abstraction object to deliver messages
reliably between tasks.
"""
class _Msg(object):
"""Private class representing a port message.
A port message may be any Python object.
"""
def __init__(self, user_msg, sync):
self._user_msg = user_msg
self._sync_msg = sync
self.reply_lock = threading.Lock()
self.reply_lock.acquire()
def get(self):
"""
Get and acknowledge message.
"""
self.reply_lock.release()
return self._user_msg
def sync(self):
"""
Wait for message acknowledgment if needed.
"""
if self._sync_msg:
self.reply_lock.acquire()
[docs]
def __init__(self, handler=None, autoclose=False):
"""
Initialize EnginePort object.
"""
EngineClient.__init__(self, None, None, False, -1, autoclose)
self.eh = handler
# ports are no subject to fanout
self.delayable = False
# Port messages queue
self._msgq = queue.Queue(DEFAULTS.port_qlimit)
# Request pipe
(readfd, writefd) = os.pipe()
# Set nonblocking flag
set_nonblock_flag(readfd)
set_nonblock_flag(writefd)
self.streams.set_stream('in', readfd, E_READ)
self.streams.set_stream('out', writefd, E_WRITE)
[docs]
def __repr__(self):
try:
fd_in = self.streams['in'].fd
except KeyError:
fd_in = None
try:
fd_out = self.streams['out'].fd
except KeyError:
fd_out = None
return "<%s at 0x%s (streams=(%s, %s))>" % (self.__class__.__name__,
id(self), fd_in, fd_out)
def _start(self):
"""Start port."""
if self.eh is not None:
self.eh.ev_port_start(self)
return self
def _close(self, abort, timeout):
"""Close port."""
if not self._msgq.empty():
# purge msgq
try:
while not self._msgq.empty():
pmsg = self._msgq.get(block=False)
LOGGER.debug('%r: dropped msg: %s', self, pmsg.get())
except queue.Empty:
pass
self._msgq = None
del self.streams['out']
del self.streams['in']
self.invalidate()
def _handle_read(self, sname):
"""
Handle a read notification. Called by the engine as the result of an
event indicating that a read is available.
"""
readbuf = self._read(sname, 4096)
for dummy_char in readbuf:
# raise Empty if empty (should never happen)
pmsg = self._msgq.get(block=False)
self.eh.ev_msg(self, pmsg.get())
[docs]
def msg(self, send_msg, send_once=False):
"""
Port message send method that will wait for acknowledgement
unless the send_once parameter if set.
May be called from another thread. Will generate ev_msg() on
Port event handler (in Port task/thread).
Return False if the message cannot be sent (eg. port closed).
"""
if self._msgq is None: # called after port closed?
return False
pmsg = EnginePort._Msg(send_msg, not send_once)
self._msgq.put(pmsg, block=True, timeout=None)
try:
ret = os.write(self.streams['out'].fd, b'M')
except OSError:
raise
pmsg.sync()
return ret == 1
[docs]
def msg_send(self, send_msg):
"""
Port message send-once method (no acknowledgement). See msg().
Return False if the message cannot be sent (eg. port closed).
"""
return self.msg(send_msg, send_once=True)