test/py/TestServer.py

#!/usr/bin/env python

#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you 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 logging
import os
import signal
import ssl
import sys
import time
from optparse import OptionParser

from util import local_libpath
sys.path.insert(0, local_libpath())
from thrift.protocol import TProtocol, TProtocolDecorator
from thrift.Thrift import TException

SCRIPT_DIR = os.path.abspath(os.path.dirname(__file__))


class TestHandler(object):
    def __init__(self, options):
        self.options = options

    def testVoid(self):
        if self.options.verbose > 1:
            logging.info('testVoid()')

    def testString(self, str):
        if self.options.verbose > 1:
            logging.info('testString(%s)' % str)
        return str

    def testBool(self, boolean):
        if self.options.verbose > 1:
            logging.info('testBool(%s)' % str(boolean).lower())
        return boolean

    def testByte(self, byte):
        if self.options.verbose > 1:
            logging.info('testByte(%d)' % byte)
        return byte

    def testI16(self, i16):
        if self.options.verbose > 1:
            logging.info('testI16(%d)' % i16)
        return i16

    def testI32(self, i32):
        if self.options.verbose > 1:
            logging.info('testI32(%d)' % i32)
        return i32

    def testI64(self, i64):
        if self.options.verbose > 1:
            logging.info('testI64(%d)' % i64)
        return i64

    def testDouble(self, dub):
        if self.options.verbose > 1:
            logging.info('testDouble(%f)' % dub)
        return dub

    def testBinary(self, thing):
        if self.options.verbose > 1:
            logging.info('testBinary()')  # TODO: hex output
        return thing

    def testUuid(self, thing):
        if self.options.verbose > 1:
            logging.info('testUuid(%s)' % thing)
        return thing

    def testStruct(self, thing):
        if self.options.verbose > 1:
            logging.info('testStruct({%s, %s, %s, %s})' % (thing.string_thing, thing.byte_thing, thing.i32_thing, thing.i64_thing))
        return thing

    def testException(self, arg):
        from ThriftTest.ttypes import Xception
        # if self.options.verbose > 1:
        logging.info('testException(%s)' % arg)
        if arg == 'Xception':
            raise Xception(errorCode=1001, message=arg)
        elif arg == 'TException':
            raise TException(message='This is a TException')

    def testMultiException(self, arg0, arg1):
        from ThriftTest.ttypes import Xtruct, Xception, Xception2

        if self.options.verbose > 1:
            logging.info('testMultiException(%s, %s)' % (arg0, arg1))
        if arg0 == 'Xception':
            raise Xception(errorCode=1001, message='This is an Xception')
        elif arg0 == 'Xception2':
            raise Xception2(
                errorCode=2002,
                struct_thing=Xtruct(string_thing='This is an Xception2'))
        return Xtruct(string_thing=arg1)

    def testOneway(self, seconds):
        if self.options.verbose > 1:
            logging.info('testOneway(%d) => sleeping...' % seconds)
        time.sleep(seconds / 3)  # be quick
        if self.options.verbose > 1:
            logging.info('done sleeping')

    def testNest(self, thing):
        if self.options.verbose > 1:
            logging.info('testNest(%s)' % thing)
        return thing

    def testMap(self, thing):
        if self.options.verbose > 1:
            logging.info('testMap(%s)' % thing)
        return thing

    def testStringMap(self, thing):
        if self.options.verbose > 1:
            logging.info('testStringMap(%s)' % thing)
        return thing

    def testSet(self, thing):
        if self.options.verbose > 1:
            logging.info('testSet(%s)' % thing)
        return thing

    def testList(self, thing):
        if self.options.verbose > 1:
            logging.info('testList(%s)' % thing)
        return thing

    def testEnum(self, thing):
        if self.options.verbose > 1:
            logging.info('testEnum(%s)' % thing)
        return thing

    def testTypedef(self, thing):
        if self.options.verbose > 1:
            logging.info('testTypedef(%s)' % thing)
        return thing

    def testMapMap(self, thing):
        if self.options.verbose > 1:
            logging.info('testMapMap(%s)' % thing)
        return {
            -4: {
                -4: -4,
                -3: -3,
                -2: -2,
                -1: -1,
            },
            4: {
                4: 4,
                3: 3,
                2: 2,
                1: 1,
            },
        }

    def testInsanity(self, argument):
        from ThriftTest.ttypes import Insanity

        if self.options.verbose > 1:
            logging.info('testInsanity(%s)' % argument)
        return {
            1: {
                2: argument,
                3: argument,
            },
            2: {6: Insanity()},
        }

    def testMulti(self, arg0, arg1, arg2, arg3, arg4, arg5):
        from ThriftTest.ttypes import Xtruct

        if self.options.verbose > 1:
            logging.info('testMulti(%s, %s, %s, %s, %s, %s)' % (arg0, arg1, arg2, arg3, arg4, arg5))
        return Xtruct(string_thing='Hello2',
                      byte_thing=arg0, i32_thing=arg1, i64_thing=arg2)


class SecondHandler(object):
    def secondtestString(self, argument):
        return "testString(\"" + argument + "\")"


