MultiPagePipeline.java

package us.codecraft.webmagic.pipeline;

import us.codecraft.webmagic.MultiPageModel;
import us.codecraft.webmagic.ResultItems;
import us.codecraft.webmagic.Task;
import us.codecraft.webmagic.utils.Experimental;
import us.codecraft.webmagic.utils.DoubleKeyMap;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * A pipeline combines the result in more than one page together.<br>
 * Used for news and articles containing more than one web page. <br>
 * MultiPagePipeline will store parts of object and output them when all parts are extracted.<br>
 *
 * @author code4crafter@gmail.com <br>
 * @since 0.2.0
 */
@Experimental
public class MultiPagePipeline implements Pipeline {

    private DoubleKeyMap<String, String, Boolean> pageMap = new DoubleKeyMap<String, String, Boolean>(ConcurrentHashMap.class);

    private DoubleKeyMap<String, String, MultiPageModel> objectMap = new DoubleKeyMap<String, String, MultiPageModel>(ConcurrentHashMap.class);

    @Override
    public void process(ResultItems resultItems, Task task) {
        Map<String, Object> resultItemsAll = resultItems.getAll();
        Iterator<Map.Entry<String, Object>> iterator = resultItemsAll.entrySet().iterator();
        while (iterator.hasNext()) {
            handleObject(iterator);
        }
    }

    private void handleObject(Iterator<Map.Entry<String, Object>> iterator) {
        Map.Entry<String, Object> objectEntry = iterator.next();
        Object o = objectEntry.getValue();
        //需要拼凑
        if (o instanceof MultiPageModel) {
            MultiPageModel multiPageModel = (MultiPageModel) o;
            //这次处理的部分，设置为完成
            pageMap.put(multiPageModel.getPageKey(), multiPageModel.getPage(), Boolean.FALSE);
            //每个key单独加锁
            synchronized (pageMap.get(multiPageModel.getPageKey())) {
                pageMap.put(multiPageModel.getPageKey(), multiPageModel.getPage(), Boolean.TRUE);
                //其他需要拼凑的部分
                if (multiPageModel.getOtherPages() != null) {
                    for (String otherPage : multiPageModel.getOtherPages()) {
                        Boolean aBoolean = pageMap.get(multiPageModel.getPageKey(), otherPage);
                        if (aBoolean == null) {
                            pageMap.put(multiPageModel.getPageKey(), otherPage, Boolean.FALSE);
                        }
                    }
                }
                //check if all pages are processed
                Map<String, Boolean> booleanMap = pageMap.get(multiPageModel.getPageKey());
                objectMap.put(multiPageModel.getPageKey(), multiPageModel.getPage(), multiPageModel);
                if (booleanMap == null) {
                    return;
                }
                // /过滤，这次完成的page item中，还未拼凑完整的item，不进入下一个pipeline
                for (Map.Entry<String, Boolean> stringBooleanEntry : booleanMap.entrySet()) {
                    if (!stringBooleanEntry.getValue()) {
                        iterator.remove();
                        return;
                    }
                }
                List<Map.Entry<String, MultiPageModel>> entryList = new ArrayList<Map.Entry<String, MultiPageModel>>();
                entryList.addAll(objectMap.get(multiPageModel.getPageKey()).entrySet());
                if (entryList.size() != 0) {
                    Collections.sort(entryList, new Comparator<Map.Entry<String, MultiPageModel>>() {
                        @Override
                        public int compare(Map.Entry<String, MultiPageModel> o1, Map.Entry<String, MultiPageModel> o2) {
                            try {
                                int i1 = Integer.parseInt(o1.getKey());
                                int i2 = Integer.parseInt(o2.getKey());
                                return i1 - i2;
                            } catch (NumberFormatException e) {
                                return o1.getKey().compareTo(o2.getKey());
                            }
                        }
                    });
                    // 合并
                    MultiPageModel value = entryList.get(0).getValue();
                    for (int i = 1; i < entryList.size(); i++) {
                        value = value.combine(entryList.get(i).getValue());
                    }
                    objectEntry.setValue(value);
                }
            }
        }

    }

}