fix race condition casued exit

Author: Twi1ight

core/utils/redis_utils.py

@@ -153,7 +153,7 @@ def create_task_from_url(self, url, add_whitelist=True, valid_url_check=True):
         :param add_whitelist: for init scan task, disabled in task result produce
         :return:
         """
-        if valid_url_check and not self.valid_task_url(url): return
+        if valid_url_check and not self.valid_task_url(url): return False
 
         logger.info('add task: %s' % url.urlstring)
         self.redis_client.lpush(self.l_url_tasks, url.urlstring)
@@ -163,6 +163,7 @@ def create_task_from_url(self, url, add_whitelist=True, valid_url_check=True):
         self.set_url_scanned(url)
         # incr req count
         self.incr_hostname_reqcount(url.hostname)
+        return True
 
     def valid_task_url(self, url):
         """

core/worker/consumer.py

@@ -10,6 +10,7 @@
 from core.spider.spider import SpiderPage
 from core.utils.log import logger
 from core.utils.redis_utils import RedisUtils
+from core.utils.url import URL
 
 
 class Consumer(object):
@@ -18,33 +19,38 @@ def __init__(self, **kwargs):
         :return: :class:Consumer object
         :rtype: Consumer
         """
+        self.context = kwargs.pop('context')
         self.__cookie_file = kwargs.pop('cookie_file')
         self.redis_handle = RedisUtils(db=kwargs.pop('redis_db'), tld=kwargs.pop('tld'))
 
-    def consume(self, tspider_context):
+    def consume(self):
         if not self.redis_handle.connected:
             logger.error('no redis connection found in consumer! exit.')
             return
         while True:
             try:
                 url = self.redis_handle.fetch_one_task()
+                with self.context['lock']:
+                    self.context['live_spider_counts'].value += 1
+                    self.context['task_counts'].value -= 1
                 logger.info('get task url: %s' % url)
-                logger.info('%d tasks left' % self.redis_handle.task_counts)
-                with tspider_context['lock']:
-                    tspider_context['live_spider_counts'].value += 1
-                self.start_spider(url, self.__cookie_file)
+                logger.info('%d tasks left' % self.context['task_counts'].value)
+                if not self.redis_handle.is_blocked(URL(url)):
+                    self.start_spider(url, self.__cookie_file)
             except:
                 logger.exception('consumer exception!')
                 if not self.redis_handle.connected:
                     logger.error('redis disconnected! reconnecting...')
                     self.redis_handle.connect()
                 time.sleep(10)
             finally:
-                with tspider_context['lock']:
-                    tspider_context['live_spider_counts'].value -= 1
+                with self.context['lock']:
+                    self.context['live_spider_counts'].value -= 1
 
     def start_spider(self, url, cookie_file=None):
         results = SpiderPage(url, cookie_file=cookie_file).spider()
+        with self.context['lock']:
+            self.context['result_counts'].value += len(results)
         for _ in results:
             self.redis_handle.insert_result(_)
 

core/worker/producer.py

@@ -26,11 +26,12 @@ def __init__(self, **kwargs):
         :return: :class:Producer object
         :rtype: Producer
         """
+        self.context = kwargs.pop('context')
         self.__mongo_db = kwargs.pop('mongo_db')
         self.mongo_handle = None
         self.redis_handle = RedisUtils(db=kwargs.pop('redis_db'), tld=kwargs.pop('tld'))
 
-    def produce(self, tspider_context):
+    def produce(self):
         # mongodb with multipleprocessing must be init after fork
         self.mongo_handle = MongoUtils(db=self.__mongo_db)
         if not self.redis_handle.connected or not self.mongo_handle.connected:
@@ -40,8 +41,9 @@ def produce(self, tspider_context):
         while True:
             try:
                 _, req = self.redis_handle.fetch_one_result()
-                remainder_result = self.redis_handle.result_counts
-                logger.debug('got req, %d results left' % remainder_result)
+                with self.context['lock']:
+                    self.context['result_counts'].value -= 1
+                logger.debug('got req, %d results left' % self.context['result_counts'].value)
                 self.proc_req(req)
             except:
                 logger.exception('produce exception!')
@@ -53,11 +55,10 @@ def produce(self, tspider_context):
                     self.mongo_handle.connect()
                 time.sleep(10)
             finally:
