Tuesday, December 10, 2013

Python socketserver and upstart-socket-bridge

In a previous post, I described how to use Upstart to monitor a port, trigger a Python script upon connection, and then let the Python application send and receive data over that port.

In a server application, Upstart can take over the bind() and listen() elements of the server daemon, sometimes eliminating the need for a daemon at all.

Here is an example of a Python socketserver, and how it changes when using Upstart.

Some complexity is reduced: The handler is simpler, it no longer needs to fork() or serve_forever(). Each connection triggers the application, so the application need only handle a single read()/write() cycle, then terminate.

Some complexity is increased: We need to override two methods of a socketserver class to make it upstart-socket-bridge compatible.




Here is the original socketserver.TCPServer serve_forever example, from the Python documentation:

import socketserver

class MyTCPHandler(socketserver.BaseRequestHandler):
    """
    The RequestHandler class for our server.

    It is instantiated once per connection to the server, and must
    override the handle() method to implement communication to the
    client.
    """

    def handle(self):
        # self.request is the TCP socket connected to the client
        self.data = self.request.recv(1024).strip()
        print("{} wrote:".format(self.client_address[0]))
        print(self.data)
        # just send back the same data, but upper-cased
        self.request.sendall(self.data.upper())

if __name__ == "__main__":
    HOST, PORT = "localhost", 9999

    # Create the server, binding to localhost on port 9999
    server = socketserver.TCPServer((HOST, PORT), MyTCPHandler)

    # Activate the server; this will keep running until you
    # interrupt the program with Ctrl-C
    server.serve_forever()


It's easy to test. Let's save the file as /tmp/socket-server-original.py

$ python3 /tmp/socket-server-original.py     # On Terminal #1

$ echo "Test String" | nc localhost 9999     # On Terminal #2
TEST STRING

The example works.




Now, let's adapt it for upstart-socket-bridge.

It takes a bit of reading through the socketserver source code to figure it out:

TCPServer's __init__ calls good old bind() and listen() using it's own methods server_bind() and server_activate(), respectively.

Both of those methods can be overridden. Better yet, they use socket module calls (socket.bind() and socket.listen()) that we already know how to replace with socket.fromfd().

So we can get TCPServer to work by overriding the two methods like this:

import os, socket

class MyServer(socketserver.TCPServer):
    """
    This class overrides two method in socketserver.TCPServer 
    and makes it Upstart-compatible
    """
    def server_bind(self):
        """ Replace the socket-created file descriptor (fd) with the Upstart-provided fd"""
        fd = int(os.environ["UPSTART_FDS"]
        self.socket = socket.fromfd(fd, socket.AF_INET, socket.SOCK_STREAM)
def server_activate(self): """This means listen(), but we don't want it do do anything""" pass


Changing the TCPServer class slightly changes the way way we activate.
We must call the new class instead of TCPServer, and we no longer need the address since bind() no longer does anything:

    server = socketserver.TCPServer((HOST, PORT), MyTCPHandler)  # OLD
    server = MyServer(None, MyTCPHandler)                        # NEW


The code that handles the request, generates the response, and sends the response has two changes. Printing on the terminal won't work, since Upstart doesn't use a terminal to start the process. So we'll delete the printing lines:

        print("{} wrote:".format(self.client_address[0]))
        print(self.data)





Finally, one more change. The script can quit after the response has been sent:

    server.serve_forever()      # OLD
    server.handle_request()     # NEW





When we put it all together, here is the new version that works with upstart-socket-bridge:

import socketserver, os, socket

class MyServer(socketserver.TCPServer):
    """
    This class overrides two method in socketserver.TCPServer 
    and makes it Upstart-compatible
    """
    def server_bind(self):
        """ Replace the socket FD with the Upstart-provided FD"""
        fd = int(os.environ["UPSTART_FDS"])
        self.socket = socket.fromfd(fd, socket.AF_INET, socket.SOCK_STREAM)

    def server_activate(self):
        pass            # Upstart takes care of listen()


class MyTCPHandler(socketserver.BaseRequestHandler):
    """
    The RequestHandler class for our server.

    It is instantiated once per connection to the server, and must
    override the handle() method to implement communication to the
    client.
    """

    def handle(self):
        # self.request is the TCP socket connected to the client
        self.data = self.request.recv(1024).strip()

        # just send back the same data, but upper-cased
        self.request.sendall(self.data.upper())

if __name__ == "__main__":
    # Create the server using the file descriptor from Upstart
    server = MyServer(None, MyTCPHandler)

    # Run once, then terminate
    server.handle_request()


Let's save the file as /tmp/socket-server-new.py
We'll need to adapt (or create) /etc/init/socket-test.conf Upstart job file.

description "upstart-socket-bridge test"
start on socket PROTO=inet PORT=34567 ADDR=127.0.0.1
setuid myusername   # Use *your* user name. Doesn't need to be root
exec /usr/bin/python3 /tmp/socket-server-new.py

And let's give it a try on the terminals:

$ python3 /tmp/socket-server-new.py           # On Terminal #1

$ echo "Test String" | nc localhost 34567     # On Terminal #2
TEST STRING

Overriding socketserver for upstart compatibility works.




After this test, cleanup is easy

$ sudo rm /etc/init/socket-test.conf
$ rm /tmp/socket-server-original.py /tmp/socket-server-new.py

No comments: