Zero-downtime restarts via interprocess descriptor transfer

takeover.py

#!/usr/bin/env python3
# -*- coding: utf8 -*-
"""
   takeover.py
   ~~~~~~~~~~~
   This script demonstrates a technique for zero-downtime restarts via
   interprocess descriptor transfer.
   
   The script operates a simple echo service on port 9999. When a new instance
   is launched, the old instance will transfer the server socket to the new
   instance over a UNIX socket connection. From that point on, new connections
   are routed to the new instance, but the old instance is still able to
   service any client connections that have already been established. When
   these connections finally drain, the old instance will shut down.

   Tested on Linux, but should work on other POSIX systems.

   Python 3 is required because Python 2 does not provide an API for the
   sendmsg() and recvmsg() system calls.

   The idea for this script came from HHVM:
     <https://github.com/facebook/hhvm/blob/a13d8a3981/hphp/runtime/server/takeover-agent.cpp>

   Copyright 2016 Ori Livneh <ori@wikimedia.org>
   
   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.

"""
import os
import select
import socket
import struct


BUF_SIZE = 4096
TAKEOVER_SOCK_PATH = './takeover.sock'
SERVER_PORT = 9999
s_fd = struct.Struct('i')


def send_fd(sock, fd, msg=b'\0'):
    if sock.family != socket.AF_UNIX:
        raise OSError('Descriptors can only be sent over UNIX domain sockets')
    control_msg = (socket.SOL_SOCKET, socket.SCM_RIGHTS, s_fd.pack(fd))
    return sock.sendmsg([msg], [control_msg])


def recv_fd(sock):
    msg, ancdata, flags, addr = sock.recvmsg(1, socket.CMSG_LEN(s_fd.size))
    cmsg_level, cmsg_type, cmsg_data = ancdata[0]
    assert (cmsg_level == socket.SOL_SOCKET and
            cmsg_type == socket.SCM_RIGHTS and
            len(cmsg_data) == s_fd.size)
    return s_fd.unpack(cmsg_data)[0]


server = None
takeover_client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
try:
    takeover_client.connect(TAKEOVER_SOCK_PATH)
except FileNotFoundError:
    pass  # There is no running instance to take over.
except ConnectionRefusedError:
    os.unlink(TAKEOVER_SOCK_PATH)
else:
    print('* acquiring server socket from running instance...')
    server = socket.socket(fileno=recv_fd(takeover_client))

if server is None:
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    print('* echo server running on port %s' % SERVER_PORT)
    server.bind(('', SERVER_PORT))
    server.listen(5)


takeover_server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
takeover_server.bind(TAKEOVER_SOCK_PATH)
takeover_server.listen(1)

sockets = [server, takeover_server]

while sockets:
    readable, *_ = select.select(sockets, [], [])
    for sock in readable:

        if sock is takeover_server:
            # A new instance is requesting the server socket
            conn, addr = takeover_server.accept()
            print('* received takeover request from %s' % (addr,))

            # We want the new instance to be able to start its own takeover
            # server right away, so we shut ours down before handing over the
            # server socket.
            takeover_server.close()
            sockets.remove(takeover_server)
            os.unlink(TAKEOVER_SOCK_PATH)

            send_fd(conn, server.fileno())
            conn.close()

            # We can close the server socket now. The kernel knows that it is
            # in flight and will it keep it open for the receiving process.
            server.close()
            sockets.remove(server)

        elif sock is server:
            # New client connection
            conn, addr = server.accept()
            print('* new connection from %s' % (addr,))
            sockets.append(conn)

        else:
            peer = sock.getpeername()
            data = sock.recv(BUF_SIZE)

            if data:
                print('* data from %s: "%s"' % (peer, data))
                sock.send(data)
            else:
                print('* closing connection: %s' % peer)
                sock.close()
                sockets.remove(sock)