pdm/cli.py

   1 """Management for daemon processes
   2
   3 This module provides some client support for the daemon management
   4 provided in the pdm.srv module.
   5 """
   6
   7 import socket, pickle, struct, select, threading
   8
   9 __all__ = ["client", "replclient", "perfclient"]
  10
  11 class protoerr(Exception):
  12     """Raised on protocol errors"""
  13     pass
  14
  15 def resolve(spec):
  16     if isinstance(spec, socket.socket):
  17         return spec
  18     sk = None
  19     try:
  20         if "/" in spec:
  21             sk = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
  22             sk.connect(spec)
  23         elif spec.isdigit():
  24             sk = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  25             sk.connect(("localhost", int(spec)))
  26         elif ":" in spec:
  27             sk = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  28             p = spec.rindex(":")
  29             sk.connect((spec[:p], int(spec[p + 1:])))
  30         else:
  31             raise Exception("Unknown target specification %r" % spec)
  32         rv = sk
  33         sk = None
  34     finally:
  35         if sk is not None: sk.close()
  36     return rv
  37
  38 class client(object):
  39     """PDM client
  40
  41     This class provides general facilities to speak to PDM servers,
  42     and is mainly intended to be subclassed to provide for the
  43     specific protocols, such as replclient and perfclient do.
  44
  45     `client' instances can be passed as arguments to select.select(),
  46     and can be used in `with' statements.
  47     """
  48     def __init__(self, sk, proto = None):
  49         """Create a client object connected to the specified
  50         server. `sk' can either be a socket object, which is used as
  51         it is, or a string specification very similar to the
  52         specification for pdm.srv.listen, so see its documentation for
  53         details. The differences are only that this function does not
  54         take arguments specific to socket creation, like the mode and
  55         group arguments for Unix sockets. If `proto' is given, that
  56         subprotocol will negotiated with the server (by calling the
  57         select() method).
  58         """
  59         self.sk = resolve(sk)
  60         self.buf = ""
  61         line = self.readline()
  62         if line != "+PDM1":
  63             raise protoerr("Illegal protocol signature")
  64         if proto is not None:
  65             self.select(proto)
  66
  67     def close(self):
  68         """Close this connection"""
  69         self.sk.close()
  70
  71     def fileno(self):
  72         """Return the file descriptor of the underlying socket."""
  73         return self.sk.fileno()
  74
  75     def readline(self):
  76         """Read a single NL-terminated line and return it."""
  77         while True:
  78             p = self.buf.find("\n")
  79             if p >= 0:
  80                 ret = self.buf[:p]
  81                 self.buf = self.buf[p + 1:]
  82                 return ret
  83             ret = self.sk.recv(1024)
  84             if ret == "":
  85                 return None
  86             self.buf += ret
  87
  88     def select(self, proto):
  89         """Negotiate the given subprotocol with the server"""
  90         if "\n" in proto:
  91             raise Exception("Illegal protocol specified: %r" % proto)
  92         self.sk.send(proto + "\n")
  93         rep = self.readline()
  94         if len(rep) < 1 or rep[0] != "+":
  95             raise protoerr("Error reply when selecting protocol %s: %s" % (proto, rep[1:]))
  96
  97     def __enter__(self):
  98         return self
  99
 100     def __exit__(self, *excinfo):
 101         self.close()
 102         return False
 103
 104 class replclient(client):
 105     """REPL protocol client
 106
 107     Implements the client side of the REPL protocol; see pdm.srv.repl
 108     for details on the protocol and its functionality.
 109     """
 110     def __init__(self, sk):
 111         """Create a connected client as documented in the `client' class."""
 112         super(replclient, self).__init__(sk, "repl")
 113
 114     def run(self, code):
 115         """Run a single block of Python code on the server. Returns
 116         the output of the command (as documented in pdm.srv.repl) as a
 117         string.
 118         """
 119         while True:
 120             ncode = code.replace("\n\n", "\n")
 121             if ncode == code: break
 122             code = ncode
 123         while len(code) > 0 and code[-1] == "\n":
 124             code = code[:-1]
 125         self.sk.send(code + "\n\n")
 126         buf = ""
 127         while True:
 128             ln = self.readline()
 129             if ln[0] == " ":
 130                 buf += ln[1:] + "\n"
 131             elif ln[0] == "+":
 132                 return buf
 133             elif ln[0] == "-":
 134                 raise protoerr("Error reply: %s" % ln[1:])
 135             else:
 136                 raise protoerr("Illegal reply: %s" % ln)
 137
 138 class perfproxy(object):
 139     def __init__(self, cl, id, proto):
 140         self.cl = cl
 141         self.id = id
 142         self.proto = proto
 143         self.subscribers = set()
 144
 145     def lookup(self, name):
 146         self.cl.lock.acquire()
 147         try:
 148             id = self.cl.nextid
 149             self.cl.nextid += 1
 150         finally:
 151             self.cl.lock.release()
 152         (proto,) = self.cl.run("lookup", id, self.id, name)
 153         proxy = perfproxy(self.cl, id, proto)
 154         self.cl.proxies[id] = proxy
 155         return proxy
 156
 157     def listdir(self):
 158         return self.cl.run("ls", self.id)[0]
 159
 160     def readattr(self):
 161         return self.cl.run("readattr", self.id)[0]
 162
 163     def attrinfo(self):
 164         return self.cl.run("attrinfo", self.id)[0]
 165
 166     def invoke(self, method, *args, **kwargs):
 167         return self.cl.run("invoke", self.id, method, args, kwargs)[0]
 168
 169     def subscribe(self, cb):
 170         if cb in self.subscribers:
 171             raise ValueError("Already subscribed")
 172         if len(self.subscribers) == 0:
 173             self.cl.run("subs", self.id)
 174         self.subscribers.add(cb)
 175
 176     def unsubscribe(self, cb):
 177         if cb not in self.subscribers:
 178             raise ValueError("Not subscribed")
 179         self.subscribers.remove(cb)
 180         if len(self.subscribers) == 0:
 181             self.cl.run("unsubs", self.id)
 182
 183     def notify(self, ev):
 184         for cb in self.subscribers:
 185             try:
 186                 cb(ev)
 187             except: pass
 188
 189     def close(self):
 190         if self.id is not None:
 191             self.cl.run("unbind", self.id)
 192             del self.cl.proxies[self.id]
 193             self.id = None
 194
 195     def __del__(self):
 196         self.close()
 197
 198     def __enter__(self):
 199         return self
 200
 201     def __exit__(self, *excinfo):
 202         self.close()
 203         return False
 204
 205 class perfclient(client):
 206     """PERF protocol client
 207
 208     Implements the client side of the PERF protocol; see pdm.srv.perf
 209     for details on the protocol and its functionality.
 210
 211     This client class implements functions for finding PERF objects on
 212     the server, and returns, for each server-side object looked up, a
 213     proxy object that mimics exactly the PERF interfaces that the
 214     object implements. As the proxy objects reference live objects on
 215     the server, they should be released when they are no longer used;
 216     they implement a close() method for that purpose, and can also be
 217     used in `with' statements.
 218
 219     See pdm.srv.perf for details on the various PERF interfaces that
 220     the proxy objects might implement.
 221     """
 222     def __init__(self, sk):
 223         """Create a connected client as documented in the `client' class."""
 224         super(perfclient, self).__init__(sk, "perf")
 225         self.nextid = 0
 226         self.lock = threading.Lock()
 227         self.proxies = {}
 228         self.names = {}
 229
 230     def send(self, ob):
 231         buf = pickle.dumps(ob)
 232         buf = struct.pack(">l", len(buf)) + buf
 233         self.sk.send(buf)
 234
 235     def recvb(self, num):
 236         buf = ""
 237         while len(buf) < num:
 238             data = self.sk.recv(num - len(buf))
 239             if data == "":
 240                 raise EOFError()
 241             buf += data
 242         return buf
 243
 244     def recv(self):
 245         return pickle.loads(self.recvb(struct.unpack(">l", self.recvb(4))[0]))
 246
 247     def event(self, id, ev):
 248         proxy = self.proxies.get(id)
 249         if proxy is None: return
 250         proxy.notify(ev)
 251
 252     def dispatch(self, timeout = None):
 253         """Wait for an incoming notification from the server, and
 254         dispatch it to the callback functions that have been
 255         registered for it. If `timeout' is specified, wait no longer
 256         than so many seconds; otherwise, wait forever. This client
 257         object may also be used as argument to select.select().
 258         """
 259         rfd, wfd, efd = select.select([self.sk], [], [], timeout)
 260         if self.sk in rfd:
 261             msg = self.recv()
 262             if msg[0] == "*":
 263                 self.event(msg[1], msg[2])
 264             else:
 265                 raise ValueError("Unexpected non-event message: %r" % msg[0])
 266
 267     def recvreply(self):
 268         while True:
 269             reply = self.recv()
 270             if reply[0] in ("+", "-"):
 271                 return reply
 272             elif reply[0] == "*":
 273                 self.event(reply[1], reply[2])
 274             else:
 275                 raise ValueError("Illegal reply header: %r" % reply[0])
 276
 277     def run(self, cmd, *args):
 278         self.lock.acquire()
 279         try:
 280             self.send((cmd,) + args)
 281             reply = self.recvreply()
 282             if reply[0] == "+":
 283                 return reply[1:]
 284             else:
 285                 raise reply[1]
 286         finally:
 287             self.lock.release()
 288
 289     def lookup(self, module, obnm):
 290         """Look up a single server-side object by the given name in
 291         the given module. Will return a new proxy object for each
 292         call when called multiple times for the same name.
 293         """
 294         self.lock.acquire()
 295         try:
 296             id = self.nextid
 297             self.nextid += 1
 298         finally:
 299             self.lock.release()
 300         (proto,) = self.run("bind", id, module, obnm)
 301         proxy = perfproxy(self, id, proto)
 302         self.proxies[id] = proxy
 303         return proxy
 304
 305     def find(self, name):
 306         """Convenience function for looking up server-side objects
 307         through PERF directories and for multiple uses. The object
 308         name can be given as "MODULE.OBJECT", which will look up the
 309         named OBJECT in the named MODULE, and can be followed by any
 310         number of slash-separated names, which will assume that the
 311         object to the left of the slash is a PERF directory, and will
 312         return the object in that directory by the name to the right
 313         of the slash. For instance, find("pdm.perf.sysres/cputime")
 314         will return the built-in attribute for reading the CPU time
 315         used by the server process.
 316         """
 317         ret = self.names.get(name)
 318         if ret is None:
 319             if "/" in name:
 320                 p = name.rindex("/")
 321                 ret = self.find(name[:p]).lookup(name[p + 1:])
 322             else:
 323                 p = name.rindex(".")
 324                 ret = self.lookup(name[:p], name[p + 1:])
 325             self.names[name] = ret
 326         return ret