def __init__(self, stream): self._stream = stream # TODO: use lockable!! self._write_semaphore = trio.Semaphore(1) # blocks reading from stream and _read_buf at the same time self._read_semaphore = trio.Semaphore(1) self._read_buf = bytearray()
async def run(self) -> None: self.logger.info( "Started Exfiltrator: %s..%s", self.start_at, "HEAD" if self.end_at is None else self.end_at, ) semaphor = trio.Semaphore( self._concurrency_factor, max_value=self._concurrency_factor ) (relay_send_channel, relay_receive_channel) = trio.open_memory_channel[Block]( self._concurrency_factor - 1 ) self.manager.run_daemon_task(self._collate_and_relay, relay_receive_channel) async def _fetch(block_number: int) -> None: # TODO: handle web3 failures self.logger.debug("Retrieving block #%d", block_number) block = await retrieve_block(self.w3, block_number) self.logger.debug("Retrieved block #%d", block_number) semaphor.release() await relay_send_channel.send(block) async with self._block_send_channel: async with relay_send_channel: async with trio.open_nursery() as nursery: while self.manager.is_running: for block_number in iter_block_numbers( self.start_at, self.end_at ): await semaphor.acquire() nursery.start_soon(_fetch, block_number)
def __init__(self, rate_limiter, stats_tracker, robots_txt_manager, crawl_db, frontier_db, extractor_db, storage_db, login_db): ''' Constructor. :param starbelly.rate_limiter.RateLimiter rate_limiter: A rate limiter. :param StatsTracker stats_tracker: A stats tracking instance. :param starbelly.robots.RobotsTxtManager robots_txt_manager: A robots.txt manager. :param starbelly.db.CrawlManagerDb crawl_db: A database layer. :param starbelly.db.CrawlFrontierDb crawl_db: A database layer. :param starbelly.db.CrawlExtractorDb crawl_db: A database layer. :param starbelly.db.CrawlStorageDb crawl_db: A database layer. :param starbelly.db.LoginDb login_db: A database layer. ''' self._rate_limiter = rate_limiter self._stats_tracker = stats_tracker self._robots_txt_manager = robots_txt_manager self._db = crawl_db self._frontier_db = frontier_db self._extractor_db = extractor_db self._storage_db = storage_db self._login_db = login_db self._download_capacity = 20 self._download_semaphore = trio.Semaphore(self._download_capacity) self._jobs = dict() self._job_state_channels = dict() self._nursery = None self._sequence = None
def __init__(self, addr, device_id, signing_key, max): self.addr = addr self.device_id = device_id self.signing_key = signing_key self.transports = [] self._closed = False self._lock = trio.Semaphore(max)
def get_mutex(self): if self.mutex == None: self.mutex = trio.Semaphore(1) else: pass log.debug('returning mutex') return self.mutex
async def open_nursery_with_capacity(conc): """A patched Trio.Nursery with child task capacity limit. Its ``start_once_acquired`` blocks when the number of running child tasks started by it exceeds ``conc``. e.g.:: async with open_nursery_with_capacity(10) as nursery: for _ in range(30): await nursery.start_once_acquired(trio.sleep, 1) """ sema = trio.Semaphore(conc) async def _rl_task(fn, *args, task_status=trio.TASK_STATUS_IGNORED): async with sema: task_status.started() await fn(*args) async def start_once_acquired(self, fn, *args): await self.start(_rl_task, fn, *args) async with trio.open_nursery() as _n: _n.start_once_acquired = types.MethodType(start_once_acquired, _n) yield _n
def __init__(self, client_factory, minimum=1, maximum=10): self.client_factory = client_factory self.minimum = minimum self.maximum = maximum self._limit = trio.Semaphore(maximum) self._free = [] self._not_free = []
def semaphore(self) -> typing.Optional[trio.Semaphore]: if not hasattr(self, "_semaphore"): max_connections = self.pool_limits.hard_limit if max_connections is None: self._semaphore = None else: self._semaphore = trio.Semaphore(max_connections, max_value=max_connections) return self._semaphore
async def run(): """ 运行主函数 """ # 创建可复用 Semaphore 实例,减少开销 sem = trio.Semaphore(MAX_CONCURRENCY) async with trio.open_nursery() as nursery: while not Q.empty(): nursery.start_soon(download, sem, await Q.get())
def __init__( self, client: AlexandriaClientAPI, advertisement_dbs: Sequence[AdvertisementDatabaseAPI], concurrency: int = 3, ) -> None: self._client = client self._advertisement_dbs = tuple(advertisement_dbs) self._concurrency_lock = trio.Semaphore(concurrency) self._ready = trio.Event()
def __init__(self, server, start_room, checker:PathChecker, n_results = 3): assert start_room is not None self.s = server self.checker = checker self.start_room = start_room self.results = [] # (destination,path) self.n_results = n_results self._n_results = trio.Semaphore(n_results) self._stall_wait = trio.Event()
async def test_socks_proxy_get(asyncio_loop, nursery): # Create a server: async def handler(stream): # Negotiate SOCKS authentication (no auth) request = await stream.receive_some(4096) assert request.startswith(b'\x05') await stream.send_all(b'\x05\x00') # Negotiate SOCKS command. request = await stream.receive_some(4096) assert request == b'\x05\x01\x00\x01\x7f\x00\x00\x01\x00\x50' await stream.send_all(b'\x05\x00\x00\x01\x7f\x00\x00\x01\x00\x50') # Get HTTP request and send response. request = await stream.receive_some(4096) assert request.startswith(b'GET /foo HTTP/1.1') assert request.endswith(b'\r\n\r\n') await stream.send_all( b'HTTP/1.1 200 OK\r\n' b'Content-type: text/html\r\n' b'\r\n' b'<html><head><title>Unit Test</title></head>\r\n' b'<body><h1>This is a unit test</h1></body></html>\r\n' b'\r\n') await stream.aclose() serve_tcp = partial(trio.serve_tcp, handler, port=0, host='127.0.0.1') socks_proxy = await nursery.start(serve_tcp) addr, port = socks_proxy[0].socket.getsockname() # Run the test: job_id = 'bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb' policy = make_policy( proxy=[{ 'proxy_url': 'socks5://{}:{}'.format(addr, port), }]) request_send, request_recv = trio.open_memory_channel(0) response_send, response_recv = trio.open_memory_channel(0) semaphore = trio.Semaphore(1) rate_limiter_reset = Mock() rate_limiter_reset.send = AsyncMock() stats = make_stats() dl = Downloader(job_id, policy, response_send, request_recv, semaphore, rate_limiter_reset, stats) nursery.start_soon(dl.run) request = make_request('http://127.0.0.1/foo', policy=policy) await request_send.send(request) response = await response_recv.receive() nursery.cancel_scope.cancel() assert response.status_code == 200 assert response.content_type == 'text/html' assert response.body.startswith(b'<html>') assert stats['item_count'] == 1 assert stats['http_success_count'] == 1 assert stats['http_status_counts'][200] == 1
async def test_http_post(nursery): # Create a server: async def handler(stream): request = b'' while True: request = await stream.receive_some(4096) if request.endswith(b'\r\n\r\n'): break assert request.startswith(b'POST /foo HTTP/1.1\r\n') await stream.send_all( b'HTTP/1.1 200 OK\r\n' b'Content-type: text/html\r\n' b'\r\n' b'<html><head><title>Unit Test</title></head>\r\n' b'<body><h1>This is a unit test</h1></body></html>\r\n' b'\r\n') await stream.aclose() serve_tcp = partial(trio.serve_tcp, handler, port=0, host='127.0.0.1') http_server = await nursery.start(serve_tcp) addr, port = http_server[0].socket.getsockname() # Run the test: job_id = 'bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb' policy = make_policy() request_send, request_recv = trio.open_memory_channel(0) response_send, response_recv = trio.open_memory_channel(0) semaphore = trio.Semaphore(1) rate_limiter_reset = Mock() rate_limiter_reset.send = AsyncMock() stats = make_stats() dl = Downloader(job_id, policy, response_send, request_recv, semaphore, rate_limiter_reset, stats) nursery.start_soon(dl.run) request = make_request('http://{}:{}/foo'.format(addr, port), method='POST', form_data={'user': '******'}) await request_send.send(request) response = await response_recv.receive() nursery.cancel_scope.cancel() assert response.status_code == 200 assert response.content_type == 'text/html' assert response.body.startswith(b'<html>') assert stats['item_count'] == 1 assert stats['http_success_count'] == 1 assert stats['http_status_counts'][200] == 1
def __init__(self, capacity): ''' Constructor. :param int capacity: The maximum number of items to buffer inside of the rate limiter. ''' self._expires = list() self._expiry_cancel_scope = None self._global_limit = None self._queues = dict() self._rate_limits = dict() self._capacity = capacity self._semaphore = trio.Semaphore(capacity) self._request_send, self._request_recv = trio.open_memory_channel(0) self._reset_send, self._reset_recv = trio.open_memory_channel(0) self._job_channels = dict()
async def retrieve_headers(w3: Web3, start_at: int, end_at: int, send_channel: trio.abc.SendChannel[BlockHeader], request_rate: int = 3) -> None: semaphor = trio.Semaphore(request_rate, max_value=request_rate) async def _fetch(block_number: int) -> None: header = await retrieve_header(w3, block_number) semaphor.release() await send_channel.send(header) async with send_channel: async with trio.open_nursery() as nursery: logger.debug('Starting retrieval of headers %d-%d', start_at, end_at) for block_number in range(start_at, end_at): await semaphor.acquire() nursery.start_soon(_fetch, block_number)
async def QueueLimiter(conc): """Rate limit for the submission of delayed functions. Grain by default enqueues delayed functions eagerly (i.e. as soon as it is called). Sometimes if we have a lot of functions that can be run in parallel, we don't want to overwhelm the queue, so we could set a rate limit check prior submission. Note that the context scope blocks until all functions submitted through it are done. Args: conc (int): maximum number of concurrently running functions e.g.:: @delayed async def dfn(x): await trio.sleep(1) return x+1 r_ = 0 async with QueueLimiter(10) as ql: for i in range(30): r_ += await ql(dfn)(i) # ^^^^^wait for submission, not for result r = await r_ """ sema = trio.Semaphore(conc) async def sema2notify(dfn, args, kwargs, task_status): async with sema: do = dfn(*args, **kwargs) task_status.started(do) await do # sema is not released until calculation is done def _rl_wrapper(dfn): async def _rl_dfn(*args, **kwargs): do = await _n.start(sema2notify, dfn, args, kwargs) return do return _rl_dfn async with trio.open_nursery() as _n: yield _rl_wrapper
async def _fetch_blocks(self) -> None: semaphor = trio.Semaphore(self._concurrency_factor, max_value=self._concurrency_factor) (relay_send_channel, relay_receive_channel) = trio.open_memory_channel[Block](0) self.manager.run_task(self._collate_and_relay, relay_receive_channel) async def _fetch(block_number: BlockNumber) -> None: # TODO: handle web3 failures self.logger.debug("Retrieving block #%d", block_number) block = await retrieve_block(self.w3, block_number) self.logger.debug("Retrieved block #%d", block_number) semaphor.release() await relay_send_channel.send(block) async with self._block_send_channel: async with relay_send_channel: async with trio.open_nursery() as nursery: for block_number in range(self.start_at, self.end_at): await semaphor.acquire() nursery.start_soon(_fetch, block_number)
async def main(paths, rate_limit): s = Session(connections=3) # mutual exclusion for getter worker mut_ex = trio.Semaphore(1) interval = 1 / rate_limit async def tick(): await trio.sleep(interval) mut_ex.release() async def getter(session, path): await mut_ex.acquire() n.start_soon(tick) logging.info(f'Sending to {path}. Current responses: {len(results)}') resp = await session.get(path) # we do not raise_for_status() here; process the batch afterwards results.append(resp) async with trio.open_nursery() as n: for path in path_list: n.start_soon(getter, s, path)
def __init__(self): self.unpaired = trio.Semaphore(0, max_value=1) self.done = trio.Event() self.ident = next(self.counter)
def __init__(self, connect_cb, max_pool): self._connect_cb = connect_cb self._transports = [] self._closed = False self._lock = trio.Semaphore(max_pool)
def CombineGroup_ctxt(exer, push_newgroup): sema = trio.Semaphore(300) # rate limit gid = 0 # monotonously incremental with __aenter__'s side-effect class __CombineGroup(object): """A wait group that submit jobs to backend GrainExecutor and collect the results. This context manager is reusable, but single-use is recommended, because creating a instance each time allows nesting multiple ones. """ def __init__(self): self.counter = 0 self.results = [] async def asubmit(self, res, fn, *args, **kwargs): await exer.asubmit(res, _grouped_task, self.gid, fn, *args, **kwargs) self.counter += 1 def submit(self, res, fn, *args, **kwargs): exer.submit(res, _grouped_task, self.gid, fn, *args, **kwargs) self.counter += 1 def start_subtask(self, fn, *args, **kwargs): # tasks requesting zero resource would be executed locally s, r = trio.open_memory_channel(1) exer.submit(ZERO, _grouped_task, self.gid, r.receive) _gn.start_soon( partial(_run_and_send_result, s, fn, *args, **kwargs)) self.counter += 1 def load_cache_or_submit(self, res, fn, *args, **kwargs): self.start_subtask(load_cache_or_exec1, res, fn, *args, **kwargs) async def __aenter__(self): # async part of __init__ if self.counter > 0: raise RuntimeError( "attempt to re-enter a CombineGroup. For recursive use, create a new instance instead." ) await sema.acquire() nonlocal gid self.gid = gid = gid + 1 s, self.resultq = trio.open_memory_channel(INFIN) push_newgroup.send_nowait((-1, (self.gid, s))) return self async def __aexit__(self, *exc): if any(exc): return False resultd = {} async with self.resultq: if self.counter > 0: async for i, r in self.resultq: resultd[i] = r self.counter -= 1 if self.counter == 0: break self.results = [ v for k, v in sorted(resultd.items(), key=lambda x: x[0]) ] sema.release() async def __Exec1(res, fn, *args, **kwargs): nonlocal gid g = gid = gid + 1 sq, rq = trio.open_memory_channel(1) push_newgroup.send_nowait((-1, (g, sq))) exer.submit(res, _grouped_task, g, fn, *args, **kwargs) async with rq, sq: return (await rq.receive())[1] global CombineGroup, Exec1 CombineGroup, Exec1 = __CombineGroup, __Exec1 try: yield finally: CombineGroup, Exec1 = None, None
async def main(): strip = False if "TRAVIS" in os.environ else None colorama.init(autoreset=True, strip=strip) parser = ArgumentParser() parser.add_argument("--limit", type=int) parser.add_argument("--workers", type=int, default=8) parser.add_argument("--post-batches", type=int, default=10) args = parser.parse_args() limit = args.limit post_batches = args.post_batches pytest_version = os.environ["PYTEST_VERSION"] # important to remove POST_KEY from environment so others cannot sniff it somehow (#26) secret = os.environ.pop("POST_KEY", None) if secret is None and limit is None: # bail out early so CI doesn't take forever for a PR limit = args.post_batches * 3 print(Fore.CYAN + "Limit forced to {} since secret is unavailable".format(limit)) tox_env = "py%d%d" % sys.version_info[:2] plugins = read_plugins_index(update_index.INDEX_FILE_NAME) if limit is not None: plugins = plugins[:limit] n_total = len(plugins) print(Fore.CYAN + f"Processing {len(plugins)} packages with {args.workers} workers") tmp = mkdtemp() async with asks.Session() as session: results_poster = ResultsPoster( session, batch_size=post_batches, tox_env=tox_env, pytest_version=pytest_version, secret=secret, ) progress_counter = ProgressCounter(n_total) semaphore = trio.Semaphore(args.workers) with working_directory(tmp): async with trio.open_nursery() as nursery: for plugin in plugins: await nursery.start( process_package, semaphore, session, results_poster, progress_counter, tox_env, pytest_version, plugin["name"], plugin["version"], plugin["description"], ) await results_poster.post_all() print() if results_poster.total_posted: print(Fore.GREEN + f"Posted {results_poster.total_posted} new results") print(Fore.GREEN + "All done, congratulations :)") shutil.rmtree(tmp, ignore_errors=True)
def _semaphore_factory(self, value: int) -> LockLike: return cast(LockLike, trio.Semaphore(value))
def semaphore(self) -> trio.Semaphore: if not hasattr(self, "_semaphore"): self._semaphore = trio.Semaphore(self.max_value, max_value=self.max_value) return self._semaphore
def __init__(self, max_meter): self.sem = trio.Semaphore( initial_value=max_meter.max_at_once, max_value=max_meter.max_at_once )