diff --git a/nodepool/cmd/builder.py b/nodepool/cmd/builder.py index 46f304a27..1b396a14f 100644 --- a/nodepool/cmd/builder.py +++ b/nodepool/cmd/builder.py @@ -15,6 +15,7 @@ import sys from nodepool import builder import nodepool.cmd +import nodepool.lib.repl class NodePoolBuilderApp(nodepool.cmd.NodepoolDaemonApp): @@ -40,6 +41,8 @@ class NodePoolBuilderApp(nodepool.cmd.NodepoolDaemonApp): parser.add_argument('--upload-workers', dest='upload_workers', default=4, help='number of upload workers', type=int) + parser.add_argument('--repl', action='store_true', + help="Start a REPL on port 3000") return parser def parse_args(self): @@ -57,6 +60,9 @@ class NodePoolBuilderApp(nodepool.cmd.NodepoolDaemonApp): signal.signal(signal.SIGINT, self.sigint_handler) self.nb.start() + if self.args.repl: + self.repl = nodepool.lib.repl.REPLServer(self.nb) + self.repl.start() while True: signal.pause() diff --git a/nodepool/cmd/launcher.py b/nodepool/cmd/launcher.py index 23c2677f8..b796e6647 100644 --- a/nodepool/cmd/launcher.py +++ b/nodepool/cmd/launcher.py @@ -21,6 +21,7 @@ import signal import nodepool.cmd import nodepool.launcher import nodepool.webapp +import nodepool.lib.repl log = logging.getLogger(__name__) @@ -38,6 +39,8 @@ class NodePoolLauncherApp(nodepool.cmd.NodepoolDaemonApp): parser.add_argument('-s', dest='secure', help='path to secure file') parser.add_argument('--no-webapp', action='store_true') + parser.add_argument('--repl', action='store_true', + help="Start a REPL on port 3000") return parser def parse_args(self): @@ -73,6 +76,10 @@ class NodePoolLauncherApp(nodepool.cmd.NodepoolDaemonApp): if not self.args.no_webapp: self.webapp.start() + if self.args.repl: + self.repl = nodepool.lib.repl.REPLServer(self.pool) + self.repl.start() + while True: signal.pause() diff --git a/nodepool/lib/__init__.py b/nodepool/lib/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/nodepool/lib/repl.py b/nodepool/lib/repl.py new file mode 100644 index 000000000..de4a9b4ac --- /dev/null +++ b/nodepool/lib/repl.py @@ -0,0 +1,81 @@ +# Licensed under the Apache License, Version 2.0 (the "License"); you may +# not use this file except in compliance with the License. You may obtain +# a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT +# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the +# License for the specific language governing permissions and limitations +# under the License. + +# Based on ASL2 code from: +# https://gist.github.com/tim-patterson/4471877 + +import sys +import io +import threading +import socketserver +import code + + +class ThreadLocalProxy(object): + def __init__(self, default): + self.files = {} + self.default = default + + def __getattr__(self, name): + obj = self.files.get(threading.current_thread(), self.default) + return getattr(obj, name) + + def register(self, obj): + self.files[threading.current_thread()] = obj + + def unregister(self): + self.files.pop(threading.current_thread()) + + +class REPLHandler(socketserver.StreamRequestHandler): + def handle(self): + sys.stdout.register(io.TextIOWrapper(self.wfile, 'utf8')) + sys.stderr.register(io.TextIOWrapper(self.wfile, 'utf8')) + sys.stdin.register(io.TextIOWrapper(self.rfile, 'utf8')) + try: + console = code.InteractiveConsole(locals()) + console.interact('Console:') + except Exception: + pass + finally: + sys.stdout.unregister() + sys.stderr.unregister() + sys.stdin.unregister() + + +class REPLThreadedTCPServer(socketserver.ThreadingMixIn, + socketserver.TCPServer): + daemon_threads = True + allow_reuse_address = True + + def __init__(self, app, *args, **kw): + self.app = app + super(REPLThreadedTCPServer, self).__init__(*args, **kw) + sys.stdout = ThreadLocalProxy(sys.stdout) + sys.stderr = ThreadLocalProxy(sys.stderr) + sys.stdin = ThreadLocalProxy(sys.stdin) + + +class REPLServer(object): + def __init__(self, app): + self.server = REPLThreadedTCPServer( + app, ('localhost', 3000), REPLHandler) + + def start(self): + self.thread = threading.Thread(target=self.server.serve_forever) + self.thread.daemon = True + self.thread.start() + + def stop(self): + self.server.shutdown() + self.server.server_close() + self.thread.join(10)