class Redirector(Greenlet):
def __init__(self, msg):
self.sock_type = msg.sock_type
self.info = msg
self.sock = socket.socket(socket.AF_INET, self.sock_type)
self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
self.links = Group()
super(Redirector, self).__init__()
def _run(self):
self.sock.bind(self.info.local)
self.sock.listen(64)
while True:
cli, addr = self.sock.accept()
self.links.start(Linker(self.info.remote, self.sock_type, cli))
def kill(self):
self.links.kill()
super(Redirector, self).kill()
self.sock.close()
def dump(self):
return dict(protocol=self.info.protocol,
local='%s:%d' % self.info.local,
remote='%s:%d' % self.info.remote)
def run(self):
import gevent.queue
from gevent.pool import Group
gq = gevent.queue.Queue()
group = Group()
group.spawn(transfer_queue, self.q, gq)
run_from_queue(group, gq)
def _push_to_target(self, targets):
"""Get a batch of elements from the queue, and push it to the targets.
This function returns True if it proceeded all the elements in
the queue, and there isn't anything more to read.
"""
if self.queue.empty():
return 0 # nothing
batch = []
pushed = 0
# collecting a batch
while len(batch) < self.batch_size:
item = self.queue.get()
if item == 'END':
pushed += 1 # the 'END' item
break
batch.append(item)
if len(batch) != 0:
greenlets = Group()
for plugin in targets:
green = greenlets.spawn(self._put_data, plugin, batch)
green.link_exception(partial(self._error,
exception.InjectError, plugin))
greenlets.join()
pushed += len(batch)
return pushed
def _enqueue_children(resources, seed, parent):
threads = Group()
for resource in [r for r in resources if r.parent == parent.name]:
thread = threads.spawn(_process_resource, resources, seed, resource)
thread.link_exception(_create_error_handler(resource.collection))
threads.join()
class WSGIServer(_WSGIServer):
handler_class = WebSocketWSGIHandler
def __init__(self, *args, **kwargs):
"""
WSGI server that simply tracks websockets
and send them a proper closing handshake
when the server terminates.
Other than that, the server is the same
as its :class:`gevent.pywsgi.WSGIServer`
base.
"""
_WSGIServer.__init__(self, *args, **kwargs)
self._websockets = Group()
def link_websocket_to_server(self, websocket):
logger.info("Managing websocket %s" % format_addresses(websocket))
self._websockets.spawn(websocket.run)
def stop(self, *args, **kwargs):
logger.info("Terminating server and all connected websockets")
for greenlet in self._websockets:
try:
websocket = greenlet._run.im_self
if websocket:
websocket.close(1001, 'Server is shutting down')
except:
pass
_WSGIServer.stop(self, *args, **kwargs)
def poll(request):
out = {}
out['res'] = {}
ids = request.GET.get('ids')
done = True
if ids:
group = Group()
greenlets = []
split = ids.split(',')
for id in split:
greenlet = CheckGreenlet(id, group)
greenlets.append(greenlet)
group.join()
for greenlet in greenlets:
out['res'][greenlet.id] = greenlet.result
if not greenlet.done:
done = False
out['done'] = done
return HttpResponse(json.dumps(out), mimetype='application/json')
class CrawlBase(object):
spider_count = 0
def __init__(self):
self.group = Group()
self.queue = Queue()
def read_seed(self, file='seeds.txt'):
with open(file) as f:
for line in f:
if len(line) > 0 and line != "\n":
yield line.strip()
else:
return
def dispatch(self):
for url in self.read_seed():
g = gevent.spawn(Spider, self, url)
self.group.add(g)
self.group.join()
def harvest(self):
try:
while True:
content = self.queue.get(timeout=2)
print(content)
except Empty:
pass
def _search(responder, id, search, max_results):
todo = dict()
group = Group()
for ctx in search.contexts:
if len(ctx.results) > max_results/len(search.contexts):
continue
todo[ctx.name] = 0
group.spawn(_run, responder, search, ctx, todo)
group.join()
for ctx in search.contexts[:]:
if not ctx.results:
search.remove_ctx(ctx)
results = list()
while search.contexts and (max_results is None or len(results) < max_results):
ctx = search.next_result_ctx()
if not ctx.results:
break
results.append(ctx.results.pop(0))
display = search.display
for ctx in search.contexts[:]:
if ctx.next is None:
search.remove_ctx(ctx)
if search.more:
cache[search.id] = search
elif search.id in cache:
del cache[search.id]
return dict(id=id, search_id=search.id, display=display, more=search.more, results=results)
def test_quiet_group():
from gevent.pool import Group
group = Group()
group.spawn(divide_by_zero)
group.spawn(divide_by_zero)
with raises_from(ZeroDivisionError, 'divide_by_zero'):
group.join(raise_error=True)
def test_greenlet(self):
queue = JoinableQueue()
requests_done = Event()
g = Greenlet(self._producer, queue, FirstService(), 'Terminator')
h = Greenlet(self._producer, queue, SecondService(), 'Terminator')
i = Greenlet(self._producer, queue, ThirdService(), 'Terminator')
requests = Group()
for request in g, h, i:
requests.add(request)
log.debug('before spawn')
c = spawn(
self._consumer,
done=requests_done,
queue=queue,
)
[r.start() for r in requests]
log.debug('after spawn')
requests.join()
requests_done.set()
log.debug('requests are done')
c.join()
log.debug('consumer is done')
class SlaveLocustRunner(DistributedLocustRunner):
def __init__(self, *args, **kwargs):
super(SlaveLocustRunner, self).__init__(*args, **kwargs)
self.client_id = socket.gethostname() + "_" + md5(str(time() + random.randint(0,10000))).hexdigest()
self.client = rpc.Client(self.master_host)
self.greenlet = Group()
self.greenlet.spawn(self.worker).link_exception()
self.client.send(Message("client_ready", None, self.client_id))
self.greenlet.spawn(self.stats_reporter).link_exception()
# register listener for when all locust users have hatched, and report it to the master node
def on_hatch_complete(count):
self.client.send(Message("hatch_complete", {"count":count}, self.client_id))
events.hatch_complete += on_hatch_complete
# register listener that adds the current number of spawned locusts to the report that is sent to the master node
def on_report_to_master(client_id, data):
data["user_count"] = self.user_count
events.report_to_master += on_report_to_master
# register listener that sends quit message to master
def on_quitting():
self.client.send(Message("quit", None, self.client_id))
events.quitting += on_quitting
# register listener thats sends locust exceptions to master
def on_locust_error(locust, e, tb):
formatted_tb = "".join(traceback.format_tb(tb))
self.client.send(Message("exception", {"msg" : str(e), "traceback" : formatted_tb}, self.client_id))
events.locust_error += on_locust_error
def worker(self):
while True:
msg = self.client.recv()
if msg.type == "hatch":
self.client.send(Message("hatching", None, self.client_id))
job = msg.data
self.hatch_rate = job["hatch_rate"]
#self.num_clients = job["num_clients"]
self.num_requests = job["num_requests"]
self.host = job["host"]
self.hatching_greenlet = gevent.spawn(lambda: self.start_hatching(locust_count=job["num_clients"], hatch_rate=job["hatch_rate"]))
elif msg.type == "stop":
self.stop()
self.client.send(Message("client_stopped", None, self.client_id))
self.client.send(Message("client_ready", None, self.client_id))
def stats_reporter(self):
while True:
data = {}
events.report_to_master.fire(self.client_id, data)
try:
self.client.send(Message("stats", data, self.client_id))
except:
logger.error("Connection lost to master server. Aborting...")
break
gevent.sleep(SLAVE_REPORT_INTERVAL)
def sample_no_result():
group = Group()
def hello_from(n):
print 'Size of group %s' % len(group)
print 'Hello %s from Greenlet %s' % (n, id(getcurrent()))
group.map(hello_from, xrange(3))
class SlaveLocustRunner(DistributedLocustRunner):
def __init__(self, *args, **kwargs):
super(SlaveLocustRunner, self).__init__(*args, **kwargs)
self.client_id = socket.gethostname() + "_" + md5(str(time() + random.randint(0,10000))).hexdigest()
self.client = zmqrpc.Client(self.master_host)
self.greenlet = Group()
self.greenlet.spawn(self.worker).link_exception()
self.client.send({"type":"client_ready", "data":self.client_id})
self.greenlet.spawn(self.stats_reporter).link_exception()
# register listener for when all locust users have hatched, and report it to the master node
def on_hatch_complete(count):
self.client.send({"type":"hatch_complete", "data":{"client_id":self.client_id, "count":count}})
events.hatch_complete += on_hatch_complete
# register listener that adds the current number of spawned locusts to the report that is sent to the master node
def on_report_to_master(client_id, data):
data["user_count"] = self.user_count
events.report_to_master += on_report_to_master
# register listener that sends quit message to master
def on_quitting():
self.client.send({"type":"quit", "data":self.client_id})
events.quitting += on_quitting
def worker(self):
while True:
msg = self.client.recv()
if msg["type"] == "hatch":
self.client.send({"type":"hatching", "data":self.client_id})
job = msg["data"]
self.hatch_rate = job["hatch_rate"]
#self.num_clients = job["num_clients"]
self.num_requests = job["num_requests"]
self.host = job["host"]
self.hatching_greenlet = gevent.spawn(lambda: self.start_hatching(locust_count=job["num_clients"], hatch_rate=job["hatch_rate"]))
elif msg["type"] == "stop":
self.stop()
self.client.send({"type":"client_stopped", "data":self.client_id})
self.client.send({"type":"client_ready", "data":self.client_id})
def stats_reporter(self):
while True:
data = {}
events.report_to_master.fire(self.client_id, data)
report = {
"client_id": self.client_id,
"data": data,
}
try:
self.client.send({"type":"stats", "data":report})
except:
logger.error("Connection lost to master server. Aborting...")
break
gevent.sleep(SLAVE_REPORT_INTERVAL)
def search(self, query, queue):
"""Schedules a search and returns the related task information."""
group = Group()
for provider in self._providers:
group.add(spawn(self._search_wrapper, provider, query, queue, self._middleware))
return group
def terminate():
g = Group()
for torrent in torrents.values():
g.spawn(torrent.save_resume_data)
try:
g.join(timeout=5)
except:
pass
for torrent in torrents.values():
torrent.remove()
def add(self, greenlet, blocking=True):
acquired = self._semaphore.acquire(blocking=blocking)
if not acquired:
return False
try:
Group.add(self, greenlet)
except:
self._semaphore.release()
raise
return True
def execute(resources, seed=None):
threads = Group()
for resource in [r for r in resources if not r.parent]:
thread = threads.spawn(_process_resource, resources, seed, resource)
thread.link_exception(_create_error_handler(resource.collection))
threads.join()
if len(_errors):
sys.exit(1)
def spawn_greenlets(conf):
"""Some sugar to wrap up all of your greenlets."""
group = Group()
for args in conf:
group.spawn(*args)
try:
while True:
gevent.sleep(1)
except KeyboardInterrupt:
pass
def add(self, greenlet):
acquired = self._semaphore.acquire(blocking=False, timeout=self._timeout)
# XXX(Mouad): Checking directly for False because DummySemaphore always
# return None https://github.com/gevent/gevent/pull/544.
if acquired is False:
raise RejectExcecutionError("No more resource available to run %r" % greenlet)
try:
Group.add(self, greenlet)
except:
self._semaphore.release()
raise
def getStories(self):
"""
return a list of story dicts
"""
logger.debug('getStories: about to get {0} stories'.format(len(self.storyids)))
group = Group()
getstory = lambda storyid: firebase.get('/v0/item', storyid)
stories = group.map(getstory, self.storyids)
return stories
def main():
group.map(hello_from, xrange(3))
print('Ordered')
ogroup = Group()
for i in ogroup.imap(intensive, xrange(3)):
print(i)
print('Unordered')
igroup = Group()
for i in igroup.imap_unordered(intensive, xrange(3)):
print(i)
def create_connection(self, address, timeout=10):
ip_list = dnslib.dnsQuery(address[0])
# 尝试连接缓存
route_list = self.get_route_order_ping(address[0],address[1],None)
if route_list:
try:
route = route_list[0]
hit_ip = route['hit_ip']
if hit_ip in ip_list:
cache_timeout = route['tcp_ping']
if cache_timeout<1000:
cache_timeout = cache_timeout * 2
else:
cache_timeout = cache_timeout+1000
cache_timeout = int(math.ceil(cache_timeout/1000.0))
start_time = int(time.time() * 1000)
sock = self._direct_create_connection(address,hit_ip, cache_timeout)
t = int(time.time() * 1000) - start_time
logging.debug(u'[upstream][RouteCache]%s 缓存记录 连接 %s(%s):%s 命中。time:%s'%(self.get_display_name(),address[0],hit_ip,address[1],t))
self.update_route_ping(address[0],address[1],t,hit_ip)
return sock
else:
logging.debug(u'[upstream][RouteCache]%s 缓存记录 连接 %s(%s):%s IP 不匹配，放弃缓存。'%(self.get_display_name(),address[0],hit_ip,address[1]))
except:
t = int(time.time() * 1000) - start_time
info = traceback.format_exc()
logging.debug(
u'[upstream][RouteCache]%s 缓存记录 连接 %s(%s):%s 失败。time:%s' % (self.get_display_name(), address[0],hit_ip, address[1],t))
logging.debug('%s\r\n\r\n' % info)
# 缓存失败，连接全部
evt = Event()
group = Group()
aync_task = DirectAsyncTask(evt, None, group)
for ip in ip_list:
group.add(gevent.spawn(self._create_connection, aync_task, address,ip, timeout))
# 所有连接失败时发出通知
gevent.spawn(self._create_connection_all_end, aync_task)
evt.wait()
if aync_task.sock:
return aync_task.sock
else:
raise UpstreamConnectError()
def test_kill_processlet_group(proc):
group = Group()
group.greenlet_class = lets.Processlet
group.spawn(raise_when_killed)
group.spawn(raise_when_killed)
group.spawn(raise_when_killed)
group.join(0)
assert len(proc.children()) == 3
group.kill()
assert len(proc.children()) == 0
for job in group:
with pytest.raises(Killed):
job.get()
assert job.exit_code == 1
def handle(self, *args, **options):
plogger = logging.getLogger(__name__)
plogger.info("Parallel Update All Feeds")
new_posts_count = 0
start = datetime.now()
feed_urls = Feed.site_objects.all().values_list("url", flat=True)
pool = Group()
for result in pool.imap_unordered(process_feed, feed_urls):
new_posts_count += result
delta = datetime.now() - start
plogger.info("Added %s posts in %d seconds" % (new_posts_count, delta.seconds))
feeds_updated.send(sender=self, instance=self)
def _main():
user_crawler_group = Group()
for _ in xrange(GREENLET_COUNT):
user_crawler_group.spawn(analyze)
with open('ids.txt') as FILE:
for line in FILE:
id = line.strip()
cursor.execute('SELECT COUNT(1) as total_count FROM tb_xweibo_user_info WHERE uid = %s' % id)
result = cursor.fetchone()
if not result['total_count']:
users_fetch_queue.put(id)
user_crawler_group.join()
class Manager(object):
def __init__(self, config_addr, keyfile,
certfile, cacerts, backlog=10):
if isinstance(config_addr, basestring):
ip, port = config_addr.split(':')
config_addr = (ip, int(port))
self.keyfile = keyfile
self.certfile = certfile
self.cacerts = cacerts
self.config_addr = config_addr
self.backlog = backlog
self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
self.end_evt = Event()
self.clients = Group()
self.redirectors = {}
self.msg_processors = {
'redirect': self.m_redirect,
'list_redirect': self.m_list_redirect,
'drop_redirect': self.m_drop_redirect,
'shutdown': self.m_shutdown,
}
logging.info('manager initialized')
def run(self):
logging.info('manager start to run')
self.sock.bind(self.config_addr)
logging.info('manager bind to: %s:%d' % self.config_addr)
self.sock.listen(self.backlog)
accept_let = gevent.spawn(self.accept_let)
self.end_evt.wait()
logging.info('shutdown evt recved')
accept_let.kill()
self.clients.kill()
def accept_let(self):
while True:
sock, addr = self.sock.accept()
try:
sock = ssl.wrap_socket(sock, keyfile=self.keyfile,
certfile=self.certfile, server_side=True,
cert_reqs=ssl.CERT_REQUIRED, ca_certs=self.cacerts)
except ssl.SSLError, e:
print e
continue
self.clients.spawn(self.client_let, sock, addr)
def test_gevent(num, depth, period, settle):
global SWITCH_COUNT
import gevent
from gevent.pool import Group
from gevent.hub import sleep
horde = Group()
for x in xrange(num):
horde.spawn(stack_filler, depth, sleep)
gevent.sleep(settle)
print("settle period over, {:.2f} sw/sec, testing".format(SWITCH_COUNT/(1.0*settle)))
SWITCH_COUNT=0
gevent.sleep(period)
print("testing period over, {:.2f} sw/sec".format(SWITCH_COUNT/(1.0*period)))
horde.kill()
def test_parallel_queue_declare(self):
conn = Connection(self.amqp_url)
conn.connect()
channel = conn.allocate_channel()
def declare(name):
return channel.queue_declare(queue=name)
g = Group()
res = g.map(declare, queues)
assert len(res) == len(queues)
assert all(isinstance(r, FrameQueueDeclareOk) for r in res)
def Demo1():
g1 = gevent.spawn(talk, 'bar')
g2 = gevent.spawn(talk, 'foo')
g3 = gevent.spawn(talk, 'fizz')
group = Group()
group.add(g1)
group.add(g2)
group.join()
group.add(g3)
group.join()
def get_multihoster_account(task, multi_match, file):
if not account.config.use_useraccounts:
print "multihoster off, use_useraccounts false"
return
group = Group()
for pool in account.manager.values():
for acc in pool:
if acc.multi_account:
from . import manager
acc.hoster = manager.find_by_name(acc.name)
group.spawn(acc.boot)
group.join()
accounts = []
best_weight = 0
hostname = file.split_url.host
for pool in account.manager.values():
for acc in pool:
if acc._private_account:
continue
if not acc.multi_account:
continue
if hasattr(acc, 'premium') and not acc.premium:
continue
print acc.premium
if not multi_match(acc, hostname):
continue
try:
weight = acc.weight
except gevent.GreenletExit:
print "greenlet exit"
continue
if weight > best_weight:
accounts = []
best_weight = weight
bisect.insort(accounts, (acc.get_task_pool(task).full() and 1 or 0, len(accounts), acc))
if accounts:
return accounts[0][2]
"""try:
file.log.info('trying multihoster {}'.format(acc.name))
acc.hoster.get_download_context(file)
except gevent.GreenletExit:
raise
except BaseException as e:
log.exception(e)"""
else:
print "multi: no accounts found"
class ProxyPool(object):
"""Proxy Pool.
"""
def __init__(self, configfile='settings.yaml'):
self.configs = self._get_configs(configfile)
self.rdb = redis.StrictRedis(db=self.configs['RDB'])
self.db_zproxy = self.configs['DB_ZPROXY']
self.pool = Pool(self.configs['MAX_CONCURRENCY'])
self.group = Group()
self.wrong_value = self.configs['WRONG_VALUE']
self.init_value = self.configs['INIT_VALUE']
def _get_configs(self, configfile):
"""Return the configuration dict"""
# XXX: getting the path of the configuraion file needs improving
configpath = os.path.join('.', configfile)
with open(configpath, 'rb') as fp:
configs = yaml.load(fp)
return configs
def get_many(self, num=3, minscore=0, maxscore=None):
"""
Return a list of proxies including at most 'num' proxies
which socres are between 'minscore' and 'mascore'.
If there's no proxies matching, return an empty list.
"""
minscore = minscore
maxscore = maxscore or self.init_value
res = self.rdb.zrange(self.db_zproxy, minscore, maxscore)
if res:
random.shuffle(res) # for getting random results
if len(res) < num:
logging.warning(
"The number of proxies you want is less than %d" % (num, ))
return [proxy for proxy in res[:num]]
else:
logging.warning(
"There're no proxies which scores are between %d and %d" %
(minscore, maxscore))
return []
def get_one(self, minscore=0, maxscore=None):
"""
Return one proxy which score is between 'minscore'
and 'maxscore'.
If there's no proxy matching, return an empty string.
"""
minscore = minscore
maxscore = maxscore or self.init_value
res = self.get_many(num=1, minscore=minscore, maxscore=maxscore)
if res:
return res[0]
else:
return ''
def crawl_proxies(self):
"""Get proxies from vairous methods.
"""
statics = []
self._crawl_proxies_sites(statics=statics)
logging.info('Having add %d proxies' % (len(statics), ))
def _crawl_proxies_sites(self, statics=[]):
"""Get proxies from web pages."""
args = ((url, val['rules'], val['proxies'], statics)
for url, val in self.configs['PROXY_SITES'].iteritems())
self.pool.map(self._crawl_proxies_one_site, args)
def _crawl_proxies_one_site(self, args):
"""Get proxies (ip:port) from url and then write them into redis."""
url = args[0]
rules = args[1]
proxies = args[2]
headers = self.configs['HEADERS']
headers['Host'] = url.split('/', 3)[2]
logging.info('Begin crawl page %s' % (url, ))
res = requests.get(url, headers=headers, proxies=proxies)
encoding = res.encoding
html = etree.HTML(res.content)
proxies = []
len_rules = len(rules)
nodes = html.xpath(rules[0])
if nodes:
if len_rules == 1:
for node in nodes:
text = str(node.text).encode(encoding).strip()
if text:
proxies.append('http://%s' % (text, ))
elif len_rules == 2:
rule_1 = rules[1].split(',')
for node in nodes:
node = node.xpath(rule_1[0])
ip = str(node[1].text).encode(encoding).strip()
port = str(node[2].text).encode(encoding).strip() or '80'
if ip:
proxies.append('http://%s:%s' % (ip, port))
for proxy in proxies:
logging.info('Get proxy %s from %s' % (proxy, url))
args[3].append(proxy)
self.rdb.zadd(self.db_zproxy, self.init_value, proxy)
def validate_proxies(self):
"""Validate whether the proxies are alive."""
maxscore = self.init_value + self.wrong_value
proxies = self.rdb.zrange(self.db_zproxy, 0, maxscore)
statics_errors = []
args = [(proxy, statics_errors) for proxy in proxies]
self.group.map(self._validate_one_proxy, args)
logging.info('Have validated %d proxies, %d errors happened.' %
(len(proxies), len(statics_errors)))
def _validate_one_proxy(self, args):
"""Validate whether the proxy is still alive."""
test_sites = self.configs['TEST_SITES']
proxy = args[0]
time_res = []
args = [(site, proxy, time_res, args[1]) for site in test_sites]
self.group.map(self._test_one_site, args)
mean_time = sum(time_res) / len(test_sites)
logging.info('Validating %s, score is %d' % (proxy, mean_time))
self.rdb.zadd(self.db_zproxy, mean_time, proxy)
def _test_one_site(self, args):
url = args[0]
headers = self.configs['HEADERS']
headers['Host'] = url.split('/', 3)[2]
proxy = args[1]
proxies = {
'http': proxy,
}
start_time = time.time()
timeout = self.configs['TEST_TIMEOUT']
error = False
with Timeout(timeout, False):
try:
res = requests.get(url, headers=headers, proxies=proxies)
except Exception as e:
error = True
logging.error('%s: Error: %s' % (proxy, e.message))
if error:
args[2].append(self.configs['WRONG_VALUE'])
args[3].append(proxy)
else:
args[2].append(time.time() - start_time)
class LocustRunner(object):
def __init__(self, locust_classes, options):
self.options = options
self.locust_classes = locust_classes
self.hatch_rate = options.hatch_rate
self.num_clients = options.num_clients
self.host = options.host
self.locusts = Group()
self.greenlet = self.locusts
self.state = STATE_INIT
self.hatching_greenlet = None
self.exceptions = {}
self.stats = global_stats
# register listener that resets stats when hatching is complete
def on_hatch_complete(user_count):
self.state = STATE_RUNNING
if self.options.reset_stats:
logger.info("Resetting stats\n")
self.stats.reset_all()
events.hatch_complete += on_hatch_complete
@property
def request_stats(self):
return self.stats.entries
@property
def errors(self):
return self.stats.errors
@property
def user_count(self):
return len(self.locusts)
def weight_locusts(self, amount):
"""
Distributes the amount of locusts for each WebLocust-class according to it's weight
returns a list "bucket" with the weighted locusts
"""
bucket = []
weight_sum = sum((locust.weight for locust in self.locust_classes if locust.task_set))
residuals = {}
for locust in self.locust_classes:
if not locust.task_set:
warnings.warn("Notice: Found Locust class (%s) got no task_set. Skipping..." % locust.__name__)
continue
if self.host is not None:
locust.host = self.host
# create locusts depending on weight
percent = locust.weight / float(weight_sum)
num_locusts = int(round(amount * percent))
bucket.extend([locust for x in xrange(0, num_locusts)])
# used to keep track of the amount of rounding was done if we need
# to add/remove some instances from bucket
residuals[locust] = amount * percent - round(amount * percent)
if len(bucket) < amount:
# We got too few locust classes in the bucket, so we need to create a few extra locusts,
# and we do this by iterating over each of the Locust classes - starting with the one
# where the residual from the rounding was the largest - and creating one of each until
# we get the correct amount
for locust in [l for l, r in sorted(residuals.items(), key=lambda x:x[1], reverse=True)][:amount-len(bucket)]:
bucket.append(locust)
elif len(bucket) > amount:
# We've got too many locusts due to rounding errors so we need to remove some
for locust in [l for l, r in sorted(residuals.items(), key=lambda x:x[1])][:len(bucket)-amount]:
bucket.remove(locust)
return bucket
def spawn_locusts(self, spawn_count=None, wait=False):
if spawn_count is None:
spawn_count = self.num_clients
bucket = self.weight_locusts(spawn_count)
spawn_count = len(bucket)
if self.state == STATE_INIT or self.state == STATE_STOPPED:
self.state = STATE_HATCHING
self.num_clients = spawn_count
else:
self.num_clients += spawn_count
logger.info("Hatching and swarming %i clients at the rate %g clients/s..." % (spawn_count, self.hatch_rate))
occurrence_count = dict([(l.__name__, 0) for l in self.locust_classes])
def hatch():
sleep_time = 1.0 / self.hatch_rate
while True:
if not bucket:
logger.info("All locusts hatched: %s" % ", ".join(["%s: %d" % (name, count) for name, count in six.iteritems(occurrence_count)]))
events.hatch_complete.fire(user_count=self.num_clients)
return
locust = bucket.pop(random.randint(0, len(bucket)-1))
occurrence_count[locust.__name__] += 1
new_locust = locust()
def start_locust(_):
try:
new_locust.run(runner=self)
except GreenletExit:
pass
self.locusts.spawn(start_locust, new_locust)
if len(self.locusts) % 10 == 0:
logger.debug("%i locusts hatched" % len(self.locusts))
gevent.sleep(sleep_time)
hatch()
if wait:
self.locusts.join()
logger.info("All locusts dead\n")
def kill_locusts(self, kill_count):
"""
Kill a kill_count of weighted locusts from the Group() object in self.locusts
"""
bucket = self.weight_locusts(kill_count)
kill_count = len(bucket)
self.num_clients -= kill_count
logger.info("Killing %i locusts" % kill_count)
dying = []
for g in self.locusts:
for l in bucket:
if l == type(g.args[0]):
dying.append(g)
bucket.remove(l)
break
for g in dying:
self.locusts.killone(g)
events.hatch_complete.fire(user_count=self.num_clients)
def start_hatching(self, locust_count=None, hatch_rate=None, wait=False):
if self.state != STATE_RUNNING and self.state != STATE_HATCHING:
self.stats.clear_all()
self.stats.start_time = time()
self.exceptions = {}
events.locust_start_hatching.fire()
# Dynamically changing the locust count
if self.state != STATE_INIT and self.state != STATE_STOPPED:
self.state = STATE_HATCHING
if self.num_clients > locust_count:
# Kill some locusts
kill_count = self.num_clients - locust_count
self.kill_locusts(kill_count)
elif self.num_clients < locust_count:
# Spawn some locusts
if hatch_rate:
self.hatch_rate = hatch_rate
spawn_count = locust_count - self.num_clients
self.spawn_locusts(spawn_count=spawn_count)
else:
events.hatch_complete.fire(user_count=self.num_clients)
else:
if hatch_rate:
self.hatch_rate = hatch_rate
if locust_count is not None:
self.spawn_locusts(locust_count, wait=wait)
else:
self.spawn_locusts(wait=wait)
def stop(self):
# if we are currently hatching locusts we need to kill the hatching greenlet first
if self.hatching_greenlet and not self.hatching_greenlet.ready():
self.hatching_greenlet.kill(block=True)
if self.options.stop_timeout:
for locust_greenlet in self.locusts:
locust = locust_greenlet.args[0]
if locust._state == LOCUST_STATE_WAITING:
locust_greenlet.kill()
else:
locust._state = LOCUST_STATE_STOPPING
if not self.locusts.join(timeout=self.options.stop_timeout):
logger.info("Not all locusts finished their tasks & terminated in %s seconds. Killing them..." % self.options.stop_timeout)
self.locusts.kill(block=True)
self.state = STATE_STOPPED
events.locust_stop_hatching.fire()
def quit(self):
self.stop()
self.greenlet.kill(block=True)
def log_exception(self, node_id, msg, formatted_tb):
key = hash(formatted_tb)
row = self.exceptions.setdefault(key, {"count": 0, "msg": msg, "traceback": formatted_tb, "nodes": set()})
row["count"] += 1
row["nodes"].add(node_id)
self.exceptions[key] = row
# coding: utf-8
import random
import gevent
from gevent.pool import Group
from gevent.local import local
group = Group()
def hello(n):
gevent.sleep(3 - n)
print('Size of group %s' % len(group))
print('Hello %d from Greenlet %s' % (n, id(gevent.getcurrent())))
return n
# ex1, Group().map()
def gevent_ex1():
''' greenlet 交替执行，但返回结果（list）与输入顺序一致 [0, 1, 2] '''
res = group.map(hello, range(3))
print(type(res), res)
# ex2, Group().imap()
def gevent_ex2():
''' 返回结果为IMap, 延迟执行 '''
res = group.imap(hello, range(3))
print(type(res), list(res))
'''
Created on Dec 22, 2015
@author: lingandcs
'''
import gevent.monkey
gevent.monkey.patch_socket()
import gevent
from gevent import getcurrent
from gevent.pool import Group
group = Group()
def hello_from(n):
print('Size of group %s' % len(group))
print('Hello from Greenlet %s' % id(getcurrent()))
group.map(hello_from, xrange(3))
def intensive(n):
gevent.sleep(3 - n)
return 'task', n
print('Ordered')
ogroup = Group()
def __init__(self):
self.wrappers = []
self.pool = Group()
class ImapSyncMonitor(BaseMailSyncMonitor):
""" Top-level controller for an account's mail sync. Spawns individual
FolderSync greenlets for each folder.
Parameters
----------
poll_frequency: Integer
Seconds to wait between polling for the greenlets spawned
heartbeat: Integer
Seconds to wait between checking on folder sync threads.
refresh_flags_max: Integer
the maximum number of UIDs for which we'll check flags
periodically.
"""
def __init__(self,
account,
heartbeat=1,
poll_frequency=30,
retry_fail_classes=[],
refresh_flags_max=2000):
self.poll_frequency = poll_frequency
self.syncmanager_lock = db_write_lock(account.namespace.id)
self.refresh_flags_max = refresh_flags_max
provider_supports_condstore = provider_info(account.provider).get(
'condstore', False)
account_supports_condstore = getattr(account, 'supports_condstore',
False)
if provider_supports_condstore or account_supports_condstore:
self.sync_engine_class = CondstoreFolderSyncEngine
else:
self.sync_engine_class = FolderSyncEngine
self.folder_monitors = Group()
BaseMailSyncMonitor.__init__(self, account, heartbeat,
retry_fail_classes)
def prepare_sync(self):
"""Ensures that canonical tags are created for the account, and gets
and save Folder objects for folders on the IMAP backend. Returns a list
of tuples (folder_name, folder_id) for each folder we want to sync (in
order)."""
with mailsync_session_scope() as db_session:
account = db_session.query(ImapAccount).get(self.account_id)
Tag.create_canonical_tags(account.namespace, db_session)
with _pool(self.account_id).get() as crispin_client:
sync_folders = crispin_client.sync_folders()
save_folder_names(log, self.account_id,
crispin_client.folder_names(), db_session)
sync_folder_names_ids = []
for folder_name in sync_folders:
try:
id_, = db_session.query(Folder.id). \
filter(Folder.name == folder_name,
Folder.account_id == self.account_id).one()
sync_folder_names_ids.append((folder_name, id_))
except NoResultFound:
log.error("Missing Folder object when starting sync",
folder_name=folder_name)
raise MailsyncError(
"Missing Folder '{}' on account {}".format(
folder_name, self.account_id))
return sync_folder_names_ids
def sync(self):
""" Start per-folder syncs. Only have one per-folder sync in the
'initial' state at a time.
"""
sync_folder_names_ids = self.prepare_sync()
for folder_name, folder_id in sync_folder_names_ids:
log.info('initializing folder sync')
thread = self.sync_engine_class(
self.account_id, folder_name, folder_id, self.email_address,
self.provider_name, self.poll_frequency, self.syncmanager_lock,
self.refresh_flags_max, self.retry_fail_classes)
thread.start()
self.folder_monitors.add(thread)
while not self._thread_polling(thread) and \
not self._thread_finished(thread) and \
not thread.ready():
sleep(self.heartbeat)
# Allow individual folder sync monitors to shut themselves down
# after completing the initial sync.
if self._thread_finished(thread) or thread.ready():
log.info('folder sync finished/killed',
folder_name=thread.folder_name)
# NOTE: Greenlet is automatically removed from the group.
self.folder_monitors.join()
);
'''
def insert_train(train):
try:
if isinstance(train.get('Tags'), list):
cursor.execute("INSERT INTO train (Id, Title, Body, Tags) VALUES (%s, %s, %s, %s)", (
train.get('Id'),
train.get('Title'),
train.get('Body'),
train.get('Tags')
))
connection.commit()
except Exception as e:
print e
g1 = Group()
TASKS = 1000
for i, train in enumerate(mongo_db.Train.find()):
work = gevent.spawn(insert_train, train)
g1.add(work)
if i % TASKS == 0:
print i
g1.join()
g1.join()
print('work done')
def __init__(self, *args, **kwargs):
super(NoWebMasterLocustRunner, self).__init__(*args, **kwargs)
self.all_stats = {}
self.greenlet = Group()
self.max_num_requests = None
self.max_seconds_elapsed = None
class Consumer(object):
"""High level NSQ consumer.
A Consumer will connect to the nsqd tcp addresses or poll the provided
nsqlookupd http addresses for the configured topic and send signals to
message handlers connected to the :attr:`on_message` signal or provided by
``message_handler``.
Messages will automatically be finished when the message handle returns
unless :meth:`message.enable_async() <gnsq.Message.enable_async>` is called.
If an exception occurs or :class:`~gnsq.errors.NSQRequeueMessage` is raised,
the message will be requeued.
The Consumer will handle backing off of failed messages up to a configurable
``max_interval`` as well as automatically reconnecting to dropped
connections.
Example usage::
from gnsq import Consumer
consumer = gnsq.Consumer('topic', 'channel', 'localhost:4150')
@consumer.on_message.connect
def handler(consumer, message):
print 'got message:', message.body
consumer.start()
:param topic: specifies the desired NSQ topic
:param channel: specifies the desired NSQ channel
:param nsqd_tcp_addresses: a sequence of string addresses of the nsqd
instances this consumer should connect to
:param lookupd_http_addresses: a sequence of string addresses of the
nsqlookupd instances this consumer should query for producers of the
specified topic
:param name: a string that is used for logging messages (defaults to
``'gnsq.consumer.{topic}.{channel}'``)
:param message_handler: the callable that will be executed for each message
received
:param max_tries: the maximum number of attempts the consumer will make to
process a message after which messages will be automatically discarded
:param max_in_flight: the maximum number of messages this consumer will
pipeline for processing. this value will be divided evenly amongst the
configured/discovered nsqd producers
:param requeue_delay: the default delay to use when requeueing a failed
message
:param lookupd_poll_interval: the amount of time in seconds between querying
all of the supplied nsqlookupd instances. A random amount of time based
on this value will be initially introduced in order to add jitter when
multiple consumers are running
:param lookupd_poll_jitter: the maximum fractional amount of jitter to add
to the lookupd poll loop. This helps evenly distribute requests even if
multiple consumers restart at the same time.
:param low_ready_idle_timeout: the amount of time in seconds to wait for a
message from a producer when in a state where RDY counts are
re-distributed (ie. `max_in_flight` < `num_producers`)
:param max_backoff_duration: the maximum time we will allow a backoff state
to last in seconds. If zero, backoff wil not occur
:param backoff_on_requeue: if ``False``, backoff will only occur on
exception
:param **kwargs: passed to :class:`~gnsq.NsqdTCPClient` initialization
"""
def __init__(self,
topic,
channel,
nsqd_tcp_addresses=[],
lookupd_http_addresses=[],
name=None,
message_handler=None,
max_tries=5,
max_in_flight=1,
requeue_delay=0,
lookupd_poll_interval=60,
lookupd_poll_jitter=0.3,
low_ready_idle_timeout=10,
max_backoff_duration=128,
backoff_on_requeue=True,
**kwargs):
if not nsqd_tcp_addresses and not lookupd_http_addresses:
raise ValueError('must specify at least one nsqd or lookupd')
self.nsqd_tcp_addresses = parse_nsqds(nsqd_tcp_addresses)
self.lookupds = parse_lookupds(lookupd_http_addresses)
self.iterlookupds = cycle(self.lookupds)
self.topic = topic
self.channel = channel
self.max_tries = max_tries
self.max_in_flight = max_in_flight
self.requeue_delay = requeue_delay
self.lookupd_poll_interval = lookupd_poll_interval
self.lookupd_poll_jitter = lookupd_poll_jitter
self.low_ready_idle_timeout = low_ready_idle_timeout
self.backoff_on_requeue = backoff_on_requeue
self.max_backoff_duration = max_backoff_duration
self.conn_kwargs = kwargs
if name:
self.name = name
else:
self.name = '%s.%s.%s' % (__name__, self.topic, self.channel)
if message_handler is not None:
self.on_message.connect(message_handler, weak=False)
self.logger = logging.getLogger(self.name)
self._state = INIT
self._redistributed_ready_event = Event()
self._connection_backoffs = defaultdict(self._create_backoff)
self._message_backoffs = defaultdict(self._create_backoff)
self._connections = {}
self._workers = Group()
self._killables = Group()
@cached_property
def on_message(self):
"""Emitted when a message is received.
The signal sender is the consumer and the ``message`` is sent as an
argument. The ``message_handler`` param is connected to this signal.
"""
return blinker.Signal(doc='Emitted when a message is received.')
@cached_property
def on_response(self):
"""Emitted when a response is received.
The signal sender is the consumer and the ``response`` is sent as an
argument.
"""
return blinker.Signal(doc='Emitted when a response is received.')
@cached_property
def on_error(self):
"""Emitted when an error is received.
The signal sender is the consumer and the ``error`` is sent as an
argument.
"""
return blinker.Signal(doc='Emitted when a error is received.')
@cached_property
def on_finish(self):
"""Emitted after a message is successfully finished.
The signal sender is the consumer and the ``message_id`` is sent as an
argument.
"""
return blinker.Signal(doc='Emitted after the a message is finished.')
@cached_property
def on_requeue(self):
"""Emitted after a message is requeued.
The signal sender is the consumer and the ``message_id`` and ``timeout``
are sent as arguments.
"""
return blinker.Signal(doc='Emitted after the a message is requeued.')
@cached_property
def on_giving_up(self):
"""Emitted after a giving up on a message.
Emitted when a message has exceeded the maximum number of attempts
(``max_tries``) and will no longer be requeued. This is useful to
perform tasks such as writing to disk, collecting statistics etc. The
signal sender is the consumer and the ``message`` is sent as an
argument.
"""
return blinker.Signal(doc='Sent after a giving up on a message.')
@cached_property
def on_auth(self):
"""Emitted after a connection is successfully authenticated.
The signal sender is the consumer and the ``conn`` and parsed
``response`` are sent as arguments.
"""
return blinker.Signal(doc='Emitted when a response is received.')
@cached_property
def on_exception(self):
"""Emitted when an exception is caught while handling a message.
The signal sender is the consumer and the ``message`` and ``error`` are
sent as arguments.
"""
return blinker.Signal(doc='Emitted when an exception is caught.')
@cached_property
def on_close(self):
"""Emitted after :meth:`close`.
The signal sender is the consumer.
"""
return blinker.Signal(doc='Emitted after the consumer is closed.')
def start(self, block=True):
"""Start discovering and listing to connections."""
if self._state == INIT:
if not any(self.on_message.receivers_for(blinker.ANY)):
raise RuntimeError('no receivers connected to on_message')
self.logger.debug('starting %s...', self.name)
self._state = RUNNING
self.query_nsqd()
if self.lookupds:
self.query_lookupd()
self._killables.add(self._workers.spawn(self._poll_lookupd))
self._killables.add(self._workers.spawn(self._poll_ready))
else:
self.logger.warning('%s already started', self.name)
if block:
self.join()
def close(self):
"""Immediately close all connections and stop workers."""
if not self.is_running:
return
self._state = CLOSED
self.logger.debug('killing %d worker(s)', len(self._killables))
self._killables.kill(block=False)
self.logger.debug('closing %d connection(s)', len(self._connections))
for conn in self._connections:
conn.close_stream()
self.on_close.send(self)
def join(self, timeout=None, raise_error=False):
"""Block until all connections have closed and workers stopped."""
self._workers.join(timeout, raise_error)
@property
def is_running(self):
"""Check if consumer is currently running."""
return self._state == RUNNING
@property
def is_starved(self):
"""Evaluate whether any of the connections are starved.
This property should be used by message handlers to reliably identify
when to process a batch of messages.
"""
return any(conn.is_starved for conn in self._connections)
@property
def total_ready_count(self):
return sum(c.ready_count for c in self._connections)
@property
def total_in_flight(self):
return sum(c.in_flight for c in self._connections)
def query_nsqd(self):
self.logger.debug('querying nsqd...')
for address in self.nsqd_tcp_addresses:
address, port = address.split(':')
self.connect_to_nsqd(address, int(port))
def query_lookupd(self):
self.logger.debug('querying lookupd...')
lookupd = next(self.iterlookupds)
try:
producers = lookupd.lookup(self.topic)['producers']
self.logger.debug('found %d producers', len(producers))
except Exception as error:
self.logger.warning('Failed to lookup %s on %s (%s)', self.topic,
lookupd.address, error)
return
for producer in producers:
self.connect_to_nsqd(producer['broadcast_address'],
producer['tcp_port'])
def _poll_lookupd(self):
try:
delay = self.lookupd_poll_interval * self.lookupd_poll_jitter
gevent.sleep(random.random() * delay)
while True:
gevent.sleep(self.lookupd_poll_interval)
self.query_lookupd()
except gevent.GreenletExit:
pass
def _poll_ready(self):
try:
while True:
if self._redistributed_ready_event.wait(5):
self._redistributed_ready_event.clear()
self._redistribute_ready_state()
except gevent.GreenletExit:
pass
def _redistribute_ready_state(self):
if not self.is_running:
return
if len(self._connections) > self.max_in_flight:
ready_state = self._get_unsaturated_ready_state()
else:
ready_state = self._get_saturated_ready_state()
for conn, count in ready_state.items():
if conn.ready_count == count:
self.logger.debug('[%s] RDY count already %d', conn, count)
continue
self.logger.debug('[%s] sending RDY %d', conn, count)
try:
conn.ready(count)
except NSQSocketError as error:
self.logger.warning('[%s] RDY %d failed (%r)', conn, count,
error)
def _get_unsaturated_ready_state(self):
ready_state = {}
active = []
for conn, state in self._connections.items():
if state == BACKOFF:
ready_state[conn] = 0
elif state in (RUNNING, THROTTLED):
active.append(conn)
random.shuffle(active)
for conn in active[self.max_in_flight:]:
ready_state[conn] = 0
for conn in active[:self.max_in_flight]:
ready_state[conn] = 1
return ready_state
def _get_saturated_ready_state(self):
ready_state = {}
active = []
now = time.time()
for conn, state in self._connections.items():
if state == BACKOFF:
ready_state[conn] = 0
elif state == THROTTLED:
ready_state[conn] = 1
elif state == RUNNING:
if (now - conn.last_message) > self.low_ready_idle_timeout:
self.logger.info(
'[%s] idle connection, giving up RDY count', conn)
ready_state[conn] = 1
else:
active.append(conn)
if not active:
return ready_state
ready_available = self.max_in_flight - sum(ready_state.values())
connection_max_in_flight = ready_available // len(active)
for conn in active:
ready_state[conn] = connection_max_in_flight
for conn in random.sample(active, ready_available % len(active)):
ready_state[conn] += 1
return ready_state
def redistribute_ready_state(self):
self._redistributed_ready_event.set()
def connect_to_nsqd(self, address, port):
if not self.is_running:
return
conn = NsqdTCPClient(address, port, **self.conn_kwargs)
if conn in self._connections:
self.logger.debug('[%s] already connected', conn)
return
self._connections[conn] = INIT
self.logger.debug('[%s] connecting...', conn)
conn.on_message.connect(self.handle_message)
conn.on_response.connect(self.handle_response)
conn.on_error.connect(self.handle_error)
conn.on_finish.connect(self.handle_finish)
conn.on_requeue.connect(self.handle_requeue)
conn.on_auth.connect(self.handle_auth)
try:
conn.connect()
conn.identify()
if conn.max_ready_count < self.max_in_flight:
msg = ('[%s] max RDY count %d < consumer max in flight %d, '
'truncation possible')
self.logger.warning(msg, conn, conn.max_ready_count,
self.max_in_flight)
conn.subscribe(self.topic, self.channel)
except NSQException as error:
self.logger.warning('[%s] connection failed (%r)', conn, error)
self.handle_connection_failure(conn)
return
# Check if we've closed since we started
if not self.is_running:
self.handle_connection_failure(conn)
return
self.logger.info('[%s] connection successful', conn)
self.handle_connection_success(conn)
def _listen(self, conn):
try:
conn.listen()
except NSQException as error:
self.logger.warning('[%s] connection lost (%r)', conn, error)
self.handle_connection_failure(conn)
def handle_connection_success(self, conn):
self._connections[conn] = THROTTLED
self._workers.spawn(self._listen, conn)
self.redistribute_ready_state()
if str(conn) not in self.nsqd_tcp_addresses:
return
self._connection_backoffs[conn].success()
def handle_connection_failure(self, conn):
del self._connections[conn]
conn.close_stream()
if not self.is_running:
return
self.redistribute_ready_state()
if str(conn) not in self.nsqd_tcp_addresses:
return
seconds = self._connection_backoffs[conn].failure().get_interval()
self.logger.debug('[%s] retrying in %ss', conn, seconds)
gevent.spawn_later(seconds, self.connect_to_nsqd, conn.address,
conn.port)
def handle_auth(self, conn, response):
metadata = []
if response.get('identity'):
metadata.append("Identity: %r" % response['identity'])
if response.get('permission_count'):
metadata.append("Permissions: %d" % response['permission_count'])
if response.get('identity_url'):
metadata.append(response['identity_url'])
self.logger.info('[%s] AUTH accepted %s', conn, ' '.join(metadata))
self.on_auth.send(self, conn=conn, response=response)
def handle_response(self, conn, response):
self.logger.debug('[%s] response: %s', conn, response)
self.on_response.send(self, response=response)
def handle_error(self, conn, error):
self.logger.debug('[%s] error: %s', conn, error)
self.on_error.send(self, error=error)
def _handle_message(self, message):
if self.max_tries and message.attempts > self.max_tries:
self.logger.warning("giving up on message '%s' after max tries %d",
message.id, self.max_tries)
self.on_giving_up.send(self, message=message)
return message.finish()
self.on_message.send(self, message=message)
if not self.is_running:
return
if message.is_async():
return
if message.has_responded():
return
message.finish()
def handle_message(self, conn, message):
self.logger.debug('[%s] got message: %s', conn, message.id)
try:
return self._handle_message(message)
except NSQRequeueMessage as error:
if error.backoff is None:
backoff = self.backoff_on_requeue
else:
backoff = error.backoff
except Exception as error:
backoff = True
self.logger.exception(
'[%s] caught exception while handling message', conn)
self.on_exception.send(self, message=message, error=error)
if not self.is_running:
return
if message.has_responded():
return
try:
message.requeue(self.requeue_delay, backoff)
except NSQException as error:
self.logger.warning('[%s] error requeueing message (%r)', conn,
error)
def _create_backoff(self):
return BackoffTimer(max_interval=self.max_backoff_duration)
def _start_backoff(self, conn):
self._connections[conn] = BACKOFF
interval = self._message_backoffs[conn].get_interval()
gevent.spawn_later(interval, self._start_throttled, conn)
self.logger.info('[%s] backing off for %s seconds', conn, interval)
self.redistribute_ready_state()
def _start_throttled(self, conn):
if self._connections.get(conn) != BACKOFF:
return
self._connections[conn] = THROTTLED
self.logger.info('[%s] testing backoff state with RDY 1', conn)
self.redistribute_ready_state()
def _complete_backoff(self, conn):
if self._message_backoffs[conn].is_reset():
self._connections[conn] = RUNNING
self.logger.info('throttle complete, resuming normal operation')
self.redistribute_ready_state()
else:
self._start_backoff(conn)
def _finish_message(self, conn, backoff):
if not self.max_backoff_duration:
return
try:
state = self._connections[conn]
except KeyError:
return
if state == BACKOFF:
return
if backoff:
self._message_backoffs[conn].failure()
self._start_backoff(conn)
elif state == THROTTLED:
self._message_backoffs[conn].success()
self._complete_backoff(conn)
def handle_finish(self, conn, message_id):
self.logger.debug('[%s] finished message: %s', conn, message_id)
self._finish_message(conn, backoff=False)
self.on_finish.send(self, message_id=message_id)
def handle_requeue(self, conn, message_id, timeout, backoff):
self.logger.debug('[%s] requeued message: %s (%s)', conn, message_id,
timeout)
self._finish_message(conn, backoff=backoff)
self.on_requeue.send(self, message_id=message_id, timeout=timeout)
class MasterLocustRunner(DistributedLocustRunner):
def __init__(self, *args, **kwargs):
super(MasterLocustRunner, self).__init__(*args, **kwargs)
class SlaveNodesDict(dict):
def get_by_state(self, state):
return [c for c in six.itervalues(self) if c.state == state]
@property
def ready(self):
return self.get_by_state(STATE_INIT)
@property
def hatching(self):
return self.get_by_state(STATE_HATCHING)
@property
def running(self):
return self.get_by_state(STATE_RUNNING)
self.clients = SlaveNodesDict()
self.server = rpc.Server(self.master_bind_host, self.master_bind_port)
self.greenlet = Group()
self.greenlet.spawn(
self.client_listener).link_exception(callback=self.noop)
# listener that gathers info on how many locust users the slaves has spawned
def on_slave_report(client_id, data):
if client_id not in self.clients:
logger.info("Discarded report from unrecognized slave %s",
client_id)
return
self.clients[client_id].user_count = data["user_count"]
events.slave_report += on_slave_report
# register listener that sends quit message to slave nodes
def on_quitting():
self.quit()
events.quitting += on_quitting
@property
def user_count(self):
return sum([c.user_count for c in six.itervalues(self.clients)])
def start_hatching(self, locust_count, hatch_rate):
num_slaves = len(self.clients.ready) + len(self.clients.running)
if not num_slaves:
logger.warning(
"You are running in distributed mode but have no slave servers connected. "
"Please connect slaves prior to swarming.")
return
self.num_clients = locust_count
slave_num_clients = locust_count // (num_slaves or 1)
slave_hatch_rate = float(hatch_rate) / (num_slaves or 1)
remaining = locust_count % num_slaves
logger.info("Sending hatch jobs to %d ready clients", num_slaves)
if self.state != STATE_RUNNING and self.state != STATE_HATCHING:
self.stats.clear_all()
self.exceptions = {}
events.master_start_hatching.fire()
for client in six.itervalues(self.clients):
data = {
"hatch_rate": slave_hatch_rate,
"num_clients": slave_num_clients,
"num_requests": self.num_requests,
"host": self.host,
"stop_timeout": None
}
if remaining > 0:
data["num_clients"] += 1
remaining -= 1
self.server.send(Message("hatch", data, None))
self.stats.start_time = time()
self.state = STATE_HATCHING
def stop(self):
for client in self.clients.hatching + self.clients.running:
self.server.send(Message("stop", None, None))
events.master_stop_hatching.fire()
def quit(self):
for client in six.itervalues(self.clients):
self.server.send(Message("quit", None, None))
self.greenlet.kill(block=True)
def client_listener(self):
while True:
msg = self.server.recv()
if msg.type == "client_ready":
id = msg.node_id
self.clients[id] = SlaveNode(id)
logger.info(
"Client %r reported as ready. Currently %i clients ready to swarm."
% (id, len(self.clients.ready)))
## emit a warning if the slave's clock seem to be out of sync with our clock
#if abs(time() - msg.data["time"]) > 5.0:
# warnings.warn("The slave node's clock seem to be out of sync. For the statistics to be correct the different locust servers need to have synchronized clocks.")
elif msg.type == "client_stopped":
del self.clients[msg.node_id]
if len(self.clients.hatching + self.clients.running) == 0:
self.state = STATE_STOPPED
logger.info("Removing %s client from running clients" %
(msg.node_id))
elif msg.type == "stats":
events.slave_report.fire(client_id=msg.node_id, data=msg.data)
elif msg.type == "hatching":
self.clients[msg.node_id].state = STATE_HATCHING
elif msg.type == "hatch_complete":
self.clients[msg.node_id].state = STATE_RUNNING
self.clients[msg.node_id].user_count = msg.data["count"]
if len(self.clients.hatching) == 0:
count = sum(c.user_count
for c in six.itervalues(self.clients))
events.hatch_complete.fire(user_count=count)
elif msg.type == "quit":
if msg.node_id in self.clients:
del self.clients[msg.node_id]
logger.info(
"Client %r quit. Currently %i clients connected." %
(msg.node_id, len(self.clients.ready)))
elif msg.type == "exception":
self.log_exception(msg.node_id, msg.data["msg"],
msg.data["traceback"])
@property
def slave_count(self):
return len(self.clients.ready) + len(self.clients.hatching) + len(
self.clients.running)
def start(self,greenlet=None): """Start the greenlet pool or add a greenlet to the pool.""" if greenlet is not None: return Group.start(self,greenlet)
def __init__(self,*args,**kwargs): Group.__init__(self,*args,**kwargs) self.open = True
#!/usr/bin/env python
#!encoding=utf8
import gevent
from gevent.pool import Group
def intensive(n):
gevent.sleep(3 - n)
return 'task', n
print('Ordered')
ogroup = Group()
x = ogroup.imap(intensive, xrange(3))
print x
for x in ogroup.imap(intensive, xrange(3)):
print x
import gevent
from gevent.pool import Group
def intensive(n):
gevent.sleep(3 - n)
return 'task', n
igroup = Group()
for i in igroup.imap_unordered(intensive, xrange(3)):
print(i)
class NoWebMasterLocustRunner(MasterLocustRunner):
def __init__(self, *args, **kwargs):
super(NoWebMasterLocustRunner, self).__init__(*args, **kwargs)
self.all_stats = {}
self.greenlet = Group()
self.max_num_requests = None
self.max_seconds_elapsed = None
def run_locustfiles(self):
"""
Iterate through all locustfiles and run each locustfile until the maximum number of requests are hit or
the maximum timeout in seconds has elapsed. When one of those conditions are met on each locustfile,
move on to the next one, storing the stats result of each locustfile run.
"""
for run_count, locustfile_key in enumerate(
self.available_locustfiles.keys()):
self.switch(locustfile_key)
if run_count > 0:
logger.info("Waiting {} seconds to cool down.".format(
self.cooldown))
gevent.sleep(self.cooldown)
super(NoWebMasterLocustRunner,
self).start_hatching(self.num_clients, self.hatch_rate)
while True:
hit_max_requests = self.max_num_requests and (
self.stats.aggregated_stats().num_requests >=
self.max_num_requests)
hit_max_elapsed_time = self.max_seconds_elapsed and (
time() - self.stats.start_time >= self.max_seconds_elapsed)
if hit_max_requests or hit_max_elapsed_time:
break
gevent.sleep(1)
self.all_stats[locustfile_key] = copy.deepcopy(self.stats)
if self.save_stats:
logfile = format_logfile(
self.statsfile_format, {
"locustfile": locustfile_key,
"date": datetime.datetime.now().isoformat()
})
save_logfile(
logfile,
json.dumps(self.stats.aggregated_stats().serialize(),
indent=4))
self.stop()
def wait_for_slaves(self, min_slaves, timeout):
"""
Wait the specified timeout seconds until the minimum number of slaves come online
:param min_slaves: Minimum number of slaves to expect
:param timeout: Max number of seconds to wait before raising an exception
:raises: polling.TimeoutException
:return: The count of slaves currently available
"""
return polling.poll(lambda: len(self.clients.ready),
check_success=lambda ready: ready >= min_slaves,
timeout=timeout,
step=1)
def slaves_start_swarming(self,
max_num_requests=None,
max_seconds_elapsed=None):
"""
Instruct the slaves to start swarming for the available locustfiles asynchronously
:param max_num_requests: Stop when the total number of requests for a locustfile reach this number
:param max_seconds_elapsed: Stop when the total elapsed seconds for a locustfile reach this number
"""
if max_num_requests is not None:
self.max_num_requests = max_num_requests
elif max_seconds_elapsed is not None:
self.max_seconds_elapsed = max_seconds_elapsed
else:
raise ValueError(
'You must specify the total number or requests to be made or a timeout'
)
self.state = STATE_HATCHING
self.greenlet.spawn(
self.run_locustfiles).link_exception(callback=self.noop)
class KafkaProductor(object):
def __init__(self, queues):
self.queues = queues
self.kafka = None
self.topic = None
self.sessions = None
self.group = None
self.rdkafka = None
self._init_parse()
def is_rdkafka(self):
try:
from pykafka import rdkafka
self.rdkafka = True
except:
self.rdkafka = False
def _init_parse(self):
self.sessions = [
self.is_connection(data['kafka'], data['topic'], data['queue'])
for data in self.queues
]
self.is_rdkafka()
def encode_topic(self, topic):
if isinstance(topic, bytes) is not True:
topic = str.encode(str(topic))
return topic
def switch(self):
sleep(0)
def is_connection(self, host, topic, queue):
try:
client = KafkaClient(hosts=host)
session = client.topics[self.encode_topic(topic)]
except:
raise KafkaConnectErr(host)
return (session, queue)
def start_a_producer(self, session):
_session, _queue = session[0], session[1]
with _session.get_producer(delivery_reports=True,
use_rdkafka=self.rdkafka) as producer:
count = 0
while True:
try:
msg = _queue.get_nowait()
except:
self.switch()
continue
count += 1
# producer.produce will only access bytes, translate the str to bytes first
producer.produce(msg.encode())
if count % 1000 == 0:
while True:
try:
msg, exc = producer.get_delivery_report(
block=False)
self.switch()
if exc is not None:
print('Failed to deliver msg {}: {}'.format(
msg.partition_key, repr(exc)))
except:
break
@property
def start(self):
self.group = Group()
for session in self.sessions:
self.group.add(spawn(self.start_a_producer, session))
self.group.join()
class Runner(object):
"""
Orchestrates the load test by starting and stopping the users.
Use one of the :meth:`create_local_runner <locust.env.Environment.create_local_runner>`,
:meth:`create_master_runner <locust.env.Environment.create_master_runner>` or
:meth:`create_worker_runner <locust.env.Environment.create_worker_runner>` methods on
the :class:`Environment <locust.env.Environment>` instance to create a runner of the
desired type.
"""
def __init__(self, environment):
self.environment = environment
self.user_greenlets = Group()
self.greenlet = Group()
self.state = STATE_INIT
self.hatching_greenlet = None
self.stepload_greenlet = None
self.current_cpu_usage = 0
self.cpu_warning_emitted = False
self.greenlet.spawn(self.monitor_cpu).link_exception(greenlet_exception_handler)
self.exceptions = {}
self.target_user_count = None
# set up event listeners for recording requests
def on_request_success(request_type, name, response_time, response_length, **kwargs):
self.stats.log_request(request_type, name, response_time, response_length)
def on_request_failure(request_type, name, response_time, response_length, exception, **kwargs):
self.stats.log_request(request_type, name, response_time, response_length)
self.stats.log_error(request_type, name, exception)
self.environment.events.request_success.add_listener(on_request_success)
self.environment.events.request_failure.add_listener(on_request_failure)
self.connection_broken = False
# register listener that resets stats when hatching is complete
def on_hatch_complete(user_count):
self.state = STATE_RUNNING
if environment.reset_stats:
logger.info("Resetting stats\n")
self.stats.reset_all()
self.environment.events.hatch_complete.add_listener(on_hatch_complete)
def __del__(self):
# don't leave any stray greenlets if runner is removed
if self.greenlet and len(self.greenlet) > 0:
self.greenlet.kill(block=False)
@property
def user_classes(self):
return self.environment.user_classes
@property
def stats(self) -> RequestStats:
return self.environment.stats
@property
def errors(self):
return self.stats.errors
@property
def user_count(self):
"""
:returns: Number of currently running users
"""
return len(self.user_greenlets)
def cpu_log_warning(self):
"""Called at the end of the test to repeat the warning & return the status"""
if self.cpu_warning_emitted:
logger.warning("CPU usage was too high at some point during the test! See https://docs.locust.io/en/stable/running-locust-distributed.html for how to distribute the load over multiple CPU cores or machines")
return True
return False
def weight_users(self, amount):
"""
Distributes the amount of users for each WebLocust-class according to it's weight
returns a list "bucket" with the weighted users
"""
bucket = []
weight_sum = sum([user.weight for user in self.user_classes])
residuals = {}
for user in self.user_classes:
if self.environment.host is not None:
user.host = self.environment.host
# create users depending on weight
percent = user.weight / float(weight_sum)
num_users = int(round(amount * percent))
bucket.extend([user for x in range(num_users)])
# used to keep track of the amount of rounding was done if we need
# to add/remove some instances from bucket
residuals[user] = amount * percent - round(amount * percent)
if len(bucket) < amount:
# We got too few User classes in the bucket, so we need to create a few extra users,
# and we do this by iterating over each of the User classes - starting with the one
# where the residual from the rounding was the largest - and creating one of each until
# we get the correct amount
for user in [l for l, r in sorted(residuals.items(), key=lambda x:x[1], reverse=True)][:amount-len(bucket)]:
bucket.append(user)
elif len(bucket) > amount:
# We've got too many users due to rounding errors so we need to remove some
for user in [l for l, r in sorted(residuals.items(), key=lambda x:x[1])][:len(bucket)-amount]:
bucket.remove(user)
return bucket
def spawn_users(self, spawn_count, hatch_rate, wait=False):
bucket = self.weight_users(spawn_count)
spawn_count = len(bucket)
if self.state == STATE_INIT or self.state == STATE_STOPPED:
self.state = STATE_HATCHING
existing_count = len(self.user_greenlets)
logger.info("Hatching and swarming %i users at the rate %g users/s (%i users already running)..." % (spawn_count, hatch_rate, existing_count))
occurrence_count = dict([(l.__name__, 0) for l in self.user_classes])
def hatch():
sleep_time = 1.0 / hatch_rate
while True:
if not bucket:
logger.info("All users hatched: %s (%i already running)" % (
", ".join(["%s: %d" % (name, count) for name, count in occurrence_count.items()]),
existing_count,
))
self.environment.events.hatch_complete.fire(user_count=len(self.user_greenlets))
return
user_class = bucket.pop(random.randint(0, len(bucket)-1))
occurrence_count[user_class.__name__] += 1
new_user = user_class(self.environment)
new_user.start(self.user_greenlets)
if len(self.user_greenlets) % 10 == 0:
logger.debug("%i users hatched" % len(self.user_greenlets))
if bucket:
gevent.sleep(sleep_time)
hatch()
if wait:
self.user_greenlets.join()
logger.info("All users stopped\n")
def stop_users(self, user_count):
"""
Stop a stop_count of weighted users from the Group() object in self.users
"""
bucket = self.weight_users(user_count)
user_count = len(bucket)
logger.info("Stopping %i users" % user_count)
to_stop = []
for g in self.user_greenlets:
for l in bucket:
user = g.args[0]
if l == type(user):
to_stop.append(user)
bucket.remove(l)
break
self.stop_user_instances(to_stop)
self.environment.events.hatch_complete.fire(user_count=self.user_count)
def stop_user_instances(self, users):
if self.environment.stop_timeout:
stopping = Group()
for user in users:
if not user.stop(self.user_greenlets, force=False):
# User.stop() returns False if the greenlet was not stopped, so we'll need
# to add it's greenlet to our stopping Group so we can wait for it to finish it's task
stopping.add(user._greenlet)
if not stopping.join(timeout=self.environment.stop_timeout):
logger.info("Not all users finished their tasks & terminated in %s seconds. Stopping them..." % self.environment.stop_timeout)
stopping.kill(block=True)
else:
for user in users:
user.stop(self.user_greenlets, force=True)
def monitor_cpu(self):
process = psutil.Process()
while True:
self.current_cpu_usage = process.cpu_percent()
if self.current_cpu_usage > 90 and not self.cpu_warning_emitted:
logging.warning("CPU usage above 90%! This may constrain your throughput and may even give inconsistent response time measurements! See https://docs.locust.io/en/stable/running-locust-distributed.html for how to distribute the load over multiple CPU cores or machines")
self.cpu_warning_emitted = True
gevent.sleep(CPU_MONITOR_INTERVAL)
def start(self, user_count, hatch_rate, wait=False):
"""
Start running a load test
:param user_count: Number of users to start
:param hatch_rate: Number of users to spawn per second
:param wait: If True calls to this method will block until all users are spawned.
If False (the default), a greenlet that spawns the users will be
started and the call to this method will return immediately.
"""
if self.state != STATE_RUNNING and self.state != STATE_HATCHING:
self.stats.clear_all()
self.exceptions = {}
self.cpu_warning_emitted = False
self.worker_cpu_warning_emitted = False
self.target_user_count = user_count
# Dynamically changing the user count
if self.state != STATE_INIT and self.state != STATE_STOPPED:
self.state = STATE_HATCHING
if self.user_count > user_count:
# Stop some users
stop_count = self.user_count - user_count
self.stop_users(stop_count)
elif self.user_count < user_count:
# Spawn some users
spawn_count = user_count - self.user_count
self.spawn_users(spawn_count=spawn_count, hatch_rate=hatch_rate)
else:
self.environment.events.hatch_complete.fire(user_count=self.user_count)
else:
self.hatch_rate = hatch_rate
self.spawn_users(user_count, hatch_rate=hatch_rate, wait=wait)
def start_stepload(self, user_count, hatch_rate, step_user_count, step_duration):
if user_count < step_user_count:
logger.error("Invalid parameters: total user count of %d is smaller than step user count of %d" % (user_count, step_user_count))
return
self.total_users = user_count
if self.stepload_greenlet:
logger.info("There is an ongoing swarming in Step Load mode, will stop it now.")
self.stepload_greenlet.kill()
logger.info("Start a new swarming in Step Load mode: total user count of %d, hatch rate of %d, step user count of %d, step duration of %d " % (user_count, hatch_rate, step_user_count, step_duration))
self.state = STATE_INIT
self.stepload_greenlet = self.greenlet.spawn(self.stepload_worker, hatch_rate, step_user_count, step_duration)
self.stepload_greenlet.link_exception(greenlet_exception_handler)
def stepload_worker(self, hatch_rate, step_users_growth, step_duration):
current_num_users = 0
while self.state == STATE_INIT or self.state == STATE_HATCHING or self.state == STATE_RUNNING:
current_num_users += step_users_growth
if current_num_users > int(self.total_users):
logger.info('Step Load is finished.')
break
self.start(current_num_users, hatch_rate)
logger.info('Step loading: start hatch job of %d user.' % (current_num_users))
gevent.sleep(step_duration)
def stop(self):
"""
Stop a running load test by stopping all running users
"""
self.state = STATE_CLEANUP
# if we are currently hatching users we need to kill the hatching greenlet first
if self.hatching_greenlet and not self.hatching_greenlet.ready():
self.hatching_greenlet.kill(block=True)
self.stop_user_instances([g.args[0] for g in self.user_greenlets])
self.state = STATE_STOPPED
self.cpu_log_warning()
def quit(self):
"""
Stop any running load test and kill all greenlets for the runner
"""
self.stop()
self.greenlet.kill(block=True)
def log_exception(self, node_id, msg, formatted_tb):
key = hash(formatted_tb)
row = self.exceptions.setdefault(key, {"count": 0, "msg": msg, "traceback": formatted_tb, "nodes": set()})
row["count"] += 1
row["nodes"].add(node_id)
self.exceptions[key] = row
class GeventThreadController(Thread):
def __init__(self, name, sleep=0.2, debug=False):
Thread.__init__(self)
self.name = name
self._q_wait_to_run = queue.Queue()
self._q_wait_to_stop = queue.Queue()
self._is_stop = False
self._greenlets = Group()
self._sleep = sleep
self._isdebug = debug
def run(self):
self._is_stop = False
self._debug("{} start".format(self.name))
while not self._is_stop:
if not self._q_wait_to_stop.empty():
self._stop_greenlet(self._q_wait_to_stop.get())
elif not self._q_wait_to_run.empty():
self._start_greenlet(self._q_wait_to_run.get())
else:
gevent.sleep(self._sleep)
# Stoping thread
self._greenlets.kill()
self._debug("{} stop".format(self.name))
def stop(self):
self.stop_nowait()
self.join()
def stop_nowait(self):
self._is_stop = True
def append(self, name, green_create_task, *args, **kwargs):
if hasattr(green_create_task, '__call__'):
self._q_wait_to_run.put(tuple((name, green_create_task, args, kwargs)))
else:
raise TypeError("must be \'function\' or \'GreenClassPack\', not {}".format(type(green_create_task)))
def stop_greenlet(self, gr_mame):
self._q_wait_to_stop.put(gr_mame)
def __getitem__(self, gr_mame):
gr_obj = self._get_greenlet(gr_mame)
if gr_obj:
return gr_obj
raise IndexError
def _start_greenlet(self, gr_create_task_tuple):
gr_task_name = gr_create_task_tuple[0]
gr_create_task = gr_create_task_tuple[1]
args = gr_create_task_tuple[2]
kwargs = gr_create_task_tuple[3]
gr_obj = gr_create_task(*args, **kwargs)
if isinstance(gr_obj, Greenlet) and not self._get_greenlet(gr_task_name):
gevent.spawn(self._task_hndler, gr_obj, gr_task_name)
else:
self._debug("Error form {0}.{1}: must be \'gevent.Greenlet\', not {2}. Or same greenlet already runnig"
.format(self.name, gr_task_name, type(gr_obj)))
def _stop_greenlet(self, gr_task_name):
gr_obj = self._get_greenlet(gr_task_name)
if gr_obj:
gr_obj.kill()
def _get_greenlet(self, name_str):
gr_name = "{0}.{1}".format(self.name, name_str)
for gr_obj in self._greenlets:
if gr_name == gr_obj.name:
return gr_obj
return None
def _task_hndler(self, gr_obj, gr_name):
gr_obj.name = "{0}.{1}".format(self.name, gr_name)
self._greenlets.add(gr_obj)
self._debug("{0}.{1}: starting".format(self.name, gr_name))
gr_obj.start()
gr_obj.join()
self._greenlets.discard(gr_obj)
self._debug("{0}.{1} is stop".format(self.name, gr_name))
def _debug(self, log_str):
if self._isdebug:
print(log_str)
def task_redis_push_multi_imap(self,
records: List[Any],
push_mode: str = "pipe") -> bool:
"""Push records to Redis in parallel.
Split ``records`` list into smaller data sets to be pushed to Redis
with multiple Greenlets. Parameter ``push_mode`` defines which data
transfer strategy to use.
Note
----
See the note for :meth:`task_redis_push_multi_spawn` for detailed explanation on
parameter ``push_mode``.
Parameters
----------
records
List of records to store in Redis.
push_mode
Specifies the data transfer strategy to use.
Returns
-------
bool
Flag indicating success or failure of the operation
"""
# noqa: D202
def chunks(l: List, n: int):
"""Yield successive n-sized chunks from l."""
for i in range(0, len(l), n):
yield l[i:i + n] # noqa: E203
result: bool = False
batch_size: int = 0
records_len: int = len(records)
greenlets_count: int = self._config["graph_greenlets_count"]
if records_len > greenlets_count:
batch_size = (((records_len << 1) // greenlets_count) + 1) >> 1
elif records_len > 0:
batch_size = records_len
if batch_size:
greenlets: Group = Group()
redis_push_func_partial: Callable = partial(
getattr(self, "_records_to_redis_{}".format(push_mode)))
greenlet_results: List = []
for g in greenlets.imap_unordered(redis_push_func_partial,
chunks(records, batch_size)):
greenlet_results.append(g)
self._logger.debug("%s Greenlets returned: %s",
len(greenlet_results), greenlet_results)
result = False if False in greenlet_results else True
self._logger.debug("Summary result: %s", result)
return result
def spawn_imap_unordered(self):
igroup = Group()
result = []
with tracer.start_active_span('test'):
for i in igroup.imap_unordered(self.make_http_call, range(3)):
result.append(i)
def stop_users(self, user_count, stop_rate=None):
"""
Stop `user_count` weighted users at a rate of `stop_rate`
"""
if user_count == 0 or stop_rate == 0:
return
bucket = self.weight_users(user_count)
user_count = len(bucket)
to_stop = []
for g in self.user_greenlets:
for l in bucket:
user = g.args[0]
if isinstance(user, l):
to_stop.append(user)
bucket.remove(l)
break
if not to_stop:
return
if stop_rate is None or stop_rate >= user_count:
sleep_time = 0
logger.info("Stopping %i users" % (user_count))
else:
sleep_time = 1.0 / stop_rate
logger.info("Stopping %i users at rate of %g users/s" %
(user_count, stop_rate))
async_calls_to_stop = Group()
stop_group = Group()
while True:
user_to_stop: User = to_stop.pop(
random.randint(0,
len(to_stop) - 1))
logger.debug("Stopping %s" % user_to_stop._greenlet.name)
if user_to_stop._greenlet is greenlet.getcurrent():
# User called runner.quit(), so dont block waiting for killing to finish"
user_to_stop._group.killone(user_to_stop._greenlet,
block=False)
elif self.environment.stop_timeout:
async_calls_to_stop.add(
gevent.spawn_later(0, User.stop, user_to_stop,
force=False))
stop_group.add(user_to_stop._greenlet)
else:
async_calls_to_stop.add(
gevent.spawn_later(0, User.stop, user_to_stop, force=True))
if to_stop:
gevent.sleep(sleep_time)
else:
break
async_calls_to_stop.join()
if not stop_group.join(timeout=self.environment.stop_timeout):
logger.info(
"Not all users finished their tasks & terminated in %s seconds. Stopping them..."
% self.environment.stop_timeout)
stop_group.kill(block=True)
logger.info("%i Users have been stopped" % user_count)
class FilterManager:
"""Manages access to filtered Ethereum events."""
def __init__(self):
self.wrappers = []
self.pool = Group()
def register(self,
filter_installer: FilterInstaller,
fmt_cls: FormatClass,
backoff: bool = True):
"""Add a new filter, with an optional associated WebsocketMessage-serializer class"""
wrapper = FilterWrapper(filter_installer, fmt_cls, backoff)
self.wrappers.append(wrapper)
logger.debug('Registered new filter: %s', wrapper)
def flush(self):
"""End all event polling, uninstall all filters and remove their corresponding wrappers"""
self.pool.kill()
self.wrappers.clear()
def fetch(self):
"""Return a queue of currently managed contract events"""
queue = Queue()
for wrapper in self.wrappers:
self.pool.spawn(wrapper.spawn_poll_loop, queue.put_nowait)
yield from queue
def setup_event_filters(self, chain: Any):
"""Setup the most common event filters"""
if len(self.wrappers) != 0:
logger.exception(
"Attempting to initialize already initialized filter manager")
self.flush()
bounty_contract = chain.bounty_registry.contract
# Setup Latest (although this could pass `w3.eth.filter` directly)
self.register(chain.w3.eth.filter,
messages.LatestEvent.make(chain.w3.eth),
backoff=False)
# messages.NewBounty shouldn't wait or back-off from new bounties.
self.register(bounty_contract.eventFilter,
messages.NewBounty,
backoff=False)
filter_events: List[FormatClass] = [
messages.FeesUpdated,
messages.WindowsUpdated,
messages.NewAssertion,
messages.NewVote,
messages.QuorumReached,
messages.SettledBounty,
messages.RevealedAssertion,
messages.Deprecated,
messages.Undeprecated,
]
for cls in filter_events:
self.register(bounty_contract.eventFilter, cls)
offer_registry = chain.offer_registry
if offer_registry and offer_registry.contract:
self.register(offer_registry.contract.eventFilter,
messages.InitializedChannel)
class LocustRunner(object):
def __init__(self, locust_classes, options):
self.options = options
self.locust_classes = locust_classes
self.hatch_rate = options.hatch_rate
self.num_clients = options.num_clients
self.num_requests = options.num_requests
self.host = options.host
self.locusts = Group()
self.state = STATE_INIT
self.hatching_greenlet = None
self.exceptions = {}
self.stats = global_stats
# register listener that resets stats when hatching is complete
def on_hatch_complete(user_count):
self.state = STATE_RUNNING
if not self.options.no_reset_stats:
logger.info("Resetting stats\n")
self.stats.reset_all()
events.hatch_complete += on_hatch_complete
@property
def request_stats(self):
return self.stats.entries
@property
def errors(self):
return self.stats.errors
@property
def user_count(self):
return len(self.locusts)
def weight_locusts(self, amount, stop_timeout=None):
"""
Distributes the amount of locusts for each WebLocust-class according to it's weight
returns a list "bucket" with the weighted locusts
"""
bucket = []
weight_sum = sum((locust.weight for locust in self.locust_classes
if locust.task_set))
for locust in self.locust_classes:
if not locust.task_set:
warnings.warn(
"Notice: Found Locust class (%s) got no task_set. Skipping..."
% locust.__name__)
continue
if self.host is not None:
locust.host = self.host
if stop_timeout is not None:
locust.stop_timeout = stop_timeout
# create locusts depending on weight
percent = locust.weight / float(weight_sum)
num_locusts = int(round(amount * percent))
bucket.extend([locust for x in xrange(0, num_locusts)])
return bucket
def spawn_locusts(self, spawn_count=None, stop_timeout=None, wait=False):
if spawn_count is None:
spawn_count = self.num_clients
if self.num_requests is not None:
self.stats.max_requests = self.num_requests
bucket = self.weight_locusts(spawn_count, stop_timeout)
spawn_count = len(bucket)
if self.state == STATE_INIT or self.state == STATE_STOPPED:
self.state = STATE_HATCHING
self.num_clients = spawn_count
else:
self.num_clients += spawn_count
logger.info(
"Hatching and swarming %i clients at the rate %g clients/s..." %
(spawn_count, self.hatch_rate))
occurence_count = dict([(l.__name__, 0) for l in self.locust_classes])
def hatch():
sleep_time = 1.0 / self.hatch_rate
while True:
if not bucket:
logger.info("All locusts hatched: %s" % ", ".join([
"%s: %d" % (name, count)
for name, count in six.iteritems(occurence_count)
]))
events.hatch_complete.fire(user_count=self.num_clients)
return
locust = bucket.pop(random.randint(0, len(bucket) - 1))
occurence_count[locust.__name__] += 1
def start_locust(_):
try:
locust().run()
except GreenletExit:
pass
new_locust = self.locusts.spawn(start_locust, locust)
if len(self.locusts) % 10 == 0:
logger.debug("%i locusts hatched" % len(self.locusts))
gevent.sleep(sleep_time)
hatch()
if wait:
self.locusts.join()
logger.info("All locusts dead\n")
def kill_locusts(self, kill_count):
"""
Kill a kill_count of weighted locusts from the Group() object in self.locusts
"""
bucket = self.weight_locusts(kill_count)
kill_count = len(bucket)
self.num_clients -= kill_count
logger.info("Killing %i locusts" % kill_count)
dying = []
for g in self.locusts:
for l in bucket:
if l == g.args[0]:
dying.append(g)
bucket.remove(l)
break
for g in dying:
self.locusts.killone(g)
events.hatch_complete.fire(user_count=self.num_clients)
def start_hatching(self, locust_count=None, hatch_rate=None, wait=False):
if self.state != STATE_RUNNING and self.state != STATE_HATCHING:
self.stats.clear_all()
self.stats.start_time = time()
self.exceptions = {}
events.locust_start_hatching.fire()
# Dynamically changing the locust count
if self.state != STATE_INIT and self.state != STATE_STOPPED:
self.state = STATE_HATCHING
if self.num_clients > locust_count:
# Kill some locusts
kill_count = self.num_clients - locust_count
self.kill_locusts(kill_count)
elif self.num_clients < locust_count:
# Spawn some locusts
if hatch_rate:
self.hatch_rate = hatch_rate
spawn_count = locust_count - self.num_clients
self.spawn_locusts(spawn_count=spawn_count)
else:
events.hatch_complete.fire(user_count=self.num_clients)
else:
if hatch_rate:
self.hatch_rate = hatch_rate
if locust_count is not None:
self.spawn_locusts(locust_count, wait=wait)
else:
self.spawn_locusts(wait=wait)
def stop(self):
# if we are currently hatching locusts we need to kill the hatching greenlet first
if self.hatching_greenlet and not self.hatching_greenlet.ready():
self.hatching_greenlet.kill(block=True)
self.locusts.kill(block=True)
self.state = STATE_STOPPED
events.locust_stop_hatching.fire()
def log_exception(self, node_id, msg, formatted_tb):
key = hash(formatted_tb)
row = self.exceptions.setdefault(key, {
"count": 0,
"msg": msg,
"traceback": formatted_tb,
"nodes": set()
})
row["count"] += 1
row["nodes"].add(node_id)
self.exceptions[key] = row
def task_redis_push_multi_spawn(self,
records: List[Any],
push_mode: str = "pipe") -> bool:
"""Push records to Redis in parallel.
Split ``records`` list into smaller data sets to be pushed to Redis
with multiple Greenlets.
Note
----
Parameter ``push_mode`` defines which data transfer strategy to use
by providing the suffix for specific ``_records_to_redis_*`` method
to be called. Currently accepts ``naive`` or ``pipe``.
Parameters
----------
records
List of records to store in Redis.
push_mode
Specifies the data transfer strategy to use.
Returns
-------
bool
Flag indicating success or failure of the operation
"""
# noqa: D202
def chunks(l: List, n: int):
"""Yield successive n-sized chunks from l."""
for i in range(0, len(l), n):
yield l[i:i + n] # noqa: E203
result: bool = False
batch_size: int = 0
records_len: int = len(records)
greenlets_count: int = self._config["graph_greenlets_count"]
if records_len > greenlets_count:
batch_size = (((records_len << 1) // greenlets_count) + 1) >> 1
elif records_len > 0:
batch_size = records_len
if records_len:
greenlets: Group = Group()
# See https://docs.python.org/3/library/functools.html#functools.partial
redis_push_func_partial: Callable = partial(
getattr(self, "_records_to_redis_{}".format(push_mode)))
greenlet_list: List = []
for batch in chunks(records, batch_size):
greenlet_list.append(
greenlets.spawn(redis_push_func_partial, records=batch))
greenlets.join(raise_error=True)
greenlet_results = [g.value for g in greenlet_list]
self._logger.debug("%s Greenlets spawned: %s",
len(greenlet_results), greenlet_results)
result = False if False in greenlet_results else True
self._logger.debug("Summary result: %s", result)
return result
def printer(msg, func, *args):
"run a function, print results"
print msg, '<request>'
res = func(*args)
print msg, '<response>', res
if __name__ == '__main__':
#from netcall import setup_logger
#setup_logger()
# Custom serializer/deserializer functions can be passed in. The server
# side ones must match.
echo = GreenRPCClient(green_env='gevent', serializer=JSONSerializer())
echo.connect('tcp://127.0.0.1:5555')
tasks = Group()
spawn = tasks.spawn
spawn(printer, "[echo] Echoing \"Hi there\"", echo.echo, "Hi there")
try:
print "Testing a remote exception...",
echo.error()
print "FAIL, no remote exception!"
except RemoteRPCError, e:
print "OK, got an expected remote exception:"
#print e.ename
print e.evalue
print e.traceback
try:
def __init__(self):
self.mods = {}
self.group = Group()
self.cmd_tables = {}
self._tasks = []
def create_connection(self, address, timeout=10):
ip_list = dnslib.dnsQuery(address[0])
# 尝试连接缓存
route_list = self.get_route_order_ping(address[0], address[1], None)
if route_list:
try:
route = route_list[0]
hit_ip = route['hit_ip']
if hit_ip in ip_list:
cache_timeout = route['tcp_ping']
if cache_timeout < 1000:
cache_timeout = cache_timeout * 2
else:
cache_timeout = cache_timeout + 1000
cache_timeout = int(math.ceil(cache_timeout / 1000.0))
start_time = int(time.time() * 1000)
sock = self._direct_create_connection(
address, hit_ip, cache_timeout)
t = int(time.time() * 1000) - start_time
logging.debug(
u'[upstream][RouteCache]%s 缓存记录 连接 %s(%s):%s 命中。time:%s'
% (self.get_display_name(), address[0], hit_ip,
address[1], t))
self.update_route_ping(address[0], address[1], t, hit_ip)
return sock
else:
logging.debug(
u'[upstream][RouteCache]%s 缓存记录 连接 %s(%s):%s IP 不匹配，放弃缓存。'
% (self.get_display_name(), address[0], hit_ip,
address[1]))
except:
t = int(time.time() * 1000) - start_time
info = traceback.format_exc()
logging.debug(
u'[upstream][RouteCache]%s 缓存记录 连接 %s(%s):%s 失败。time:%s' %
(self.get_display_name(), address[0], hit_ip, address[1],
t))
logging.debug('%s\r\n\r\n' % info)
# 缓存失败，连接全部
evt = Event()
group = Group()
aync_task = DirectAsyncTask(evt, None, group)
for ip in ip_list:
group.add(
gevent.spawn(self._create_connection, aync_task, address, ip,
timeout))
# 所有连接失败时发出通知
gevent.spawn(self._create_connection_all_end, aync_task)
evt.wait()
if aync_task.sock:
return aync_task.sock
else:
raise UpstreamConnectError()
class CommonInterface:
'''通用接口类，所有新增接口都补充到该类中'''
logger = Glogger.getLogger()
def __init__(self):
self.client = None
self.greenlets = Group()
self.loop_check_interval = 1 / 100
self.hb_interval = 30
self.hb_version = 0
self.app_id = "123"
self.cli_type = 123
#包头中的sequence字段，每发送后自增
self.seq = 0
#包头中的target字段
self.target = 0
def connect_to_svc(self, lbsip='172.16.16.25', port=6000):
'''1.连接lbs获取acc地址列表 2.连接acc获取并更新心跳策略 '''
self.etablish_con(ip=lbsip, port=port)
acc_list = self.get_acc_list()
acc_ip, acc_port = self.get_first_accaddr(acc_list)
self.etablish_con(ip=acc_ip, port=acc_port, acc=True)
self.start_heart_beat()
self.update_heartbeat()
self.report_hb_version()
self.client_access()
def connect_to_acc(self, ip, port):
'''连接指定的acc server'''
self.etablish_con(ip, port, acc=True)
self.start_heart_beat()
self.report_hb_version()
self.update_heartbeat()
self.client_access()
def get_first_accaddr(self, acc_list):
try:
acc_ip = unit32_to_ip(acc_list['acc_addrs'][0]['ip'])
acc_port = int(acc_list['acc_addrs'][0]['ports'][0])
except KeyError:
self.close()
raise KeyError(
"Received acc list is empty, unable to continue,exit!"
) from None
return acc_ip, acc_port
def start_heart_beat(self):
def _heart_beat():
self.logger.debug("Heart beat thread started!")
try:
while self.client:
self.send_data(CommandMapping.HeartBeat, {})
gevent.sleep(self.hb_interval)
self.logger.debug("Socket close, stop sending heart beat.")
except Exception as e:
self.logger.error("HeartBeat greenlet error:{}".format(e.args))
self.greenlets.spawn(_heart_beat)
def get_response(self, cmd, timeout=5):
'''获取指定cmd的响应结果'''
try:
timer = gevent.Timeout(timeout)
timer.start()
while True:
if cmd in self.client.callback_dict:
result = self.client.callback_dict.pop(cmd)
return result
gevent.sleep(0.01)
except gevent.Timeout as t:
if t is timer:
return {'timeout': timeout}
else:
raise
finally:
timer.cancel()
def close(self):
if self.client:
self.client.close()
self.client = None
self.greenlets.kill()
def update_data(self, d1, d2):
for key in d2.keys():
if key not in d1:
raise KeyError("Default param:{} has no key {}".format(
d1, key))
d1.update(d2)
return d1
def get_heartbeat_conf(self, **kwargs):
'''
根据客户端类型获取对应的心跳配置
@return: {"version":1,"send_sec":15,"check_sec":30,"bt_timeout":60}
'''
data = {"cli_type": 0}
self.update_data(data, kwargs)
self.send_data(CommandMapping.GetHeartBeatConf, data)
result = self.get_response(CommandMapping.GetHeartBeatConfRsp)
return result
def update_heartbeat(self):
'''根据获取的心跳配置更新本地的心跳发送策略'''
try:
conf = self.get_heartbeat_conf()['body']
self.hb_interval = int(conf["send_sec"])
self.hb_version = int(conf["version"])
except KeyError:
self.logger.error("Get HB conf timeout, use the default 30s")
def report_hb_version(self):
'''获取到心跳配置后上报心跳版本号'''
data = {"version": self.hb_version}
self.send_data(CommandMapping.ClientReportHbConf, data)
return self.get_response(CommandMapping.ClientReportHbConfRsp)
def client_access(self, **kwargs):
'''客户端接入，要确保所有传入的app_id和cli_type都一致'''
data = {
"app_id": self.app_id,
"cli_type": self.cli_type,
"cli_version": "16",
"cli_os_ver": 0,
"cli_mac": "test",
# "device_id": b'119'
}
self.update_data(data, kwargs)
self.send_data(CommandMapping.ClientAccess, data)
return self.get_response(CommandMapping.ClientAccessRsp)
def login(self, **kwargs):
'''账号登陆,该服务暂时未启用，后期可能会用到'''
data = {
"account_type": 2,
"account": "123123",
"auth_ticket": md5('hello'.encode()).hexdigest(),
"online_status": 1,
"cli_type": 0,
# "device_id": bin_to_b64str(b'123'),
"app_id": '0x0001'
}
self.update_data(data, kwargs)
self.send_data(CommandMapping.UserLogin, data)
return self.get_response(CommandMapping.UserLoginRsp)
def anonymous_login(self, **kwargs):
'''匿名登陆'''
data = {
"uid": 666,
"online_status": 1,
"cli_type": self.cli_type,
# "device_id": b'234',
"app_id": self.app_id
}
self.update_data(data, kwargs)
self.send_data(CommandMapping.AnonymousLogin, data)
return self.get_response(CommandMapping.AnonymousLoginRsp)
def get_anonymous_uid(self, **kwargs):
'''获取匿名登陆账号'''
self.send_data(CommandMapping.GetAnonymousUid, {})
return self.get_response(CommandMapping.GetAnonymousUidRsp)
def token_login(self, **kwargs):
'''token登陆'''
data = {
"uid": self.uid,
"token": self.token,
"online_status": 1,
"cli_type": self.cli_type,
# "device_id": bin_to_b64str(b'123'),
"app_id": self.app_id
}
self.update_data(data, kwargs)
self.send_data(CommandMapping.TokenLogin, data)
return self.get_response(CommandMapping.TokenLoginRsp)
def login_complete(self):
self.send_data(CommandMapping.LoginComplete, {})
def logout(self):
'''用户登出'''
self.send_data(CommandMapping.UserLogout, {})
response = self.get_response(CommandMapping.UserLogoutRsp)
if "timeout" in response:
return False
return True
def update_token(self, **kwargs):
'''通过refresh_token更新token'''
data = {"refresh_token": "123", "cli_type": 0, "device_id": "123"}
self.update_data(data, kwargs)
self.send_data(CommandMapping.UpdateToken, data)
return self.get_response(CommandMapping.UpdateTokenRsp)
def update_refresh_token(self, **kwargs):
'''通过refresh token更新refresh token'''
data = {
"refresh_token": "123",
"cli_type": self.cli_type,
# "device_id": '123',
"app_id": self.app_id
}
self.update_data(data, kwargs)
self.send_data(CommandMapping.UpdateRefreshToken, data)
return self.get_response(CommandMapping.UpdateRefreshTokenRsp)
def enter_class(self, cid, **kwargs):
'''进入教室'''
data = {"user_name": "user1", "cus_data": "test"}
self.target = cid
self.update_data(data, kwargs)
self.send_data(CommandMapping.EnterClass, data)
return self.get_response(CommandMapping.EnterClassRsp)
def enter_class_notify(self):
'''接收某个用户进入教室的通知'''
return self.get_response(CommandMapping.EnterClassNotify)
def enter_class_complete(self):
self.send_data(CommandMapping.EnterClassComplete, {})
def leave_class(self):
'''退出教室'''
self.send_data(CommandMapping.LeaveClass, {})
return self.get_response(CommandMapping.LeaveClassRsp)
def leave_class_notify(self):
'''退出教室通知'''
return self.get_response(CommandMapping.LeaveClassNotify)
def force_leave_class(self, **kwargs):
'''强制某个用户退出教室'''
data = {"uid": 123, "reason": 0}
self.update_data(data, kwargs)
self.send_data(CommandMapping.ForceLeaveClass, data)
return self.get_response(CommandMapping.ForceLeaveClassRsp)
def force_leave_class_notify(self):
'''强制退出教室通知'''
return self.get_response(CommandMapping.ForceLeaveClassNotify)
def shutdown_class(self, **kwargs):
'''关闭教室'''
data = {"reason": 0}
self.update_data(data, kwargs)
self.send_data(CommandMapping.CloseClass, data)
return self.get_response(CommandMapping.CloseClassRsp)
def shutdown_class_notify(self):
'''关闭教室通知'''
return self.get_response(CommandMapping.CloseClassNotify)
def send_chat_message(self, **kwargs):
'''用户在教室内发送聊天消息'''
data = {
"cli_seq": 0,
"send_tm": int(time.time()),
"type": 0,
"option": "font=宋体",
"chat_msg": "hello world"
}
self.update_data(data, kwargs)
self.send_data(CommandMapping.UserSendMessage, data)
return self.get_response(CommandMapping.UserSendMessageRsp)
def receive_chat_message(self):
'''用户在教室内接收消息'''
message = self.get_response(CommandMapping.UserReceiveMessage)
if "timeout" not in message:
data = {"msg_id": message["body"]["msg_id"]}
self.send_data(CommandMapping.UserReceiveMessageRsp, data)
return message
def receive_offline_chat_message(self):
'''用户接收离线消息'''
return self.get_response(CommandMapping.ReceiveOfflineMessage)
def hand_up(self):
'''举手'''
self.send_data(CommandMapping.HandUp, {})
return self.get_response(CommandMapping.HandUpRsp)
def hand_down(self):
'''放手'''
self.send_data(CommandMapping.HandDown, {})
return self.get_response(CommandMapping.HandDownRsp)
def change_speak_list(self, **kwargs):
'''修改发言列表'''
data = {"op_type": 1, "target_uid": 123}
self.update_data(data, kwargs)
self.send_data(CommandMapping.ChangeSpeakOrder, data)
return self.get_response(CommandMapping.ChangeSpeakOrderRsp)
def change_speak_list_notify(self):
'''修改发言列表通知'''
return self.get_response(CommandMapping.ChangeSpeakOrderNotify)
def http_request(self, url, method, data):
self.logger.info("Send request:URL-{} Method-{} Data-{}".format(
url, method, data))
response = requests.request(method, url, params=data, json=data)
status = response.status_code
if status == requests.codes.ok:
r = response.json()
self.logger.info("Reponse from {}: {}".format(url, r))
else:
r = str(response.content)
self.logger.error("Request fail:http_code:{},msg:{}".format(
status, r))
return r
class HttpLoginError(Exception):
pass
def http_login(self, **kwds):
'''http账号密码登陆接口'''
url = "http://api.mebutoo.com/namelogin"
data = {
"uname": "test1916",
"password": "******",
"online_status": 1,
"ctype": self.cli_type,
"app_id": self.app_id
}
self.update_data(data, kwds)
response = self.http_request(url, 'POST', data)
if response.get('rsp_code') != 0:
raise HttpLoginError(response)
self.token = response['token']
self.refresh_token = response['refresh_token']
self.uid = response['uid']
return response
def http_regist(self, **kwds):
'''http用户注册接口'''
url = "http://api.mebutoo.com/nameregist"
data = {
"uname": "test1916",
"password": "******",
# "device_id": "123",
"app_id": self.app_id,
"client_version": "1.0",
"client_type": self.cli_type
}
self.update_data(data, kwds)
response = self.http_request(url, 'POST', data)
return response
def http_anonymous_login(self, **kwds):
url = "http://api.mebutoo.com/anylogin"
data = {
"uid": self.http_get_anonymous_uid(),
"online_status": 1,
"ctype": self.cli_type,
# "deviceid": "123",
"app_id": self.app_id
}
self.update_data(data, kwds)
response = self.http_request(url, 'POST', data)
if response.get('rsp_code') != 0:
raise HttpLoginError(response)
self.token = response['token']
self.refresh_token = response['refresh_token']
self.uid = response['uid']
return response
def http_get_anonymous_uid(self, **kwds):
url = "http://api.mebutoo.com/anyregist"
data = {"app_id": self.app_id}
self.update_data(data, kwds)
r = self.http_request(url, 'GET', data)
uid = r['uid'] if r.get('rsp_code', 0) == 0 else -1
return uid
def http_get_state(self, **kwds):
url = "http://api.mebutoo.com/oauthstate"
data = {
"cli_type": self.cli_type,
# "device_id":"123",
"app_id": self.app_id,
"target": 123
}
self.update_data(data, kwds)
return self.http_request(url, 'GET', data)
def start(self):
self.group = Group()
for session in self.sessions:
self.group.add(spawn(self.start_a_producer, session))
self.group.join()
def __init__(self):
self.group = Group()
self.channel: GeventChannel[
Tuple[Callable, List[Any], Dict[str, Any]]
] = GeventChannel()
self.spawn_work_greenlet = gevent.spawn(self._spawn_work)
class SlaveLocustRunner(DistributedLocustRunner):
def __init__(self, *args, **kwargs):
super(SlaveLocustRunner, self).__init__(*args, **kwargs)
self.client_id = socket.gethostname() + "_" + md5(
str(time() +
random.randint(0, 10000)).encode('utf-8')).hexdigest()
self.client = rpc.Client(self.master_host, self.master_port)
self.greenlet = Group()
self.greenlet.spawn(self.worker).link_exception(callback=self.noop)
self.client.send(Message("client_ready", None, self.client_id))
self.greenlet.spawn(
self.stats_reporter).link_exception(callback=self.noop)
# register listener for when all locust users have hatched, and report it to the master node
def on_hatch_complete(user_count):
self.client.send(
Message("hatch_complete", {"count": user_count},
self.client_id))
events.hatch_complete += on_hatch_complete
# register listener that adds the current number of spawned locusts to the report that is sent to the master node
def on_report_to_master(client_id, data):
data["user_count"] = self.user_count
events.report_to_master += on_report_to_master
# register listener that sends quit message to master
def on_quitting():
self.client.send(Message("quit", None, self.client_id))
events.quitting += on_quitting
# register listener thats sends locust exceptions to master
def on_locust_error(locust_instance, exception, tb):
formatted_tb = "".join(traceback.format_tb(tb))
self.client.send(
Message("exception", {
"msg": str(exception),
"traceback": formatted_tb
}, self.client_id))
events.locust_error += on_locust_error
def worker(self):
while True:
msg = self.client.recv()
if msg.type == "hatch":
self.client.send(Message("hatching", None, self.client_id))
job = msg.data
self.hatch_rate = job["hatch_rate"]
#self.num_clients = job["num_clients"]
self.num_requests = job["num_requests"]
self.host = job["host"]
self.hatching_greenlet = gevent.spawn(
lambda: self.start_hatching(locust_count=job["num_clients"
],
hatch_rate=job["hatch_rate"]))
elif msg.type == "stop":
self.stop()
self.client.send(
Message("client_stopped", None, self.client_id))
self.client.send(Message("client_ready", None, self.client_id))
elif msg.type == "quit":
logger.info("Got quit message from master, shutting down...")
self.stop()
self.greenlet.kill(block=True)
def stats_reporter(self):
while True:
data = {}
events.report_to_master.fire(client_id=self.client_id, data=data)
try:
self.client.send(Message("stats", data, self.client_id))
except:
logger.error("Connection lost to master server. Aborting...")
break
gevent.sleep(SLAVE_REPORT_INTERVAL)
