exit when task done

Author: Twi1ight

core/worker/consumer.py

@@ -21,23 +21,27 @@ def __init__(self, **kwargs):
         self.__cookie_file = kwargs.pop('cookie_file')
         self.redis_handle = RedisUtils(db=kwargs.pop('redis_db'), tld=kwargs.pop('tld'))
 
-    def consume(self):
+    def consume(self, tspider_context):
         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()
                 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)
             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
 
     def start_spider(self, url, cookie_file=None):
         results = SpiderPage(url, cookie_file=cookie_file).spider()

core/worker/producer.py

@@ -30,7 +30,7 @@ def __init__(self, **kwargs):
         self.mongo_handle = None
         self.redis_handle = RedisUtils(db=kwargs.pop('redis_db'), tld=kwargs.pop('tld'))
 
-    def produce(self):
+    def produce(self, tspider_context):
         # 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,7 +40,8 @@ def produce(self):
         while True:
             try:
                 _, req = self.redis_handle.fetch_one_result()
-                logger.debug('got req, %d results left' % self.redis_handle.result_counts)
+                remainder_result = self.redis_handle.result_counts
+                logger.debug('got req, %d results left' % remainder_result)
                 self.proc_req(req)
             except:
                 logger.exception('produce exception!')
@@ -51,6 +52,12 @@ def produce(self):
                     logger.error('mongodb disconnected! reconnecting...')
                     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()
 
     def proc_req(self, req):
         try:

doc/CHANGELOG.md

@@ -1,3 +1,10 @@
+### (2017-04-11)
+
+- 修复自适应协议（//开头的url）导致的异常
+- 修复urlparse后用urlunsplit重新组合url导致的错误
+- 修复request捕获的url数据未加type字段
+- 新增任务扫描完成后退出功能
+
 ### (2017-04-05)
 
 - 添加MutationObserver支持

tspider.py

@@ -7,16 +7,17 @@
 Copyright (c) 2016-2017 [email protected] (http://twi1ight.com/)
 See the file 'doc/COPYING' for copying permission
 """
+import os
 import sys
 import argparse
-from multiprocessing import Process
+from multiprocessing import Process, Value, Lock, Event
 
 from core.utils.log import logger
 from core.utils.redis_utils import RedisUtils
 from core.utils.url import URL
 from core.worker.consumer import Consumer
 from core.worker.producer import Producer
-from settings import VERSION, RedisConf, MongoConf
+from settings import VERSION, RedisConf, MongoConf, TMPDIR_PATH
 
 
 def cmdparse():
@@ -52,25 +53,29 @@ def cmdparse():
 
 if __name__ == '__main__':
     args = cmdparse()
-    producer_pool = []
-    consumer_pool = []
     redis_handle = RedisUtils(db=args.redis_db)
     if args.keepon:
         redis_handle.restore_startup_params(args)
         logger.info(args)
 
+    for f in os.listdir(TMPDIR_PATH):
+        os.remove(os.path.join(TMPDIR_PATH, f))
+    tspider_context = {}
+    tspider_context['live_spider_counts'] = Value('i', 0)
+    tspider_context['task_done'] = Event()
+    tspider_context['lock'] = Lock()
     kwargs = {'tld': args.tld, 'cookie_file': args.cookie_file,
               '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)
+        proc = Process(name='consumer-%d' % _, target=worker, args=(tspider_context,))
+        proc.daemon = True
         proc.start()
-        consumer_pool.append(proc)
     for _ in range(args.producer):
         worker = Producer(**kwargs).produce
-        proc = Process(name='producer-%d' % _, target=worker)
+        proc = Process(name='producer-%d' % _, target=worker, args=(tspider_context,))
+        proc.daemon = True
         proc.start()
-        producer_pool.append(proc)
 
     if not args.keepon:
         redis_handle.flushdb()
@@ -88,5 +93,4 @@ def cmdparse():
             producer.create_task_from_file(target)
 
     redis_handle.close()
-    map(lambda x: x.join(), consumer_pool)
-    map(lambda x: x.join(), producer_pool)
+    tspider_context['task_done'].wait()