gsh/stdin.py

   1 # This program is free software; you can redistribute it and/or modify
   2 # it under the terms of the GNU General Public License as published by
   3 # the Free Software Foundation; either version 2 of the License, or
   4 # (at your option) any later version.
   5 #
   6 # This program is distributed in the hope that it will be useful,
   7 # but WITHOUT ANY WARRANTY; without even the implied warranty of
   8 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   9 # GNU Library General Public License for more details.
  10 #
  11 # You should have received a copy of the GNU General Public License
  12 # along with this program; if not, write to the Free Software
  13 # Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
  14 #
  15 # See the COPYING file for license information.
  16 #
  17 # Copyright (c) 2006, 2007, 2008 Guillaume Chazarain <guichaz@gmail.com>
  18
  19 import asyncore
  20 import errno
  21 import os
  22 import readline # Just to say we want to use it with raw_input
  23 import signal
  24 import socket
  25 import subprocess
  26 import sys
  27 import tempfile
  28 import termios
  29 from threading import Thread, Event, Lock
  30
  31 from gsh import dispatchers, remote_dispatcher
  32 from gsh.console import console_output, set_last_status_length
  33 from gsh import completion
  34
  35 class input_buffer(object):
  36     """The shared input buffer between the main thread and the stdin thread"""
  37     def __init__(self):
  38         self.lock = Lock()
  39         self.buf = ''
  40
  41     def add(self, data):
  42         """Add data to the buffer"""
  43         self.lock.acquire()
  44         try:
  45             self.buf += data
  46         finally:
  47             self.lock.release()
  48
  49     def get(self):
  50         """Get the content of the buffer"""
  51         self.lock.acquire()
  52         try:
  53             data = self.buf
  54             if data:
  55                 self.buf = ''
  56                 return data
  57         finally:
  58             self.lock.release()
  59
  60 def ignore_sigchld(ignore):
  61     """Typically we don't want to create zombie. But when executing a user
  62     command (!command) the subprocess module relies on zombies not being
  63     automatically reclaimed"""
  64     if ignore:
  65         signal.signal(signal.SIGCHLD, signal.SIG_IGN)
  66         # Reclaim previously created zombies
  67         try:
  68             while os.waitpid(-1, os.WNOHANG) != (0, 0):
  69                 pass
  70         except OSError, e:
  71             if e.errno != errno.ECHILD:
  72                 raise
  73     else:
  74         signal.signal(signal.SIGCHLD, signal.SIG_DFL)
  75
  76 def process_input_buffer():
  77     """Send the content of the input buffer to all remote processes, this must
  78     be called in the main thread"""
  79     from gsh.control_commands_helpers import handle_control_command
  80     data = the_stdin_thread.input_buffer.get()
  81     if not data:
  82         return
  83
  84     remote_dispatcher.log('> ' + data)
  85
  86     if data.startswith(':'):
  87         handle_control_command(data[1:-1])
  88         return
  89
  90     if data.startswith('!'):
  91         ignore_sigchld(False)
  92         try:
  93             retcode = subprocess.call(data[1:], shell=True)
  94         except OSError, e:
  95             if e.errno == errno.EINTR:
  96                 console_output('Child was interrupted\n')
  97                 retcode = 0
  98             else:
  99                 raise
 100         ignore_sigchld(True)
 101         if retcode > 0:
 102             console_output('Child returned %d\n' % retcode)
 103         elif retcode < 0:
 104             console_output('Child was terminated by signal %d\n' % -retcode)
 105         return
 106
 107     for r in dispatchers.all_instances():
 108         try:
 109             r.dispatch_command(data)
 110         except asyncore.ExitNow, e:
 111             raise e
 112         except Exception, msg:
 113             console_output('%s for %s, disconnecting\n' % (msg, r.display_name))
 114             r.disconnect()
 115         else:
 116             if r.enabled and r.state is remote_dispatcher.STATE_IDLE:
 117                 r.change_state(remote_dispatcher.STATE_RUNNING)
 118
 119 # The stdin thread uses a synchronous (with ACK) socket to communicate with the
 120 # main thread, which is most of the time waiting in the poll() loop.
 121 # Socket character protocol:
 122 # d: there is new data to send
 123 # A: ACK, same reply for every message, communications are synchronous, so the
 124 # stdin thread sends a character to the socket, the main thread processes it,
 125 # sends the ACK, and the stdin thread can go on.
 126
 127 class socket_notification_reader(asyncore.dispatcher):
 128     """The socket reader in the main thread"""
 129     def __init__(self):
 130         asyncore.dispatcher.__init__(self, the_stdin_thread.socket_read)
 131
 132     def _do(self, c):
 133         if c == 'd':
 134             process_input_buffer()
 135         else:
 136             raise Exception, 'Unknown code: %s' % (c)
 137
 138     def handle_read(self):
 139         """Handle all the available character commands in the socket"""
 140         while True:
 141             try:
 142                 c = self.recv(1)
 143             except socket.error, why:
 144                 if why[0] == errno.EWOULDBLOCK:
 145                     return
 146                 else:
 147                     raise
 148             else:
 149                 self._do(c)
 150                 self.socket.setblocking(True)
 151                 self.send('A')
 152                 self.socket.setblocking(False)
 153
 154     def writable(self):
 155         """Our writes are blocking"""
 156         return False
 157
 158 def write_main_socket(c):
 159     """Synchronous write to the main socket, wait for ACK"""
 160     the_stdin_thread.socket_write.send(c)
 161     while True:
 162         try:
 163             the_stdin_thread.socket_write.recv(1)
 164         except socket.error, e:
 165             assert e[0] == errno.EINTR
 166         else:
 167             break
 168
 169 #
 170 # This file descriptor is used to interrupt readline in raw_input().
 171 # /dev/null is not enough as it does not get out of a 'Ctrl-R' reverse-i-search.
 172 # A Ctrl-C seems to make raw_input() return in all cases, and avoids printing
 173 # a newline
 174 tempfile_fd, tempfile_name = tempfile.mkstemp()
 175 os.remove(tempfile_name)
 176 os.write(tempfile_fd, chr(3))
 177
 178 def get_stdin_pid(cached_result=None):
 179     """Try to get the PID of the stdin thread, otherwise get the whole process
 180     ID"""
 181     if cached_result is None:
 182         try:
 183             tasks = os.listdir('/proc/self/task')
 184         except OSError, e:
 185             if e.errno != errno.ENOENT:
 186                 raise
 187             cached_result = os.getpid()
 188         else:
 189             tasks.remove(str(os.getpid()))
 190             assert len(tasks) == 1
 191             cached_result = int(tasks[0])
 192     return cached_result
 193
 194 def interrupt_stdin_thread():
 195     """The stdin thread may be in raw_input(), get out of it"""
 196     dupped_stdin = os.dup(0) # Backup the stdin fd
 197     assert not the_stdin_thread.interrupt_asked # Sanity check
 198     the_stdin_thread.interrupt_asked = True # Not user triggered
 199     os.lseek(tempfile_fd, 0, 0) # Rewind in the temp file
 200     os.dup2(tempfile_fd, 0) # This will make raw_input() return
 201     pid = get_stdin_pid()
 202     os.kill(pid, signal.SIGWINCH) # Try harder to wake up raw_input()
 203     the_stdin_thread.out_of_raw_input.wait() # Wait for this return
 204     the_stdin_thread.interrupt_asked = False # Restore sanity
 205     os.dup2(dupped_stdin, 0) # Restore stdin
 206     os.close(dupped_stdin) # Cleanup
 207
 208 echo_enabled = True
 209 def set_echo(echo):
 210     global echo_enabled
 211     if echo != echo_enabled:
 212         fd = sys.stdin.fileno()
 213         attr = termios.tcgetattr(fd)
 214         if echo:
 215             attr[3] |= termios.ECHO
 216         else:
 217             attr[3] &= ~termios.ECHO
 218         termios.tcsetattr(fd, termios.TCSANOW, attr)
 219         echo_enabled = echo
 220
 221 class stdin_thread(Thread):
 222     """The stdin thread, used to call raw_input()"""
 223     def __init__(self):
 224         Thread.__init__(self, name='stdin thread')
 225         completion.install_completion_handler()
 226
 227     @staticmethod
 228     def activate(interactive):
 229         """Activate the thread at initialization time"""
 230         the_stdin_thread.input_buffer = input_buffer()
 231         if interactive:
 232             the_stdin_thread.raw_input_wanted = Event()
 233             the_stdin_thread.in_raw_input = Event()
 234             the_stdin_thread.out_of_raw_input = Event()
 235             the_stdin_thread.out_of_raw_input.set()
 236             s1, s2 = socket.socketpair()
 237             the_stdin_thread.socket_read, the_stdin_thread.socket_write = s1, s2
 238             the_stdin_thread.interrupt_asked = False
 239             the_stdin_thread.setDaemon(True)
 240             the_stdin_thread.start()
 241             the_stdin_thread.socket_notification = socket_notification_reader()
 242             the_stdin_thread.prepend_text = None
 243             readline.set_pre_input_hook(the_stdin_thread.prepend_previous_text)
 244
 245     def prepend_previous_text(self):
 246         if self.prepend_text:
 247             readline.insert_text(self.prepend_text)
 248             readline.redisplay()
 249             self.prepend_text = None
 250
 251     def want_raw_input(self):
 252         nr, total = dispatchers.count_awaited_processes()
 253         if nr:
 254             prompt = 'waiting (%d/%d)> ' % (nr, total)
 255         else:
 256             prompt = 'ready (%d)> ' % total
 257         self.prompt = prompt
 258         set_last_status_length(len(prompt))
 259         self.raw_input_wanted.set()
 260         self.socket_notification.handle_read()
 261         self.in_raw_input.wait()
 262         self.raw_input_wanted.clear()
 263
 264     def no_raw_input(self):
 265         if not self.out_of_raw_input.isSet():
 266             interrupt_stdin_thread()
 267
 268     # Beware of races
 269     def run(self):
 270         while True:
 271             self.raw_input_wanted.wait()
 272             self.out_of_raw_input.set()
 273             self.in_raw_input.set()
 274             self.out_of_raw_input.clear()
 275             cmd = None
 276             try:
 277                 cmd = raw_input(self.prompt)
 278             except EOFError:
 279                 if self.interrupt_asked:
 280                     cmd = readline.get_line_buffer()
 281                 else:
 282                     cmd = chr(4) # Ctrl-D
 283             if self.interrupt_asked:
 284                 self.prepend_text = cmd
 285                 cmd = None
 286             self.in_raw_input.clear()
 287             self.out_of_raw_input.set()
 288             if cmd:
 289                 if echo_enabled:
 290                     completion.add_to_history(cmd)
 291                 else:
 292                     completion.remove_last_history_item()
 293             set_echo(True)
 294             if cmd is not None:
 295                 self.input_buffer.add(cmd + '\n')
 296                 write_main_socket('d')
 297
 298 the_stdin_thread = stdin_thread()