多线程批量执行任务

  1. package com.gitee.kooder.utils;
  3. import java.util.List;
  4. import java.util.concurrent.ForkJoinTask;
  5. import java.util.concurrent.RecursiveAction;
  6. import java.util.function.Consumer;
  8. /**
  9.  * Batch task action
  10.  * @author Winter Lau<javayou@gmail.com>
  11.  */
  12. public final class BatchTaskRunner extends RecursiveAction {
  14.     protected int threshold = 5;
  15.     protected List taskList;
  16.     Consumer<List> action;
  18.     /**
  19.      * @param taskList      任务列表
  20.      * @param threshold     每个线程处理的任务数
  21.      */
  22.     private BatchTaskRunner(List taskList, int threshold, Consumer action) {
  23.         this.taskList = taskList;
  24.         this.threshold = threshold;
  25.         this.action = action;
  26.     }
  28.     /**
  29.      * 多线程批量执行任务
  30.      * @param taskList
  31.      * @param threshold
  32.      * @param action
  33.      */
  34.     public static <T> void execute(List<T> taskList, int threshold, Consumer<List<T>> action) {
  35.         new BatchTaskRunner(taskList, threshold, action).invoke();
  36.     }
  38.     @Override
  39.     protected void compute() {
  40.         if (taskList.size() <= threshold) {
  41.             this.action.accept(taskList);
  42.         }
  43.         else {
  44.             this.splitFromMiddle(taskList);
  45.         }
  46.     }
  48.     /**
  49.      * 任务中分
  50.      * @param list
  51.      */
  52.     private void splitFromMiddle(List list) {
  53.         int middle = (int)Math.ceil(list.size() / 2.0);
  54.         List leftList = list.subList(0, middle);
  55.         List RightList = list.subList(middle, list.size());
  56.         BatchTaskRunner left = newInstance(leftList);
  57.         BatchTaskRunner right = newInstance(RightList);
  58.         ForkJoinTask.invokeAll(left, right);
  59.     }
  61.     private BatchTaskRunner newInstance(List taskList) {
  62.         return new BatchTaskRunner(taskList, threshold, action);
  63.     }
  65. }
  1.   /**
  2.      * 读取指定目录下所有 json 文件并写入索引
  3.      * @param type
  4.      * @param action
  5.      * @param path
  6.      * @param thread_count
  7.      * @return file count
  8.      */
  9.     private static int importJsonInPath(String type, String action, Path path, int thread_count) throws IOException {
  10.         final AtomicInteger fc = new AtomicInteger(0);
  11.         thread_count = Math.min(MAX_THREAD_COUNT, Math.max(thread_count, 1));
  12.         try (
  13.                 IndexWriter writer = StorageFactory.getIndexWriter(type);
  14.                 TaxonomyWriter taxonomyWriter = StorageFactory.getTaxonomyWriter(type);
  15.                 Stream<Path> pathStream = Files.list(path);
  16.         ) {
  17.             List<Path> allFiles = pathStream.filter(p -> p.toString().endsWith(".json") && !Files.isDirectory(p)).collect(Collectors.toList());
  18.             int threshold = Math.max(allFiles.size()/thread_count, 1);
  19.             BatchTaskRunner.execute(allFiles, threshold, files -> {
  20.                 files.forEach( jsonFile -> {
  21.                     importJsonFile(type, action, jsonFile, writer, taxonomyWriter);
  22.                     fc.addAndGet(1);
  23.                 });
  24.             });
  25.         }
  26.         return fc.get();
  27.     }

