Merge "Add REPL"

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()

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()
 

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)