summaryrefslogtreecommitdiffstats
path: root/scripts/powerline-daemon
diff options
context:
space:
mode:
Diffstat (limited to '')
-rwxr-xr-xscripts/powerline-daemon495
1 files changed, 495 insertions, 0 deletions
diff --git a/scripts/powerline-daemon b/scripts/powerline-daemon
new file mode 100755
index 0000000..f15ac69
--- /dev/null
+++ b/scripts/powerline-daemon
@@ -0,0 +1,495 @@
+#!/usr/bin/env python
+# vim:fileencoding=utf-8:noet
+from __future__ import (unicode_literals, division, absolute_import, print_function)
+
+import socket
+import os
+import errno
+import sys
+import fcntl
+import atexit
+import stat
+
+from argparse import ArgumentParser
+from select import select
+from signal import signal, SIGTERM
+from time import sleep
+from functools import partial
+from io import BytesIO
+from threading import Event
+from itertools import chain
+from logging import StreamHandler
+
+from powerline.shell import ShellPowerline
+from powerline.commands.main import finish_args, write_output
+from powerline.lib.monotonic import monotonic
+from powerline.lib.encoding import get_preferred_output_encoding, get_preferred_arguments_encoding, get_unicode_writer
+from powerline.bindings.wm import wm_threads
+
+from powerline.commands.main import get_argparser as get_main_argparser
+from powerline.commands.daemon import get_argparser as get_daemon_argparser
+
+
+USE_FILESYSTEM = not sys.platform.lower().startswith('linux')
+
+
+class NonInteractiveArgParser(ArgumentParser):
+ def print_usage(self, file=None):
+ raise Exception(self.format_usage())
+
+ def print_help(self, file=None):
+ raise Exception(self.format_help())
+
+ def exit(self, status=0, message=None):
+ pass
+
+ def error(self, message):
+ raise Exception(self.format_usage())
+
+
+EOF = b'EOF\0\0'
+
+
+class State(object):
+ __slots__ = ('powerlines', 'logger', 'config_loader', 'started_wm_threads',
+ 'ts_shutdown_event')
+
+ def __init__(self, **kwargs):
+ self.logger = None
+ self.config_loader = None
+ self.started_wm_threads = {}
+ self.powerlines = {}
+ self.ts_shutdown_event = Event()
+
+
+HOME = os.path.expanduser('~')
+
+
+class NonDaemonShellPowerline(ShellPowerline):
+ def get_log_handler(self):
+ return StreamHandler()
+
+
+def start_wm(args, environ, cwd, is_daemon, state):
+ wm_name = args.ext[0][3:]
+ if wm_name in state.started_wm_threads:
+ return b''
+ thread_shutdown_event = Event()
+ thread = wm_threads[wm_name](
+ thread_shutdown_event=thread_shutdown_event,
+ pl_shutdown_event=state.ts_shutdown_event,
+ pl_config_loader=state.config_loader,
+ )
+ thread.start()
+ state.started_wm_threads[wm_name] = (thread, thread_shutdown_event)
+ return b''
+
+
+def render(args, environ, cwd, is_daemon, state):
+ segment_info = {
+ 'getcwd': lambda: cwd,
+ 'home': environ.get('HOME', HOME),
+ 'environ': environ,
+ 'args': args,
+ }
+ key = (
+ args.ext[0],
+ args.renderer_module,
+ tuple(args.config_override) if args.config_override else None,
+ tuple(args.theme_override) if args.theme_override else None,
+ tuple(args.config_path) if args.config_path else None,
+ environ.get('POWERLINE_THEME_OVERRIDES', ''),
+ environ.get('POWERLINE_CONFIG_OVERRIDES', ''),
+ environ.get('POWERLINE_CONFIG_PATHS', ''),
+ )
+
+ PowerlineClass = ShellPowerline if is_daemon else NonDaemonShellPowerline
+ powerline = None
+ try:
+ powerline = state.powerlines[key]
+ except KeyError:
+ try:
+ powerline = state.powerlines[key] = PowerlineClass(
+ args,
+ logger=state.logger,
+ config_loader=state.config_loader,
+ run_once=False,
+ shutdown_event=state.ts_shutdown_event,
+ )
+ if state.logger is None:
+ state.logger = powerline.logger
+ if state.config_loader is None:
+ state.config_loader = powerline.config_loader
+ except SystemExit:
+ # Somebody thought raising system exit was a good idea,
+ return ''
+ except Exception as e:
+ if powerline:
+ powerline.pl.exception('Failed to render {0}: {1}', str(key), str(e))
+ else:
+ return 'Failed to render {0}: {1}'.format(str(key), str(e))
+ s = BytesIO()
+ write_output(args, powerline, segment_info, get_unicode_writer(stream=s))
+ s.seek(0)
+ return s.read()
+
+
+def eintr_retry_call(func, *args, **kwargs):
+ while True:
+ try:
+ return func(*args, **kwargs)
+ except EnvironmentError as e:
+ if getattr(e, 'errno', None) == errno.EINTR:
+ continue
+ raise
+
+
+def do_read(conn, timeout=2.0):
+ ''' Read data from the client. If the client fails to send data within
+ timeout seconds, abort. '''
+ read = []
+ end_time = monotonic() + timeout
+ while not read or not read[-1].endswith(b'\0\0'):
+ r, w, e = select((conn,), (), (conn,), timeout)
+ if e:
+ return
+ if monotonic() > end_time:
+ return
+ if not r:
+ continue
+ x = eintr_retry_call(conn.recv, 4096)
+ if x:
+ read.append(x)
+ else:
+ break
+ return b''.join(read)
+
+
+def do_write(conn, result):
+ try:
+ eintr_retry_call(conn.sendall, result)
+ except Exception:
+ pass
+
+
+def safe_bytes(o, encoding=get_preferred_output_encoding()):
+ '''Return bytes instance without ever throwing an exception.'''
+ try:
+ try:
+ # We are assuming that o is a unicode object
+ return o.encode(encoding, 'replace')
+ except Exception:
+ # Object may have defined __bytes__ (python 3) or __str__ method
+ # (python 2)
+ # This also catches problem with non_ascii_bytes.encode('utf-8')
+ # that first tries to decode to UTF-8 using ascii codec (and fails
+ # in this case) and then encode to given encoding: errors= argument
+ # is not used in the first stage.
+ return bytes(o)
+ except Exception as e:
+ return safe_bytes(str(e), encoding)
+
+
+def parse_args(req, parser, encoding=get_preferred_arguments_encoding()):
+ args = [x.decode(encoding) for x in req.split(b'\0') if x]
+ numargs = int(args[0], 16)
+ shell_args = parser.parse_args(args[1:numargs + 1])
+ cwd = args[numargs + 1]
+ environ = dict(((k, v) for k, v in (x.partition('=')[0::2] for x in args[numargs + 2:])))
+ cwd = cwd or environ.get('PWD', '/')
+ return shell_args, environ, cwd
+
+
+def get_answer(req, is_daemon, argparser, state):
+ try:
+ args, environ, cwd = parse_args(req, argparser)
+ finish_args(argparser, environ, args, is_daemon=True)
+ if args.ext[0].startswith('wm.'):
+ return safe_bytes(start_wm(args, environ, cwd, is_daemon, state))
+ else:
+ return safe_bytes(render(args, environ, cwd, is_daemon, state))
+ except Exception as e:
+ return safe_bytes(str(e))
+
+
+def do_one(sock, read_sockets, write_sockets, result_map, is_daemon, argparser,
+ state):
+ r, w, e = select(
+ tuple(read_sockets) + (sock,),
+ tuple(write_sockets),
+ tuple(read_sockets) + tuple(write_sockets) + (sock,),
+ 60.0
+ )
+
+ if sock in e:
+ # We cannot accept any more connections, so we exit
+ raise SystemExit(1)
+
+ for s in e:
+ # Discard all broken connections to clients
+ s.close()
+ read_sockets.discard(s)
+ write_sockets.discard(s)
+
+ for s in r:
+ if s == sock:
+ # A client wants to connect
+ conn, _ = eintr_retry_call(sock.accept)
+ read_sockets.add(conn)
+ else:
+ # A client has sent some data
+ read_sockets.discard(s)
+ req = do_read(s)
+ if req == EOF:
+ raise SystemExit(0)
+ elif req:
+ ans = get_answer(req, is_daemon, argparser, state)
+ result_map[s] = ans
+ write_sockets.add(s)
+ else:
+ s.close()
+
+ for s in w:
+ # A client is ready to receive the result
+ write_sockets.discard(s)
+ result = result_map.pop(s)
+ try:
+ do_write(s, result)
+ finally:
+ s.close()
+
+
+def shutdown(sock, read_sockets, write_sockets, state):
+ '''Perform operations necessary for nicely shutting down daemon
+
+ Specifically it
+
+ #. Closes all sockets.
+ #. Notifies segments based on
+ :py:class:`powerline.lib.threaded.ThreadedSegment` and WM-specific
+ threads that daemon is shutting down.
+ #. Waits for threads to finish, but no more then 2 seconds total.
+ #. Waits so that total execution time of this function is 2 seconds in order
+ to allow ThreadedSegments to finish.
+ '''
+ total_wait_time = 2
+ shutdown_start_time = monotonic()
+
+ for s in chain((sock,), read_sockets, write_sockets):
+ s.close()
+
+ # Notify ThreadedSegments
+ state.ts_shutdown_event.set()
+ for thread, shutdown_event in state.started_wm_threads.values():
+ shutdown_event.set()
+
+ for thread, shutdown_event in state.started_wm_threads.values():
+ wait_time = total_wait_time - (monotonic() - shutdown_start_time)
+ if wait_time > 0:
+ thread.join(wait_time)
+
+ wait_time = total_wait_time - (monotonic() - shutdown_start_time)
+ sleep(wait_time)
+
+
+def main_loop(sock, is_daemon):
+ sock.listen(128)
+ sock.setblocking(0)
+
+ read_sockets, write_sockets = set(), set()
+ result_map = {}
+ parser = get_main_argparser(NonInteractiveArgParser)
+ state = State()
+ try:
+ try:
+ while True:
+ do_one(
+ sock, read_sockets, write_sockets, result_map,
+ is_daemon=is_daemon,
+ argparser=parser,
+ state=state,
+ )
+ except KeyboardInterrupt:
+ raise SystemExit(0)
+ except SystemExit as e:
+ shutdown(sock, read_sockets, write_sockets, state)
+ raise e
+ return 0
+
+
+def daemonize(stdin=os.devnull, stdout=os.devnull, stderr=os.devnull):
+ try:
+ pid = os.fork()
+ if pid > 0:
+ # exit first parent
+ raise SystemExit(0)
+ except OSError as e:
+ sys.stderr.write("fork #1 failed: %d (%s)\n" % (e.errno, e.strerror))
+ raise SystemExit(1)
+
+ # decouple from parent environment
+ os.chdir("/")
+ os.setsid()
+ os.umask(0)
+
+ # do second fork
+ try:
+ pid = os.fork()
+ if pid > 0:
+ # exit from second parent
+ raise SystemExit(0)
+ except OSError as e:
+ sys.stderr.write("fork #2 failed: %d (%s)\n" % (e.errno, e.strerror))
+ raise SystemExit(1)
+
+ # Redirect standard file descriptors.
+ si = open(stdin, 'rb')
+ so = open(stdout, 'a+b')
+ se = open(stderr, 'a+b', 0)
+ os.dup2(si.fileno(), sys.stdin.fileno())
+ os.dup2(so.fileno(), sys.stdout.fileno())
+ os.dup2(se.fileno(), sys.stderr.fileno())
+ return True
+
+
+def check_existing(address):
+ if USE_FILESYSTEM:
+ # We cannot bind if the socket file already exists so remove it, we
+ # already have a lock on pidfile, so this should be safe.
+ try:
+ os.unlink(address)
+ except EnvironmentError:
+ pass
+
+ sock = socket.socket(family=socket.AF_UNIX)
+ try:
+ sock.bind(address)
+ except socket.error as e:
+ if getattr(e, 'errno', None) == errno.EADDRINUSE:
+ return None
+ raise
+ return sock
+
+
+def kill_daemon(address):
+ sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+ try:
+ try:
+ eintr_retry_call(sock.connect, address)
+ except socket.error:
+ return False
+ else:
+ eintr_retry_call(sock.sendall, EOF)
+ finally:
+ sock.close()
+ return True
+
+
+def cleanup_lockfile(pidfile, fd, *args):
+ try:
+ # Remove the directory entry for the lock file
+ os.unlink(pidfile)
+ # Close the file descriptor
+ os.close(fd)
+ except EnvironmentError:
+ pass
+ if args:
+ # Called in signal handler
+ raise SystemExit(1)
+
+
+def lockpidfile(pidfile):
+ fd = os.open(
+ pidfile,
+ os.O_WRONLY | os.O_CREAT,
+ stat.S_IRUSR | stat.S_IWUSR | stat.S_IRGRP | stat.S_IROTH
+ )
+ try:
+ fcntl.lockf(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+ except EnvironmentError:
+ os.close(fd)
+ return None
+ os.lseek(fd, 0, os.SEEK_SET)
+ os.ftruncate(fd, 0)
+ os.write(fd, ('%d' % os.getpid()).encode('ascii'))
+ os.fsync(fd)
+ cleanup = partial(cleanup_lockfile, pidfile, fd)
+ signal(SIGTERM, cleanup)
+ atexit.register(cleanup)
+ return fd
+
+
+def main():
+ parser = get_daemon_argparser()
+ args = parser.parse_args()
+ is_daemon = False
+ address = None
+ pidfile = None
+
+ if args.socket:
+ address = args.socket
+ if not USE_FILESYSTEM:
+ address = '\0' + address
+ else:
+ if USE_FILESYSTEM:
+ address = '/tmp/powerline-ipc-%d'
+ else:
+ # Use the abstract namespace for sockets rather than the filesystem
+ # (Available only in linux)
+ address = '\0powerline-ipc-%d'
+
+ address = address % os.getuid()
+
+ if USE_FILESYSTEM:
+ pidfile = address + '.pid'
+
+ if args.kill:
+ if args.foreground or args.replace:
+ parser.error('--kill and --foreground/--replace cannot be used together')
+ if kill_daemon(address):
+ if not args.quiet:
+ print ('Kill command sent to daemon, if it does not die in a couple of seconds use kill to kill it')
+ raise SystemExit(0)
+ else:
+ if not args.quiet:
+ print ('No running daemon found')
+ raise SystemExit(1)
+
+ if args.replace:
+ while kill_daemon(address):
+ if not args.quiet:
+ print ('Kill command sent to daemon, waiting for daemon to exit, press Ctrl-C to terminate wait and exit')
+ sleep(2)
+
+ if USE_FILESYSTEM and not args.foreground:
+ # We must daemonize before creating the locked pidfile, unfortunately,
+ # this means further print statements are discarded
+ is_daemon = daemonize()
+
+ if USE_FILESYSTEM:
+ # Create a locked pid file containing the daemon’s PID
+ if lockpidfile(pidfile) is None:
+ if not args.quiet:
+ sys.stderr.write(
+ 'The daemon is already running. Use %s -k to kill it.\n' % (
+ os.path.basename(sys.argv[0])))
+ raise SystemExit(1)
+
+ # Bind to address or bail if we cannot bind
+ sock = check_existing(address)
+ if sock is None:
+ if not args.quiet:
+ sys.stderr.write(
+ 'The daemon is already running. Use %s -k to kill it.\n' % (
+ os.path.basename(sys.argv[0])))
+ raise SystemExit(1)
+
+ if not USE_FILESYSTEM and not args.foreground:
+ # We daemonize on linux
+ is_daemon = daemonize()
+
+ return main_loop(sock, is_daemon)
+
+
+if __name__ == '__main__':
+ main()