redis队列

  1. package com.gitee.kooder.queue;
  3. import java.util.Collection;
  4. import java.util.List;
  6. /**
  7.  * 定义了获取索引任务的队列接口
  8.  * @author Winter Lau<javayou@gmail.com>
  9.  */
  10. public interface Queue extends AutoCloseable{
  12.     /**
  13.      * 队列的唯一名称
  14.      * @return
  15.      */
  16.     String type();
  18.     /**
  19.      * 添加任务到队列
  20.      * @param tasks
  21.      */
  22.     void push(Collection<QueueTask> tasks) ;
  24.     /**
  25.      * 从队列获取任务
  26.      * @return
  27.      */
  28.     List<QueueTask> pop(int count) ;
  30. }
  1. package com.gitee.kooder.queue;
  3. import com.gitee.kooder.core.KooderConfig;
  4. import org.apache.commons.lang3.StringUtils;
  6. import java.util.Properties;
  8. /**
  9.  * 队列工厂
  10.  * @author Winter Lau<javayou@gmail.com>
  11.  */
  12. public class QueueFactory {
  14.     static QueueProvider provider;
  16.     static {
  17.         Properties props = KooderConfig.getQueueProperties();
  18.         String type = StringUtils.trim(props.getProperty("provider"));
  19.         if("redis".equalsIgnoreCase(type))
  20.             provider = new RedisQueueProvider(props);
  21.         else if("embed".equalsIgnoreCase(type))
  22.             provider = new EmbedQueueProvider(props);
  23.     }
  25.     public final static QueueProvider getProvider() {
  26.         return provider;
  27.     }
  29. }
  1. package com.gitee.kooder.queue;
  3. import com.gitee.kooder.core.Constants;
  5. import java.util.Arrays;
  6. import java.util.List;
  8. /**
  9.  * 定义了获取索引任务的队列接口
  10.  * @author Winter Lau<javayou@gmail.com>
  11.  */
  12. public interface QueueProvider extends AutoCloseable {
  14.     List<String> TYPES = Arrays.asList(Constants.TYPE_REPOSITORY, Constants.TYPE_ISSUE, Constants.TYPE_CODE);
  16.     /**
  17.      * Provider 唯一标识
  18.      * @return
  19.      */
  20.     String name();
  22.     /**
  23.      * 获取支持的所有任务类型
  24.      * @return
  25.      */
  26.     default List<String> getAllTypes() {
  27.         return TYPES;
  28.     }
  30.     /**
  31.      * 获取某个任务类型的队列
  32.      * @param type
  33.      * @return
  34.      */
  35.     Queue queue(String type);
  37. }
  1. package com.gitee.kooder.queue;
  3. import io.lettuce.core.RedisClient;
  4. import io.lettuce.core.RedisURI;
  5. import io.lettuce.core.api.StatefulRedisConnection;
  6. import io.lettuce.core.api.sync.RedisCommands;
  7. import org.apache.commons.lang3.math.NumberUtils;
  8. import org.slf4j.Logger;
  9. import org.slf4j.LoggerFactory;
  11. import java.util.ArrayList;
  12. import java.util.Collection;
  13. import java.util.List;
  14. import java.util.Properties;
  16. /**
  17.  * 使用 Redis 队列
  18.  * @author Winter Lau<javayou@gmail.com>
  19.  */
  20. public class RedisQueueProvider implements QueueProvider {
  22.     private final static Logger log = LoggerFactory.getLogger(RedisQueueProvider.class);
  24.     private String host;
  25.     private int port;
  26.     private int database;
  27.     private String baseKey;
  28.     private String username;
  29.     private String password;
  31.     private RedisClient client;
  33.     /**
  34.      * Connect to redis
  35.      * @param props
  36.      */
  37.     public RedisQueueProvider(Properties props) {
  38.         this.host = props.getProperty("redis.host", "127.0.0.1");
  39.         this.port = NumberUtils.toInt(props.getProperty("redis.port"), 6379);
  40.         this.database = NumberUtils.toInt(props.getProperty("redis.database"), 1);
  41.         this.baseKey = props.getProperty("redis.key", "gsearch-queue");
  42.         this.username = props.getProperty("username");
  43.         this.password = props.getProperty("password");
  45.         RedisURI uri = RedisURI.create(host,port);
  46.         uri.setDatabase(this.database);
  47.         if(password != null)
  48.             uri.setPassword(password.toCharArray());
  49.         if(username != null)
  50.             uri.setUsername(username);
  52.         this.client = RedisClient.create(uri);
  54.         log.info("Connected to {} at {}}:{}}\n", getRedisVersion(), this.host, this.port);
  56.     }
  58.     private String getRedisVersion() {
  59.         try (StatefulRedisConnection<String, String> connection = client.connect()) {
  60.             RedisCommands<String, String> cmd = connection.sync();
  61.             return cmd.info("redis_version");
  62.         }
  63.     }
  65.     @Override
  66.     public String name() {
  67.         return "redis";
  68.     }
  70.     @Override
  71.     public Queue queue(String type) {
  72.         return new Queue() {
  74.             private String key = type + '@' + baseKey;
  76.             @Override
  77.             public String type() {
  78.                 return type;
  79.             }
  81.             @Override
  82.             public void push(Collection<QueueTask> tasks) {
  83.                 try (StatefulRedisConnection<String, String> connection = client.connect()) {
  84.                     RedisCommands<String, String> cmd = connection.sync();
  85.                     cmd.rpush(key, tasks.stream().map(t -> t.json()).toArray(String[]::new));
  86.                 }
  87.             }
  89.             @Override
  90.             public List<QueueTask> pop(int count) {
  91.                 String json = null;
  92.                 List<QueueTask> tasks = new ArrayList<>();
  93.                 try (StatefulRedisConnection<String, String> connection = client.connect()) {
  94.                     RedisCommands<String, String> cmd = connection.sync();
  95.                     do{
  96.                         json = cmd.lpop(key);
  97.                         if(json == null)
  98.                             break;
  99.                         QueueTask task = QueueTask.parse(json);
  100.                         if(task != null)
  101.                             tasks.add(task);
  102.                     }while(tasks.size() < count);
  103.                 }
  104.                 return tasks;
  105.             }
  107.             @Override
  108.             public void close() {}
  109.         };
  110.     }
  112.     @Override
  113.     public void close() {
  114.         client.shutdown();
  115.     }
  116. }
  1. package com.gitee.kooder.queue;
  3. import com.gitee.kooder.core.KooderConfig;
  4. import org.apache.commons.lang3.math.NumberUtils;
  5. import org.infobip.lib.popout.FileQueue;
  6. import org.slf4j.Logger;
  7. import org.slf4j.LoggerFactory;
  9. import java.io.IOException;
  10. import java.nio.file.Files;
  11. import java.nio.file.Path;
  12. import java.util.*;
  13. import java.util.concurrent.ConcurrentHashMap;
  15. /**
  16.  * 实现 Gitee Search 内嵌式的队列，不依赖第三方服务，通过 HTTP 方式提供对象获取
  17.  * @author Winter Lau<javayou@gmail.com>
  18.  */
  19. public class EmbedQueueProvider implements QueueProvider {
  21.     private final static Logger log = LoggerFactory.getLogger(EmbedQueueProvider.class);
  23.     private Map<String, FileQueue<QueueTask>> fileQueues = new ConcurrentHashMap<>();
  25.     public EmbedQueueProvider(Properties props) {
  26.         int batch_size = NumberUtils.toInt(props.getProperty("embed.batch_size", "10000"), 10000);
  28.         Path path = checkoutPath(KooderConfig.getPath(props.getProperty("embed.path")));
  29.         for(String type : getAllTypes()) {
  30.             Path typePath = checkoutPath(path.resolve(type));
  31.             fileQueues.put(type, FileQueue.<QueueTask>batched().name(type)
  32.                     .folder(typePath)
  33.                     .restoreFromDisk(true)
  34.                     .batchSize(batch_size)
  35.                     .build());
  36.         }
  37.     }
  39.     private static Path checkoutPath(Path path) {
  40.         if(!Files.exists(path) || !Files.isDirectory(path)) {
  41.             log.warn("Path '{}' for queue storage not exists, created it!", path);
  42.             try {
  43.                 Files.createDirectories(path);
  44.             } catch(IOException e) {
  45.                 log.error("Failed to create directory '{}'", path, e);
  46.             }
  47.         }
  48.         return path;
  49.     }
  51.     /**
  52.      * 队列的唯一名称
  53.      *
  54.      * @return
  55.      */
  56.     @Override
  57.     public String name() {
  58.         return "embed";
  59.     }
  61.     /**
  62.      * 获取某个任务类型的队列
  63.      *
  64.      * @param type
  65.      * @return
  66.      */
  67.     @Override
  68.     public Queue queue(String type) {
  69.         return new Queue() {
  70.             @Override
  71.             public String type() {
  72.                 return type;
  73.             }
  75.             @Override
  76.             public void push(Collection<QueueTask> tasks) {
  77.                 fileQueues.get(type).addAll(tasks);
  78.             }
  80.             @Override
  81.             public List<QueueTask> pop(int count) {
  82.                 List<QueueTask> tasks = new ArrayList<>();
  83.                 QueueTask task;
  84.                 while(tasks.size() < count && (task = fileQueues.get(type).poll()) != null)
  85.                     tasks.add(task);
  86.                 return tasks;
  87.             }
  89.             @Override
  90.             public void close() {
  91.                 fileQueues.get(type).close();
  92.             }
  93.         };
  94.     }
  96.     @Override
  97.     public void close() {
  98.         fileQueues.values().forEach(q -> q.close());
  99.     }
  100. }
  1. package com.gitee.kooder.queue;
  3. import java.io.IOException;
  4. import java.io.Serializable;
  5. import java.util.ArrayList;
  6. import java.util.Arrays;
  7. import java.util.List;
  8. import java.util.Map;
  10. import com.fasterxml.jackson.annotation.JsonIgnore;
  11. import com.fasterxml.jackson.annotation.JsonProperty;
  12. import com.fasterxml.jackson.core.type.TypeReference;
  13. import com.gitee.kooder.models.CodeRepository;
  14. import com.gitee.kooder.core.Constants;
  15. import com.gitee.kooder.index.IndexManager;
  17. import com.gitee.kooder.models.Issue;
  18. import com.gitee.kooder.models.Repository;
  19. import com.gitee.kooder.models.Searchable;
  20. import com.gitee.kooder.utils.JsonUtils;
  21. import org.apache.commons.beanutils.BeanUtils;
  22. import org.apache.lucene.facet.taxonomy.TaxonomyWriter;
  23. import org.apache.lucene.index.IndexWriter;
  24. import org.slf4j.Logger;
  25. import org.slf4j.LoggerFactory;
  27. /**
  28.  * 队列中的任务
  29.  * @author Winter Lau<javayou@gmail.com>
  30.  */
  31. public class QueueTask implements Serializable {
  33.     private transient final static Logger log = LoggerFactory.getLogger(QueueTask.class);
  35.     public transient final static List<String> types = Arrays.asList(
  36.             Constants.TYPE_CODE,
  37.             Constants.TYPE_REPOSITORY,
  38.             Constants.TYPE_ISSUE,
  39.             Constants.TYPE_PR,
  40.             Constants.TYPE_COMMIT,
  41.             Constants.TYPE_WIKI,
  42.             Constants.TYPE_USER
  43.     );
  45.     public transient final static String ACTION_ADD            = "add"; //添加
  46.     public transient final static String ACTION_UPDATE         = "update"; //修改
  47.     public transient final static String ACTION_DELETE         = "delete"; //删除
  49.     private String type;    //对象类型
  50.     private String action;  //动作（添加、删除、修改）
  51.     private List<Searchable> objects = new ArrayList<>();    //objects list
  53.     public QueueTask(){}
  55.     public static void push(String type, String action, Searchable...obj){
  56.         QueueTask task = new QueueTask();
  57.         task.type = type;
  58.         task.action = action;
  59.         task.objects.addAll(Arrays.asList(obj));
  60.         QueueFactory.getProvider().queue(type).push(Arrays.asList(task));
  61.     }
  63.     public static void add(String type, Searchable...obj) {
  64.         push(type, ACTION_ADD, obj);
  65.     }
  67.     public static void update(String type, Searchable...obj) {
  68.         push(type, ACTION_UPDATE, obj);
  69.     }
  71.     public static void delete(String type, Searchable...obj) {
  72.         push(type, ACTION_DELETE, obj);
  73.     }
  75.     public String getType() {
  76.         return type;
  77.     }
  79.     public void setType(String type) {
  80.         this.type = type;
  81.     }
  83.     public final static boolean isAvailType(String p_type) {
  84.         return (p_type!=null) && types.contains(p_type.toLowerCase());
  85.     }
  87.     public final static boolean isAvailAction(String p_action) {
  88.         return ACTION_ADD.equalsIgnoreCase(p_action) || ACTION_DELETE.equalsIgnoreCase(p_action) || ACTION_UPDATE.equalsIgnoreCase(p_action);
  89.     }
  91.     public boolean isCodeTask() {
  92.         return Constants.TYPE_CODE.equals(type);
  93.     }
  95.     public String getAction() {
  96.         return action;
  97.     }
  99.     public void setAction(String action) {
  100.         this.action = action;
  101.     }
  103.     public List<Searchable> getObjects() {
  104.         return objects;
  105.     }
  108.     @JsonProperty("objects")
  109.     public void readObjects(Map<String,Object>[] values) throws Exception {
  110.         for(Map<String, Object> value : values) {
  111.             Searchable obj = null;
  112.             switch(type){
  113.                 case Constants.TYPE_CODE:
  114.                     obj = new CodeRepository();
  115.                     break;
  116.                 case Constants.TYPE_REPOSITORY:
  117.                     obj = new Repository();
  118.                     break;
  119.                 case Constants.TYPE_ISSUE:
  120.                     obj = new Issue();
  121.             }
  122.             BeanUtils.populate(obj, value);
  123.             objects.add(obj);
  124.         }
  125.     }
  127.     public void addObject(Searchable obj) {
  128.         objects.add(obj);
  129.     }
  131.     @JsonIgnore
  132.     public void setJsonObjects(String json) {
  133.         TypeReference typeRefer;
  134.         switch(type) {
  135.             case Constants.TYPE_CODE:
  136.                 typeRefer = new TypeReference<List<CodeRepository>>(){};
  137.                 break;
  138.             case Constants.TYPE_REPOSITORY:
  139.                 typeRefer = new TypeReference<List<Repository>>() {};
  140.                 break;
  141.             case Constants.TYPE_ISSUE:
  142.                 typeRefer = new TypeReference<List<Issue>>() {};
  143.                 break;
  144.             default:
  145.                 throw new IllegalArgumentException("Illegal task type: " + type);
  146.         }
  147.         this.objects = (List<Searchable>)JsonUtils.readValue(json, typeRefer);
  148.     }
  150.     /**
  151.      * 写入索引库
  152.      * @return
  153.      * @exception
  154.      */
  155.     public int write() throws IOException {
  156.         return IndexManager.write(this);
  157.     }
  159.     /**
  160.      * 用于多线程环境下共享 IndexWriter 写入
  161.      * @param i_writer
  162.      * @param t_writer
  163.      * @return
  164.      * @throws IOException
  165.      */
  166.     public int write(IndexWriter i_writer, TaxonomyWriter t_writer) throws IOException {
  167.         return IndexManager.write(this, i_writer, t_writer);
  168.     }
  170.     /**
  171.      * 生成 json
  172.      * @return
  173.      */
  174.     public String json() {
  175.         return JsonUtils.toJson(this);
  176.     }
  178.     public static QueueTask parse(String json) {
  179.         return JsonUtils.readValue(json, QueueTask.class);
  180.     }
  182.     @Override
  183.     public String toString() {
  184.         return "QueueTask{" +
  185.                 "type='" + type + '\'' +
  186.                 ", action='" + action + '\'' +
  187.                 ", objects=" + objects +
  188.                 '}';
  189.     }
  191.     public static void main(String[] args) {
  192.         String json = "{\"type\":\"code\",\"action\":\"add\",\"objects\":[{\"id\":379,\"doc_id\":0,\"doc_score\":0.0,\"enterprise\":10,\"scm\":\"git\",\"vender\":\"gitea\",\"name\":\"xxxxx\",\"url\":\"http://git.xxxxxx.com:3000/xxxx/xxxxx\",\"timestamp\":0,\"document\":{\"fields\":[{\"char_sequence_value\":\"379\"},{\"char_sequence_value\":\"gitea\"},{\"char_sequence_value\":\"10\"},{\"char_sequence_value\":\"http://git.xxxxx.com:3000/xxxx/xxxxx\"},{\"char_sequence_value\":\"xxxxx\"},{\"char_sequence_value\":\"git\"},{\"char_sequence_value\":\"1620462113883\"}]},\"relative_path\":\"000/000/000/xxxxx_379\",\"id_as_string\":\"379\"}],\"code_task\":true}";
  193.         QueueTask task = parse(json);
  194.         System.out.println(task);
  195.     }
  197. }
  1. public static void push(String type, String action, Searchable...obj){
  2.         QueueTask task = new QueueTask();
  3.         task.type = type;
  4.         task.action = action;
  5.         task.objects.addAll(Arrays.asList(obj));
  6.         QueueFactory.getProvider().queue(type).push(Arrays.asList(task));
  7.     }
  1.   List<QueueTask> tasks = provider.queue(type).pop(batch_fetch_count);