Python gevent.pool.join() Examples

def download_delta_data(stocks, pool_size=40):
    """
    Download delta data for all stocks collections of all libraries.
    :param stocks: stock code list.
    :param pool_size: the pool size of gevent.pool.Pool.
    :return: None
    """

    pool = gevent.pool.Pool(pool_size)
    for i in range(len(stocks) // pool_size + 1):
        start = i * pool_size
        end = (i + 1) * pool_size
        lst = stocks[start:end]
        logger.debug(f'download delta data for stock list: {lst}')
        for stock in lst:
            pool.spawn(bdt.TsHisData.download_one_delta_data, stock)
        pool.join(timeout=30) 

def run(self):
        lookups = self.rectypes or ['CNAME', 'A', 'AAAA']
        dnsname = self.domain
        if self.name is None:
            # Top-level, needs extra queries
            lookups += ['MX', 'SOA', 'NS', 'SRV', 'TXT', 'SPF', 'RRSIG', 'DS',
                        'DLV', 'DNSKEY']
        else:
            dnsname = '.'.join([self.name, dnsname])
        for query_type in set(lookups):
            resp = None
            LOG.debug("Checking %s %s", dnsname, query_type)
            try:
                resp = self.bruter.query(dnsname, query_type)
            except DNSException:
                continue
            except Exception:
                LOG.exception("While resolving %s %s", dnsname, query_type)
                continue
            self.bruter.on_result(self.domain, self.name, query_type, resp)
        self.bruter.on_finish() 

def run(self):
        if not self._find_wildcards():
            return
        pool = gevent.pool.Pool(self.options.concurrency)
        namegen = DNSTesterGenerator(self, self.domains, self.names)
        LOG.info("Starting DNS brute force (%d tests)", namegen.total)
        self.finished = 0
        if self.progress:
            self.progress.start(namegen.total)
        try:
            for tester in namegen.all():
                pool.add(gevent.spawn(tester.run))
        except KeyboardInterrupt:
            print("Ctrl+C caught... stopping")
        pool.join()
        if self.progress:
            self.progress.finish() 

def start(self):
        pool = gevent.pool.Pool(size=self.concurrency)
        try:
            for i in xrange(1, self.num_connectors + 1):
                pool.spawn(self.connector)
                time.sleep(self.spawn_interval)

            pool.join()
        except KeyboardInterrupt:
            pass 

def schedule_green_jobs(
    fns, concurrency=DEFAULT_THREADS, 
    progress=None, total=None
  ):
  import gevent.pool

  if total is None:
    try:
      total = len(fns)
    except TypeError: # generators don't have len
      pass

  pbar = tqdm(total=total, desc=progress, disable=(not progress))
  results = []
  
  def updatefn(fn):
    def realupdatefn():
      res = fn()
      pbar.update(1)
      results.append(res)
    return realupdatefn

  pool = gevent.pool.Pool(concurrency)
  for fn in fns:
    pool.spawn( updatefn(fn) )

  pool.join()
  pool.kill()
  pbar.close()

  return results 

def run_items(self, items, session, workers=None):
        import gevent
        import gevent.monkey
        import gevent.pool
        gevent.monkey.patch_all()
        pool = gevent.pool.Pool(size=workers)
        for index, item in enumerate(items):
            pool.spawn(self._run_next_item, session, item, index)
        pool.join() 

def wait(green_thread, *args, **kwargs):
    if CONCURRENCY_LIBRARY == 'eventlet':
        return green_thread.wait(*args, **kwargs)
    elif CONCURRENCY_LIBRARY == 'gevent':
        return green_thread.join(*args, **kwargs)
    else:
        raise ValueError('Unsupported concurrency library') 

def green_pool_wait_all(pool):
    """
    Wait for all the green threads in the pool to finish.
    """
    if CONCURRENCY_LIBRARY == 'eventlet':
        return pool.waitall()
    elif CONCURRENCY_LIBRARY == 'gevent':
        # NOTE: This mimicks eventlet.waitall() functionallity better than
        # pool.join()
        return all(gl.ready() for gl in pool.greenlets)
    else:
        raise ValueError('Unsupported concurrency library') 

def check(self):
        pool = gevent.pool.Pool(len(self.proxies))
        for index, ip in enumerate(self.proxies):
            pool.apply_async(self.ip_delay, (index, ip))
        pool.join()
        self.trigger.emit([0]) 

def _output_result(self, domain, name, query_type, result):
        """
        Output results, in various formats, to necessary places
        """
        # To console
        if name is None:
            dnsname = domain
        else:
            dnsname = '.'.join([name, domain])
        res_keys = ' '.join(['='.join([key, str(value)])
                             for key, value in result.items()])
        info = ' '.join([dnsname, query_type, res_keys])
        if not self.options.quiet:
            print(info)
        #
        # Shit out same as console, but to file
        output = self.options.output
        if output:
            output.write(info + "\n")
            output.flush()
        #
        # Optionally shit out JSON
        outjson = self.options.json
        if outjson:
            outdict = result.copy()
            outdict['_type'] = query_type
            outdict['_domain'] = domain
            outdict['_name'] = name
            outdict.update(self.options.extra)
            if name and name[0] == '*':
                outdict['_wildcard'] = True
            outjson.write(json.dumps(outdict) + "\n")
            outjson.flush() 

def _dnsresp_to_dict(self, obj):
        """
        Converts DNS reponse into a normalised dictionary
        """
        rdtype = obj.rdtype
        if rdtype in (dns.rdatatype.A, dns.rdatatype.AAAA):
            return dict(host=obj.address)
        elif rdtype == dns.rdatatype.SOA:
            return dict(retry=obj.retry, serial=obj.serial, expires=obj.expire,
                        refresh=obj.refresh, minttl=obj.minimum,
                        hostmaster=str(obj.rname), nsname=str(obj.mname))
        elif rdtype == dns.rdatatype.NS:
            return dict(host=str(obj.target))
        elif rdtype == dns.rdatatype.MX:
            return dict(priority=obj.preference, host=str(obj.exchange))
        elif rdtype == dns.rdatatype.CNAME:
            return dict(cname=str(obj.target))
        elif rdtype in (dns.rdatatype.TXT, dns.rdatatype.SPF):
            return dict(text=" ".join(obj.strings))
        elif rdtype == dns.rdatatype.SRV:
            return dict(priority=obj.priority, host=str(obj.target), port=obj.port,
                        weight=obj.weight)
        elif rdtype == dns.rdatatype.DS:
            return dict(keytag=obj.key_tag, hashtype=obj.digest_type,
                        hash=hexlify(obj.digest))
        elif rdtype == dns.rdatatype.DLV:
            return dict(keytag=obj.key_tag, hashtype=obj.digest_type)
        elif rdtype == dns.rdatatype.DNSKEY:
            return dict(keytag=dns.dnssec.key_id(obj), protocol=obj.protocol,
                        flags=obj.flags, algorithm=obj.algorithm,
                        length=keylength(obj.algorithm, obj.key),
                        key=hexlify(obj.key))
        raise RuntimeError("Unknown DNS response type %r" % (obj,))
        #  'RRSIG', 'DS', 'DLV', 'DNSKEY', 'NSEC', 'NSEC3', 'NSEC3PARAM']
        # TODO: add DS, DLV, RRSIG, NSEC, NSEC3, PTR, DNSKEY, SSHFP, NAPTR 

def _find_wildcards(self):
        """
        Queries some random non-existant records to reduce false positives.
        Returns True if process can continue, otherwise false.
        """
        wildcard_count = self.options.wildcard_tests
        if wildcard_count < 1:
            return True
        total_queries = len(self.domains) * wildcard_count
        LOG.info("Eliminating wildcard responses (%d tests)", total_queries)
        is_ok = False
        # Setup pool and progress
        pool = gevent.pool.Pool(self.options.concurrency)
        if self.progress:
            self.progress.start(total_queries)
        self.finished = 0
        try:
            for domain in self.domains:
                LOG.debug("Checking wildcard domain: %s", domain)
                names = [rand_name() for _ in range(0, wildcard_count)]
                for name in names:
                    pool.add(gevent.spawn(self._test_wildcard, domain, name))
            is_ok = True
        except KeyboardInterrupt:
            print("Ctrl+C caught... stopping")
        pool.join()
        if self.progress:
            self.progress.finish()
        return is_ok 

def fetch_image(url, folder, filename, max_size=20 * 1024, formats=['png'], dimensions=(48, 48), fetch_timeout=1):
    def make_data_dir(subfolder):
        path = os.path.join(config.data_dir, subfolder)
        if not os.path.exists(path):
            os.makedirs(path)
        return path

    try:
        # fetch the image data
        try:
            r = grequests.map((grequests.get(url, timeout=fetch_timeout, headers={'Connection': 'close'}, verify=False, stream=True),))[0]
            if r is None:
                raise Exception("result is None")

            raw_image_data = r.iter_content(chunk_size=max_size)  # read up to max_size
        except Exception as e:
            raise Exception("Got fetch_image request error: %s" % e)
        else:
            if r.status_code != 200:
                raise Exception("Bad status code returned from fetch_image: '%s'" % (r.status_code))
        finally:
            if r:
                r.close()

        # decode image data
        try:
            image = Image.open(io.StringIO(raw_image_data))
        except Exception as e:
            raise Exception("Unable to parse image data at: %s" % url)
        if image.format.lower() not in formats:
            raise Exception("Image is not a PNG: %s (got %s)" % (url, image.format))
        if image.size != dimensions:
            raise Exception("Image size is not 48x48: %s (got %s)" % (url, image.size))
        if image.mode not in ['RGB', 'RGBA']:
            raise Exception("Image mode is not RGB/RGBA: %s (got %s)" % (url, image.mode))
        imagePath = make_data_dir(folder)
        imagePath = os.path.join(imagePath, filename + '.' + image.format.lower())
        image.save(imagePath)
        os.system("exiftool -q -overwrite_original -all= %s" % imagePath)  # strip all metadata, just in case
        return True
    except Exception as e:
        logger.warn(e)
        return False