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)