├── MANIFEST.in ├── README.md ├── README.zh.md ├── examples ├── protocols │ ├── durian.py │ └── simple.py └── scripts │ ├── example.yml │ └── heartbeat.yml ├── setup.cfg ├── setup.py └── swarm ├── __init__.py ├── command.py ├── errors.py ├── fakeclient.py ├── protocol.py ├── script.py └── server.py /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include MANIFEST.in 2 | include *.md *.markdown 3 | recursive-include examples *.py *.yml 4 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | `swarm` 2 | ======= 3 | 4 | 5 | `swarm` is a simple benchmarking framework built upon `gevent`. It can be 6 | used to generate massive simultaneous and persistent TCP connections to a 7 | server, while each connection interacts with the server using your custom 8 | protocol. 9 | 10 | `swarm` does not create any connections on startup, you have to telnet to it 11 | to control its behavior. 12 | 13 | 14 | 15 | How to install 16 | -------------- 17 | 18 | python setup.py install 19 | 20 | 21 | 22 | How to use 23 | ---------- 24 | 25 | 1. Implement the protocol module for your server. This module should be 26 | consist of functions you want to use in the script descripted in step 2. 27 | `swarm` provides several examples in `examples/protocol`: 28 | 29 | # examples/protocols/durian.py 30 | from swarm.protocol import reply_parser_crlf 31 | 32 | 33 | def heartbeat(client): 34 | def make_request(): 35 | return "heartbeat\r\n" 36 | 37 | client.send_for_reply(make_request(), reply_parser_crlf()) 38 | 39 | 40 | def enter_chatroom(client): 41 | def make_request(): 42 | return "enter_chatroom\r\n" 43 | 44 | client.send_for_reply(make_request(), reply_parser_crlf()) 45 | 46 | 47 | def leave_chatroom(client): 48 | def make_request(): 49 | return "leave_chatroom\r\n" 50 | 51 | client.send_for_reply(make_request(), reply_parser_crlf()) 52 | 53 | 54 | def close_connection(client): 55 | client.close_connection() 56 | 57 | 58 | 2. Write a script (in yaml) that defines actions for each connection to 59 | perform in order. This script will be executed by every connection repeatedly. 60 | Also, under the `examples` directory, an example script using above custom 61 | protocol is provided. 62 | 63 | # examples/scripts/example.yml 64 | 65 | protocol: durian 66 | actions: 67 | - command: heartbeat 68 | args: 69 | uid: 123 70 | rounds: [2, 10] 71 | sleep_between_rounds: [1, 5] 72 | sleep_after_action: [1, 5] 73 | 74 | - command: create_chatroom 75 | args: 76 | uid: 123 77 | rid: 7 78 | rounds: 1 79 | sleep_between_rounds: 0 80 | sleep_after_action: [1, 5] 81 | 82 | 83 | 3. Run `swarm` with your script (Check out `swarm --help` first). 84 | 85 | 86 | swarm --listen 127.0.0.1:8400 --server 127.0.0.1:8080 \ 87 | --script /path/to/examples/scripts/example.yml \ 88 | --protocol-dir /path/to/examples/protocols/ 89 | 90 | 4. Use your favorite TCP client (`nc`, `telnet` etc.) to connect to `swarm` 91 | listening address, and then you will see what to do next. 92 | 93 | % telnet 127.0.0.1 8400 94 | laptop:notes:% telnet 127.0.0.1 8412 95 | Trying 127.0.0.1... 96 | Connected to 127.0.0.1. 97 | Escape character is '^]'. 98 | 99 | =============================== 100 | Welcome to Swarm Remote Console 101 | =============================== 102 | 103 | Usage: 104 | incr 1000\n -- Start 1000 new connections 105 | decr 100\n -- Close 100 random connections 106 | stop\n -- Stop swarm benchmarker remotely 107 | quit\n -- Close this control session 108 | help\n -- Help 109 | 110 | 5. Checkout the output of `swarm`. The running state will be printed 111 | periodically. 112 | 113 | 114 | 115 | API 116 | --- 117 | 118 | In order to provide the ability to communicate with your server though TCP 119 | connection for your protocol's implementation, `swarm` executes each action 120 | (the function defined in protocol module) with a instance of `FakeClient` as 121 | the first argument, which has the following network-relative methods: 122 | 123 | * `send_for_reply(request, reply_parser)` - Send the request and wait for the 124 | reply. 125 | 126 | * `request` should be a string (binary or not) containing the whole 127 | request. 128 | 129 | * `reply_parser` should be a function object which is used to parse the 130 | reply. It tells `swarm` how much bytes of data left to read for this reply 131 | though its return value. 132 | 133 | * `send_noreply(request)` - Send the request and no reply needed. 134 | 135 | * `request` should be a string (binary or not) containing the whole 136 | request. 137 | 138 | * `close_connection()` - Just close the underlying TCP connection. `swarm` will 139 | create a new connection and start a new iteration for the script. 140 | 141 | 142 | 143 | ###Notice 144 | 145 | * Don't do any blocking operations `gevent` does not support in your protocol. 146 | You've been warned. 147 | 148 | 149 | 150 | Directives 151 | ---------- 152 | 153 | Once your protocol module is ready, you can tell every connection how to act 154 | by writing a script. 155 | 156 | You have to learn some yaml first, and then you can use the directives 157 | provided by `swarm` to make up your benchmark plan: 158 | 159 | * `protocol` - protocol module name. The location of its parent directory is 160 | specified with command line option `--protocol-dir` of `swarm`. 161 | 162 | * `actions` - array of actions. Each action object is consist of: 163 | 164 | * `command` - One of the functions from the protocol module. 165 | 166 | * `args` - Extra arguments (beside the instance of `FakeClient`) this 167 | command can accept. 168 | 169 | * `rounds` - How many rounds to execute for this action per iteration. This 170 | directive accepts both integer and two-elements array (`swarm` takes it as 171 | a range and select a integer in it randomly). 172 | 173 | * `sleep_between_rounds` - How many seconds to sleep between every round. 174 | This directive accepts both integer and two-elements array (`swarm` takes 175 | it as a range and select a integer in it randomly). 176 | 177 | * `sleep_after_action` - How many seconds to sleep when current action 178 | completed. This directive accepts both integer and two-elements array 179 | (`swarm` taks it as a range and select a integer in it randomly). 180 | 181 | 182 | 183 | Tuning your OS 184 | -------------- 185 | 186 | 187 | To reduce memory usage, `swarm` configures every connection with minimum value 188 | for send/receve buffer by default. 189 | 190 | To increase concurrency of `swarm`, you still have to 191 | 192 | * raise the open files limit, 193 | 194 | * increase the system port range, 195 | 196 | or `swarm` will fail you at some point. :) 197 | 198 | 199 | 200 | TODO 201 | ---- 202 | 203 | Please do tell me what you need, my goddess. 204 | 205 | 206 | -------------------------------------------------------------------------------- /README.zh.md: -------------------------------------------------------------------------------- 1 | `swarm` 2 | ======= 3 | 4 | 5 | `swarm` 是一个简单的使用 `gevent` 开发的支持自定义协议的长连接压测框架。 6 | 7 | * 通过开发针对自定义协议进行解析的 Python 模块，`swarm` 可以用来对使用自定义 8 | 协议的长连接服务器进行压力测试； 9 | 10 | * 使用 yaml 格式的脚本，使用者可以定义每个连接的执行过程； 11 | 12 | 13 | [完整文档](https://github.com/duhoobo/swarm/blob/master/README.md) 14 | -------------------------------------------------------------------------------- /examples/protocols/durian.py: -------------------------------------------------------------------------------- 1 | from swarm.protocol import reply_parser_crlf 2 | 3 | 4 | def heartbeat(client): 5 | def make_request(): 6 | return "heartbeat\r\n" 7 | 8 | client.send_for_reply(make_request(), reply_parser_crlf()) 9 | 10 | 11 | def enter_chatroom(client): 12 | def make_request(): 13 | return "enter_chatroom\r\n" 14 | 15 | client.send_for_reply(make_request(), reply_parser_crlf()) 16 | 17 | 18 | def leave_chatroom(client): 19 | def make_request(): 20 | return "leave_chatroom\r\n" 21 | 22 | client.send_for_reply(make_request(), reply_parser_crlf()) 23 | 24 | 25 | def close_connection(client): 26 | client.close_connection() 27 | -------------------------------------------------------------------------------- /examples/protocols/simple.py: -------------------------------------------------------------------------------- 1 | 2 | def heartbeat(client, uid): 3 | def make_request(): 4 | return "hello %s\r\n" % str(uid) 5 | 6 | client.send_noreply(make_request()) 7 | 8 | 9 | def close_connection(client): 10 | client.close_connection() 11 | -------------------------------------------------------------------------------- /examples/scripts/example.yml: -------------------------------------------------------------------------------- 1 | protocol: durian 2 | actions: 3 | - command: heartbeat 4 | args: 5 | uid: 123 6 | rounds: [2, 10] 7 | sleep_between_rounds: [1, 5] 8 | sleep_after_action: [1, 5] 9 | 10 | - command: create_chatroom 11 | args: 12 | uid: 123 13 | rid: 7 14 | rounds: 1 15 | sleep_between_rounds: 0 16 | sleep_after_action: [1, 5] 17 | -------------------------------------------------------------------------------- /examples/scripts/heartbeat.yml: -------------------------------------------------------------------------------- 1 | protocol: simple 2 | actions: 3 | - command: heartbeat 4 | args: 5 | uid: 123 6 | rounds: [2, 5] 7 | sleep_between_rounds: [1, 3] 8 | sleep_after_action: [1, 2] 9 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/duhoobo/swarm/1db380a5e34bdb482ea47dfa34a35c08cc01e6fd/setup.cfg -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | import os.path 4 | 5 | from setuptools import setup, find_packages 6 | 7 | 8 | ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__))) 9 | 10 | 11 | setup( 12 | name="swarm", 13 | version="0.1", 14 | author="Hungpo DU", 15 | author_email="alecdu@gmail.com", 16 | url="https://github.com/duhoobo/swarm", 17 | description="A simple TCP benchmarking framework built upon gevent", 18 | long_description=open(os.path.join(ROOT, "README.md")).read(), 19 | license="PSF", 20 | keywords="gevent benchmark framework", 21 | 22 | packages=find_packages( 23 | exclude=("tests", "tests.*") 24 | ), 25 | 26 | include_package_data=True, 27 | 28 | install_requires=[ 29 | "click>=6.2", 30 | "gevent>=1.0.2", 31 | "PyYAML>=3.11", 32 | ], 33 | entry_points={ 34 | "console_scripts": [ 35 | "swarm = swarm.server:run" 36 | ] 37 | }, 38 | 39 | classifiers=[], 40 | ) 41 | -------------------------------------------------------------------------------- /swarm/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/duhoobo/swarm/1db380a5e34bdb482ea47dfa34a35c08cc01e6fd/swarm/__init__.py -------------------------------------------------------------------------------- /swarm/command.py: -------------------------------------------------------------------------------- 1 | 2 | class Command(object): 3 | pass 4 | 5 | 6 | class CommandINCR(Command): 7 | def __init__(self, amount): 8 | self.amount = amount 9 | 10 | 11 | class CommandDECR(Command): 12 | def __init__(self, amount): 13 | self.amount = amount 14 | 15 | 16 | class CommandQUIT(Command): 17 | pass 18 | 19 | 20 | class CommandSTATUS(Command): 21 | def __init__(self, id, status, prev_status): 22 | self.id = id 23 | self.status = status 24 | self.prev_status = prev_status 25 | -------------------------------------------------------------------------------- /swarm/errors.py: -------------------------------------------------------------------------------- 1 | 2 | class BenchException(Exception): 3 | pass 4 | 5 | 6 | class ServerClosed(BenchException): 7 | pass 8 | 9 | 10 | class CloseForcibly(BenchException): 11 | pass 12 | 13 | 14 | class ActionTimeout(BenchException): 15 | pass 16 | 17 | 18 | class InvalidReply(BenchException): 19 | pass 20 | 21 | 22 | class InvalidScript(BenchException): 23 | pass 24 | -------------------------------------------------------------------------------- /swarm/fakeclient.py: -------------------------------------------------------------------------------- 1 | import time 2 | from random import randint 3 | import socket 4 | from socket import ( 5 | create_connection, SOL_SOCKET, SO_RCVBUF, SO_SNDBUF 6 | ) 7 | from gevent import Greenlet, GreenletExit 8 | 9 | import swarm.script as script 10 | from swarm.command import CommandSTATUS 11 | from swarm.errors import BenchException, ServerClosed, CloseForcibly 12 | 13 | 14 | INITIAL = 100000 15 | STARTED = INITIAL + 1 16 | CONNECT = INITIAL + 2 17 | ACTING = INITIAL + 3 18 | STANDBY = INITIAL + 4 19 | TIMEDOUT = INITIAL + 5 20 | FATAL = INITIAL + 6 21 | KILLED = INITIAL + 7 22 | 23 | _readable_status = { 24 | INITIAL: "initial", 25 | STARTED: "started", 26 | CONNECT: "connect", 27 | ACTING: "acting", 28 | STANDBY: "standby", 29 | TIMEDOUT: "timedout", 30 | FATAL: "fatal", 31 | KILLED: "killed" 32 | } 33 | 34 | 35 | def readable_status(status): 36 | return _readable_status.get(status, str(status)) 37 | 38 | 39 | class FakeClient(object): 40 | """ 41 | A fake client with persistent connection. 42 | 43 | Driven by a dedicated greenlet, it will die trying to operate by the rules 44 | from the script orderly, round and round. 45 | """ 46 | 47 | def __init__(self, swarm, server, script_): 48 | self._swarm = swarm 49 | self._server = server 50 | self._socket = None 51 | self._greenlet = Greenlet(self._run) 52 | 53 | self._status = INITIAL 54 | self._prev_status = None 55 | self._script = script_ 56 | self._id = id(self) 57 | 58 | def _report(self, status): 59 | """ 60 | Report to swarm immediately on status change 61 | """ 62 | if status != self._status: 63 | self._swarm.commit(CommandSTATUS(self._id, status, self._status)) 64 | self._status, self._prev_status = (status, self._status) 65 | 66 | def _reconnect_server(self): 67 | """ 68 | Die trying 69 | """ 70 | while True: 71 | try: 72 | # To scatter connect requests 73 | time.sleep(randint(1, 20)) 74 | 75 | self._report(CONNECT) 76 | self._disconnect_server() 77 | self._socket = create_connection(self._server, 3) 78 | self._socket.setsockopt(SOL_SOCKET, SO_RCVBUF, 128) 79 | self._socket.setsockopt(SOL_SOCKET, SO_SNDBUF, 1024) 80 | 81 | break 82 | 83 | except socket.error as e: 84 | # A fact: `socket.timeout`, `socket.herror`, and 85 | # `socket.gaierror` are all subclasses of `socket.error`. 86 | self._report(e.args[0]) 87 | continue 88 | 89 | def _disconnect_server(self): 90 | if self._socket: 91 | self._socket.close() 92 | self._socket = None 93 | 94 | def _run(self): 95 | try: 96 | self._report(STARTED) 97 | self._reconnect_server() 98 | 99 | while True: 100 | try: 101 | self._report(ACTING) 102 | script.execute(self, self._script) 103 | self._report(STANDBY) 104 | 105 | except (socket.error, BenchException) as e: 106 | self._report(e.args[0]) 107 | self._reconnect_server() 108 | 109 | except GreenletExit: 110 | self._report(KILLED) 111 | 112 | except: 113 | self._report(FATAL) 114 | # let gevent print this exception 115 | raise 116 | 117 | finally: 118 | self._disconnect_server() 119 | 120 | def start(self): 121 | self._greenlet.start() 122 | 123 | def stop(self): 124 | self._greenlet.kill() 125 | self._greenlet.join() 126 | 127 | def send_for_reply(self, data, reply_parser): 128 | """ 129 | Called by object of Script. 130 | 131 | Exceptions raised here should be handled in `_run`. 132 | """ 133 | self._socket.sendall(data) 134 | 135 | need = reply_parser(None) 136 | while need > 0: 137 | data = self._socket.recv(need) 138 | if not data: 139 | raise ServerClosed("server closed") 140 | 141 | need = reply_parser(data) 142 | 143 | def send_noreply(self, data): 144 | self._socket.sendall(data) 145 | 146 | def close_connection(self): 147 | raise CloseForcibly("client closed") 148 | -------------------------------------------------------------------------------- /swarm/protocol.py: -------------------------------------------------------------------------------- 1 | 2 | def reply_parser_crlf(): 3 | # To allow nested function to modify local variables defined in outer 4 | # scope. 5 | last_byte = [""] 6 | 7 | def reply_parser(data=None): 8 | if data is None: 9 | return 1 10 | elif data == "\r": 11 | last_byte[0] = data 12 | return 1 13 | elif data == "\n" and last_byte[0] == "\r": 14 | return 0 15 | else: 16 | return 1 17 | 18 | return reply_parser 19 | -------------------------------------------------------------------------------- /swarm/script.py: -------------------------------------------------------------------------------- 1 | import yaml 2 | from time import sleep 3 | from inspect import getcallargs 4 | from random import randint 5 | from importlib import import_module 6 | 7 | from swarm.errors import InvalidScript 8 | 9 | 10 | options = { 11 | "command": { 12 | "must": True, "type": str, "default": None, 13 | "range": False, 14 | }, 15 | "args": { 16 | "must": False, "type": dict, "default": {}, 17 | "range": False, 18 | }, 19 | "rounds": { 20 | "must": False, "type": (int, list), "default": 1, 21 | "range": True, 22 | }, 23 | "sleep_between_rounds": { 24 | "must": False, "type": (int, list), "default": 1, 25 | "range": True, 26 | }, 27 | "sleep_after_action": { 28 | "must": False, "type": (int, list), "default": 1, 29 | "range": True, 30 | }, 31 | } 32 | 33 | 34 | def _random_or_straight(value): 35 | if isinstance(value, list): 36 | return randint(value[0], value[1]) 37 | else: 38 | return value 39 | 40 | 41 | def execute(client, script): 42 | protocol = script["protocol"] 43 | 44 | for action in script["actions"]: 45 | name = action["command"] 46 | args = action["args"] 47 | rounds = _random_or_straight(action["rounds"]) 48 | sleep_between_rounds = action["sleep_between_rounds"] 49 | sleep_after_action = action["sleep_after_action"] 50 | 51 | call = getattr(protocol, name) 52 | 53 | for n in xrange(rounds): 54 | call(client, **args) 55 | sleep(_random_or_straight(sleep_between_rounds)) 56 | 57 | sleep(_random_or_straight(sleep_after_action)) 58 | 59 | 60 | def load(script_file): 61 | with open(script_file, "r") as f: 62 | script = yaml.load(f) 63 | 64 | try: 65 | protocol_module_name = script["protocol"] 66 | except: 67 | raise InvalidScript("protocol is missing") 68 | 69 | try: 70 | protocol = import_module(protocol_module_name) 71 | except ImportError: 72 | raise InvalidScript("protocol `%s` not exists" 73 | % script["protocol"]) 74 | 75 | for n, action in enumerate(script["actions"]): 76 | for key, value in options.iteritems(): 77 | if value["must"] and key not in action: 78 | raise InvalidScript("%s is required in action %d" % 79 | (key, n)) 80 | if key in action: 81 | if not isinstance(action[key], value["type"]): 82 | raise InvalidScript("Invalid type for %s in" 83 | "action %d" % (key, n)) 84 | 85 | if value["range"] and isinstance(action[key], list): 86 | a, b = action[key] 87 | if (not isinstance(a, int) or not isinstance(b, int) 88 | or a > b): 89 | raise InvalidScript("Invalid range for %s in" 90 | "action %d" % (key, n)) 91 | else: 92 | action[key] = value["default"] 93 | try: 94 | func = getattr(protocol, action["command"]) 95 | getcallargs(func, None, **action["args"]) 96 | except AttributeError: 97 | raise InvalidScript("command '%s' not implemented yet" % 98 | action["action"]) 99 | except TypeError: 100 | raise InvalidScript("Invalid args for command '%s'" % 101 | action["command"]) 102 | 103 | script["protocol"] = protocol 104 | 105 | if not script: 106 | raise InvalidScript("script file empty") 107 | 108 | return script 109 | -------------------------------------------------------------------------------- /swarm/server.py: -------------------------------------------------------------------------------- 1 | import gevent.monkey 2 | gevent.monkey.patch_all() 3 | 4 | import sys 5 | import click 6 | from resource import getrlimit, setrlimit, RLIMIT_NOFILE 7 | 8 | from signal import ( 9 | signal, SIG_IGN, SIGHUP, SIGINT, SIGTERM, SIGPIPE 10 | ) 11 | from gevent import get_hub 12 | from gevent.server import StreamServer 13 | from gevent.queue import Queue 14 | from collections import defaultdict 15 | 16 | try: 17 | from cStringIO import StringIO as BufferIO 18 | except ImportError: 19 | from StringIO import StringIO as BufferIO 20 | 21 | import swarm.script as script 22 | from swarm.fakeclient import FakeClient, readable_status 23 | from swarm.command import ( 24 | CommandINCR, CommandDECR, CommandQUIT, CommandSTATUS 25 | ) 26 | 27 | 28 | class TCPGateway(StreamServer): 29 | banner = """ 30 | =============================== 31 | Welcome to Swarm Remote Console 32 | =============================== 33 | """ 34 | 35 | usage = ( 36 | "Usage:\n" 37 | "incr 1000\\n -- Start 1000 new connections\n" 38 | "decr 100\\n -- Close 100 random connections\n" 39 | "stop\\n -- Stop swarm benchmarker remotely\n" 40 | "quit\\n -- Close this control session\n" 41 | "help\\n -- Help\n" 42 | ) 43 | 44 | def __init__(self, swarm, listener): 45 | super(TCPGateway, self).__init__(listener, handle=self._handler, 46 | spawn=2) 47 | self._swarm = swarm 48 | print "Listening on %s" % str(listener) 49 | 50 | def _handler(self, client_sock, address): 51 | """ 52 | Called by a new greenlet. 53 | 54 | `client_sock` will be closed by `StreamServer` after this handler 55 | exits. But, this greenlet/connection won't be kill naturally if we 56 | don't use `pool` with StreamServer`. 57 | """ 58 | try: 59 | self._welcome(client_sock) 60 | 61 | request = BufferIO() 62 | while True: 63 | char = client_sock.recv(1) 64 | if not char: 65 | # peer closed 66 | break 67 | 68 | if char == "\n": 69 | if not self._process_request(client_sock, 70 | request.getvalue()): 71 | break 72 | request.seek(0) 73 | request.truncate() 74 | else: 75 | request.write(char) 76 | 77 | except Exception as e: 78 | # just let this greenlet die 79 | print e 80 | pass 81 | 82 | def _welcome(self, client_sock): 83 | client_sock.sendall(self.banner + "\n") 84 | client_sock.sendall(self.usage + "\n") 85 | 86 | def _process_request(self, client_sock, request): 87 | if not request: 88 | return True 89 | 90 | print request 91 | 92 | parts = request.split() 93 | cmd = parts[0].lower() 94 | standby, reply = True, None 95 | 96 | while True: 97 | if cmd in ["incr", "decr"]: 98 | try: 99 | amount = int(parts[1]) 100 | except: 101 | reply = "Invalid Command" 102 | break 103 | 104 | self._swarm.commit( 105 | CommandINCR(amount) if cmd == "incr" else 106 | CommandDECR(-amount) 107 | ) 108 | 109 | reply = "OK" 110 | 111 | elif cmd == "stop": 112 | self._swarm.commit(CommandQUIT()) 113 | reply = "OK" 114 | 115 | elif cmd == "quit": 116 | standby = False 117 | 118 | else: 119 | reply = self.usage 120 | 121 | break 122 | 123 | if reply: 124 | client_sock.sendall(reply + "\n") 125 | 126 | return standby 127 | 128 | 129 | class Counters(object): 130 | def __init__(self): 131 | self._total = 0 132 | self._realtime = defaultdict(set) 133 | self._aggregates = defaultdict(int) 134 | 135 | def update(self, id, status, prev_status): 136 | if prev_status: 137 | self._realtime[prev_status].discard(id) 138 | self._realtime[status].add(id) 139 | self._aggregates[status] += 1 140 | 141 | @property 142 | def total(self): 143 | return self._total 144 | 145 | @total.setter 146 | def total(self, value): 147 | self._total = value 148 | 149 | def reset(self): 150 | self._total = 0 151 | self._realtime = defaultdict(set) 152 | self._aggregates = defaultdict(int) 153 | 154 | def __str__(self): 155 | r = ", ".join(["%s: %d" % (readable_status(s), len(t)) for s, t in 156 | self._realtime.iteritems() if t]) 157 | a = ", ".join(["%s: %d" % (readable_status(s), c) for s, c in 158 | self._aggregates.iteritems()]) 159 | 160 | return ("Realtime: total: %d, %s\nAggregates: %s" 161 | % (self._total, r or "empty", a or "empty")) 162 | 163 | def __repr__(self): 164 | return "" 165 | 166 | 167 | class Swarm(object): 168 | def __init__(self, listener, server, script_file): 169 | self._continue = True 170 | self._stat_timer = None 171 | self._server = server 172 | 173 | # To save some memory, make all fakeclients share the same script 174 | self._script = script.load(script_file) 175 | self._raise_rlimit() 176 | self._register_signals() 177 | 178 | self._tcpgateway = TCPGateway(self, listener) 179 | self._commandq = Queue() 180 | self._counters = Counters() 181 | self._fakeclients = {} 182 | 183 | print "Try to bench server at %s using %s" % (server, script_file) 184 | 185 | def _raise_rlimit(self): 186 | soft, hard = getrlimit(RLIMIT_NOFILE) 187 | print "Default RLIMIT_NOFILE: (%d, %d)" % (soft, hard) 188 | setrlimit(RLIMIT_NOFILE, (hard, hard)) 189 | 190 | def _start_timers(self): 191 | loop = get_hub().loop 192 | self._stat_timer = loop.timer(5, 10) 193 | self._stat_timer.start(self._regularly_stat) 194 | 195 | def _stop_timers(self): 196 | self._stat_timer.stop() 197 | 198 | def _register_signals(self): 199 | signal(SIGHUP, SIG_IGN) 200 | signal(SIGPIPE, SIG_IGN) 201 | signal(SIGINT, self._stop) 202 | signal(SIGTERM, self._stop) 203 | 204 | def _regularly_stat(self): 205 | print "<%s" % ("-" * 10,) 206 | print self._counters 207 | print "%s>" % ("-" * 20,) 208 | 209 | def _process_command(self, command): 210 | if isinstance(command, (CommandINCR, CommandDECR)): 211 | self._spawn(command.amount) 212 | 213 | elif isinstance(command, CommandQUIT): 214 | self._continue = False 215 | 216 | elif isinstance(command, CommandSTATUS): 217 | self._counters.update(command.id, command.status, 218 | command.prev_status) 219 | 220 | else: 221 | # Just ignore unrecognized commands 222 | pass 223 | 224 | def _spawn(self, amount=1): 225 | toadd, amount = (amount > 0), abs(amount) 226 | 227 | if toadd: 228 | for n in xrange(amount): 229 | client = FakeClient(self, self._server, self._script) 230 | client.start() 231 | self._fakeclients[id(client)] = client 232 | self._counters.total += 1 233 | 234 | else: 235 | for n in xrange(amount): 236 | try: 237 | (__, client) = self._fakeclients.popitem() 238 | client.stop() 239 | self._counters.total -= 1 240 | except: 241 | break 242 | 243 | def _reap(self): 244 | for __, client in self._fakeclients.iteritems(): 245 | client.stop() 246 | 247 | self._counters.reset() 248 | 249 | def commit(self, command, block=True, timeout=None): 250 | self._commandq.put(command, block, timeout) 251 | 252 | def run_forever(self): 253 | self._tcpgateway.start() 254 | self._start_timers() 255 | 256 | while self._continue: 257 | command = self._commandq.get(block=True, timeout=None) 258 | self._process_command(command) 259 | 260 | print "Stopping..." 261 | 262 | self._stop_timers() 263 | self._tcpgateway.stop() 264 | self._reap() 265 | 266 | def _stop(self, signum, frame): 267 | print "Got signal %s, to quit" % signum 268 | self.commit(CommandQUIT()) 269 | 270 | 271 | class AddressParamType(click.ParamType): 272 | name = "ADDRESS" 273 | 274 | def convert(self, value, param, ctx): 275 | try: 276 | ip, port = value.split(":") 277 | return (ip, int(port)) 278 | except: 279 | self.fail("%s is not a valid address" % value, param, ctx) 280 | 281 | 282 | ADDRESS = AddressParamType() 283 | 284 | 285 | @click.command(context_settings=dict(help_option_names=["-h", "--help"])) 286 | @click.option("--listen", "-l", required=False, type=ADDRESS, 287 | default="127.0.0.1:8411", metavar="", 288 | help="Address the console is listening on") 289 | @click.option("--server", "-r", required=True, type=ADDRESS, 290 | metavar="", 291 | help="Address of remote server you want to bench") 292 | @click.option("--script", "-s", required=True, 293 | type=click.Path(exists=True, dir_okay=False), 294 | help="Script file which defines a series of actions") 295 | @click.option("--protocol-dir", "-p", required=True, 296 | type=click.Path(exists=True, dir_okay=True, resolve_path=True), 297 | help="Directory where you put your custom protocol modules") 298 | def run(listen, server, script, protocol_dir): 299 | """ 300 | `swarm` is a simple benchmarking framework built upon `gevent`. It can be 301 | used to generate massive simultaneous and persistent TCP connections to a 302 | server, while each connection interacts with the server using your custom 303 | protocol. 304 | """ 305 | if protocol_dir not in sys.path: 306 | sys.path.append(protocol_dir) 307 | 308 | swarm = Swarm(listen, server, script) 309 | swarm.run_forever() 310 | --------------------------------------------------------------------------------