# LAST_SEQID is a global because we have one transport and multiple protocols
# running on it (when multiplexed)
LAST_SEQID = None


class TPedanticSequenceIdProtocolWrapper(TProtocolDecorator.TProtocolDecorator):
    """
    Wraps any protocol with sequence ID checking: looks for outbound
    uniqueness as well as request/response alignment.
    """

    def __init__(self, protocol):
        # TProtocolDecorator.__new__ does all the heavy lifting
        pass

    def readMessageBegin(self):
        global LAST_SEQID
        (name, type, seqid) =\
            super(TPedanticSequenceIdProtocolWrapper, self).readMessageBegin()
        if LAST_SEQID is not None and LAST_SEQID == seqid:
            raise TProtocol.TProtocolException(
                TProtocol.TProtocolException.INVALID_DATA,
                "We received the same seqid {0} twice in a row".format(seqid))
        LAST_SEQID = seqid
        return (name, type, seqid)


def make_pedantic(proto):
    """ Wrap a protocol in the pedantic sequence ID wrapper. """
    # NOTE: this is disabled for now as many clients send seqid
    #       of zero and that is okay, need a way to identify
    #       clients that MUST send seqid unique to function right
    #       or just force all implementations to send unique seqids (preferred)
    return proto  # TPedanticSequenceIdProtocolWrapper(proto)


class TPedanticSequenceIdProtocolFactory(TProtocol.TProtocolFactory):
    def __init__(self, encapsulated):
        super(TPedanticSequenceIdProtocolFactory, self).__init__()
        self.encapsulated = encapsulated

    def getProtocol(self, trans):
        return make_pedantic(self.encapsulated.getProtocol(trans))


def main(options):
    # common header allowed client types
    allowed_client_types = [
        THeaderTransport.THeaderClientType.HEADERS,
        THeaderTransport.THeaderClientType.FRAMED_BINARY,
        THeaderTransport.THeaderClientType.UNFRAMED_BINARY,
        THeaderTransport.THeaderClientType.FRAMED_COMPACT,
        THeaderTransport.THeaderClientType.UNFRAMED_COMPACT,
    ]

    # set up the protocol factory form the --protocol option
    prot_factories = {
        'accel': TBinaryProtocol.TBinaryProtocolAcceleratedFactory(),
        'multia': TBinaryProtocol.TBinaryProtocolAcceleratedFactory(),
        'accelc': TCompactProtocol.TCompactProtocolAcceleratedFactory(),
        'multiac': TCompactProtocol.TCompactProtocolAcceleratedFactory(),
        'binary': TPedanticSequenceIdProtocolFactory(TBinaryProtocol.TBinaryProtocolFactory()),
        'multi': TPedanticSequenceIdProtocolFactory(TBinaryProtocol.TBinaryProtocolFactory()),
        'compact': TCompactProtocol.TCompactProtocolFactory(),
        'multic': TCompactProtocol.TCompactProtocolFactory(),
        'header': THeaderProtocol.THeaderProtocolFactory(allowed_client_types),
        'multih': THeaderProtocol.THeaderProtocolFactory(allowed_client_types),
        'json': TJSONProtocol.TJSONProtocolFactory(),
        'multij': TJSONProtocol.TJSONProtocolFactory(),
    }
    pfactory = prot_factories.get(options.proto, None)
    if pfactory is None:
        raise AssertionError('Unknown --protocol option: %s' % options.proto)
    try:
        pfactory.string_length_limit = options.string_limit
        pfactory.container_length_limit = options.container_limit
    except Exception:
        # Ignore errors for those protocols that does not support length limit
        pass

    # get the server type (TSimpleServer, TNonblockingServer, etc...)
    if len(args) > 1:
        raise AssertionError('Only one server type may be specified, not multiple types.')
    server_type = args[0]
    if options.trans == 'http':
        server_type = 'THttpServer'

    # Set up the handler and processor objects
    handler = TestHandler(options)
    processor = ThriftTest.Processor(handler)

    if options.proto.startswith('multi'):
        secondHandler = SecondHandler()
        secondProcessor = SecondService.Processor(secondHandler)

        multiplexedProcessor = TMultiplexedProcessor()
        multiplexedProcessor.registerDefault(processor)
        multiplexedProcessor.registerProcessor('ThriftTest', processor)
        multiplexedProcessor.registerProcessor('SecondService', secondProcessor)
        processor = multiplexedProcessor

    global server

    # Handle THttpServer as a special case
    if server_type == 'THttpServer':
        if options.ssl:
            __certfile = os.path.join(os.path.dirname(SCRIPT_DIR), "keys", "server.crt")
            __keyfile = os.path.join(os.path.dirname(SCRIPT_DIR), "keys", "server.key")
            server = THttpServer.THttpServer(processor, ('', options.port), pfactory, cert_file=__certfile, key_file=__keyfile)
        else:
            server = THttpServer.THttpServer(processor, ('', options.port), pfactory)
        server.serve()
        sys.exit(0)

    # set up server transport and transport factory

    host = None
    if options.ssl:
        from thrift.transport import TSSLSocket
        keys_dir = os.path.join(os.path.dirname(SCRIPT_DIR), 'keys')
        ca_certs = os.path.join(keys_dir, 'client.pem')
        certfile = os.path.join(keys_dir, 'server.crt')
        keyfile = os.path.join(keys_dir, 'server.key')
        ssl_version = getattr(ssl, 'PROTOCOL_TLS_SERVER', ssl.PROTOCOL_TLSv1)
        transport = TSSLSocket.TSSLServerSocket(
            host,
            options.port,
            certfile=certfile,
            keyfile=keyfile,
            ca_certs=ca_certs,
            cert_reqs=ssl.CERT_REQUIRED,
            ssl_version=ssl_version,
        )
    else:
        transport = TSocket.TServerSocket(host, options.port, options.domain_socket)
    tfactory = TTransport.TBufferedTransportFactory()
    if options.trans == 'buffered':
        tfactory = TTransport.TBufferedTransportFactory()
    elif options.trans == 'framed':
        tfactory = TTransport.TFramedTransportFactory()
    elif options.trans == '':
        raise AssertionError('Unknown --transport option: %s' % options.trans)
    else:
        tfactory = TTransport.TBufferedTransportFactory()
    # if --zlib, then wrap server transport, and use a different transport factory
    if options.zlib:
        if server_type != "TProcessPoolServer":
            transport = TZlibTransport.TZlibTransport(transport)  # wrap with zlib
        # Avoid wrapping the server transport for process pools; TZlibTransport isn't picklable on spawn.
        tfactory = TZlibTransport.TZlibTransportFactory()

    # do server-specific setup here:
    if server_type == "TNonblockingServer":
        server = TNonblockingServer.TNonblockingServer(processor, transport, inputProtocolFactory=pfactory)
    elif server_type == "TProcessPoolServer":
        import signal
        from thrift.server import TProcessPoolServer
        server = TProcessPoolServer.TProcessPoolServer(processor, transport, tfactory, pfactory)
        server.setNumWorkers(5)

        def set_alarm():
            def clean_shutdown(signum, frame):
                for worker in server.workers:
                    if options.verbose > 0:
                        logging.info('Terminating worker: %s' % worker)
                    worker.terminate()
                if options.verbose > 0:
                    logging.info('Requesting server to stop()')
                try:
                    server.stop()
                except Exception:
                    pass
            signal.signal(signal.SIGALRM, clean_shutdown)
            signal.alarm(4)
        set_alarm()
    else:
        # look up server class dynamically to instantiate server
        ServerClass = getattr(TServer, server_type)
        server = ServerClass(processor, transport, tfactory, pfactory)
    # enter server main loop
    server.serve()


