package com.javaweb.spider.scheduler;
|
|
import com.javaweb.spider.domain.ExitWayEnum;
|
import com.javaweb.spider.domain.SpiderConfig;
|
import us.codecraft.webmagic.Request;
|
import us.codecraft.webmagic.Task;
|
import us.codecraft.webmagic.scheduler.DuplicateRemovedScheduler;
|
import us.codecraft.webmagic.scheduler.MonitorableScheduler;
|
|
import java.util.concurrent.BlockingQueue;
|
import java.util.concurrent.LinkedBlockingQueue;
|
|
/**
|
* 自定义的调度器，主要用来处理url限制爬取的条数
|
*/
|
public class CountDownScheduler extends DuplicateRemovedScheduler implements MonitorableScheduler {
|
private BlockingQueue<Request> queue = new LinkedBlockingQueue<>();
|
private int realUrlCount = -1;
|
|
public CountDownScheduler(SpiderConfig spiderConfig) {
|
if (ExitWayEnum.URL_COUNT.toString().equals(spiderConfig.getExitWay())) {
|
// 实际抓取的url数量包括入口页面
|
this.realUrlCount = spiderConfig.getCount() + spiderConfig.getEntryUrlsList().size();
|
}
|
}
|
|
@Override
|
public void pushWhenNoDuplicate(Request request, Task task) {
|
// 当程序退出方式非URL_COUNT时按照正常逻辑处理
|
if (realUrlCount == -1) {
|
this.queue.add(request);
|
return;
|
}
|
// 在有效期内(realUrlCount > 0)，每次push url时realUrlCount - 1， 当 realUrlCount <= 0 时，当前Scheduler将不再收录新的url
|
if (realUrlCount <= 0) {
|
return;
|
}
|
realUrlCount--;
|
this.queue.add(request);
|
}
|
|
@Override
|
public Request poll(Task task) {
|
return (Request) this.queue.poll();
|
}
|
|
@Override
|
public int getLeftRequestsCount(Task task) {
|
return this.queue.size();
|
}
|
|
@Override
|
public int getTotalRequestsCount(Task task) {
|
return this.getDuplicateRemover().getTotalRequestsCount(task);
|
}
|
}
|