-                if remainder_result == 0 and self.redis_handle.task_counts == 0:
-                    with tspider_context['lock']:
-                        live_spider_counts = tspider_context['live_spider_counts'].value
-                    if live_spider_counts == 0:
-                        tspider_context['task_done'].set()
+                with self.context['lock']:
+                    if self.context['result_counts'].value == 0:
+                        if self.context['live_spider_counts'].value == 0 and self.context['task_counts'].value == 0:
+                            self.context['task_done'].set()
 
     def proc_req(self, req):
         try:
@@ -99,16 +100,18 @@ def proc_req(self, req):
         elif method == 'GET':
             # new host found, add index page to task queue
             if self.redis_handle.get_hostname_reqcount(url.hostname) == 0:
-                self.redis_handle.create_task_from_url(URL(url.index_page), add_whitelist=False)
+                self.create_task_from_url(URL(url.index_page), add_whitelist=False)
             # check url validation inside create_url_task
-            self.redis_handle.create_task_from_url(url, add_whitelist=False)
+            self.create_task_from_url(url, add_whitelist=False)
         else:
             # not GET nor POST
             logger.error('HTTP Verb %s found!' % method)
             logger.debug(data)
 
-    def create_task_from_url(self, url):
-        self.redis_handle.create_task_from_url(url)
+    def create_task_from_url(self, url, **kwargs):
+        with self.context['lock']:
+            if self.redis_handle.create_task_from_url(url, **kwargs):
+                self.context['task_counts'].value += 1
 
     def create_task_from_file(self, fileobj):
         """
@@ -121,7 +124,7 @@ def create_task_from_file(self, fileobj):
                 line = line.strip()
                 if not line: continue
                 url = URL(line)
-                self.redis_handle.create_task_from_url(url)
+                self.create_task_from_url(url)
 
 
 if __name__ == '__main__':

tools/block_domain.py

@@ -8,28 +8,9 @@
 See the file 'doc/COPYING' for copying permission
 """
 import sys
-import traceback
 
-from core.utils.url import URL
 from core.utils.redis_utils import RedisUtils
 
-
-def remove_from_tasklist(domain):
-    urls = []
-    while True:
-        try:
-            urlstring = r.fetch_one_task(timeout=3)
-            url = URL(urlstring)
-            if r.is_blocked(url):
-                continue
-            urls.append(urlstring)
-        except:
-            break
-
-    for url in urls:
-        r.create_task_from_url(URL(url), add_whitelist=False, valid_url_check=False)
-
-
 if __name__ == '__main__':
     if len(sys.argv) != 3:
         print 'usage: block_domain.py db target.com'
@@ -38,9 +19,4 @@ def remove_from_tasklist(domain):
     domain = sys.argv[2]
     r = RedisUtils(db=db)
     r.add_blocklist(domain)
-
-    if r.redis_client.hexists(r.h_blocklist, domain):
-        remove_from_tasklist(domain)
-        print 'add success!'
-    else:
-        print 'add failed!'
+    print 'add success!'

tspider.py

@@ -62,18 +62,20 @@ def cmdparse():
             os.remove(os.path.join(TMPDIR_PATH, f))
     tspider_context = {
         'live_spider_counts': Value('i', 0),
+        'task_counts': Value('i', 0),
+        'result_counts': Value('i', 0),
         'task_done': Event(),
         'lock': Lock()}
-    kwargs = {'tld': args.tld, 'cookie_file': args.cookie_file,
+    kwargs = {'tld': args.tld, 'cookie_file': args.cookie_file, 'context': tspider_context,
               'redis_db': args.redis_db, 'mongo_db': args.mongo_db}
     for _ in range(args.consumer):
         worker = Consumer(**kwargs).consume
-        proc = Process(name='consumer-%d' % _, target=worker, args=(tspider_context,))
+        proc = Process(name='consumer-%d' % _, target=worker)
         proc.daemon = True
         proc.start()
     for _ in range(args.producer):
         worker = Producer(**kwargs).produce
-        proc = Process(name='producer-%d' % _, target=worker, args=(tspider_context,))
+        proc = Process(name='producer-%d' % _, target=worker)
         proc.daemon = True
         proc.start()