def exit_gracefully(signum, frame):
    print("SIGINT received\n")
    server.shutdown()   # doesn't work properly, yet
    sys.exit(0)


if __name__ == '__main__':
    signal.signal(signal.SIGINT, exit_gracefully)

    parser = OptionParser()
    parser.add_option('--libpydir', type='string', dest='libpydir',
                      help='include this directory to sys.path for locating library code')
    parser.add_option('--genpydir', type='string', dest='genpydir',
                      default='gen-py',
                      help='include this directory to sys.path for locating generated code')
    parser.add_option("--port", type="int", dest="port",
                      help="port number for server to listen on")
    parser.add_option("--zlib", action="store_true", dest="zlib",
                      help="use zlib wrapper for compressed transport")
    parser.add_option("--ssl", action="store_true", dest="ssl",
                      help="use SSL for encrypted transport")
    parser.add_option('-v', '--verbose', action="store_const",
                      dest="verbose", const=2,
                      help="verbose output")
    parser.add_option('-q', '--quiet', action="store_const",
                      dest="verbose", const=0,
                      help="minimal output")
    parser.add_option('--protocol', dest="proto", type="string",
                      help="protocol to use, one of: accel, accelc, binary, compact, json, multi, multia, multiac, multic, multih, multij")
    parser.add_option('--transport', dest="trans", type="string",
                      help="transport to use, one of: buffered, framed, http")
    parser.add_option('--domain-socket', dest="domain_socket", type="string",
                      help="Unix domain socket path")
    parser.add_option('--container-limit', dest='container_limit', type='int', default=None)
    parser.add_option('--string-limit', dest='string_limit', type='int', default=None)
    parser.set_defaults(port=9090, verbose=1, proto='binary', transport='buffered')
    options, args = parser.parse_args()

    # Print TServer log to stdout so that the test-runner can redirect it to log files
    logging.basicConfig(level=options.verbose)

    sys.path.insert(0, os.path.join(SCRIPT_DIR, options.genpydir))

    from ThriftTest import ThriftTest, SecondService
    from thrift.TMultiplexedProcessor import TMultiplexedProcessor
    from thrift.transport import THeaderTransport
    from thrift.transport import TTransport
    from thrift.transport import TSocket
    from thrift.transport import TZlibTransport
    from thrift.protocol import TBinaryProtocol
    from thrift.protocol import TCompactProtocol
    from thrift.protocol import THeaderProtocol
    from thrift.protocol import TJSONProtocol
    from thrift.server import TServer, TNonblockingServer, THttpServer

    sys.exit(main(options))