PoolService.java

package com.viontech.match.service;

import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch._types.ElasticsearchException;
import co.elastic.clients.elasticsearch._types.mapping.AllField;
import co.elastic.clients.elasticsearch._types.mapping.Property;
import co.elastic.clients.elasticsearch._types.mapping.TypeMapping;
import co.elastic.clients.elasticsearch.core.BulkResponse;
import co.elastic.clients.elasticsearch.core.CountRequest;
import co.elastic.clients.elasticsearch.core.CountResponse;
import co.elastic.clients.elasticsearch.indices.*;
import co.elastic.clients.transport.endpoints.BooleanResponse;
import com.alibaba.fastjson.JSONObject;
//import com.fasterxml.jackson.databind.ObjectMapper;
import com.viontech.keliu.model.Person;
import com.viontech.match.config.Constant;
import com.viontech.match.entity.PoolInfo;
import com.viontech.match.entity.vo.RequestVo;
import com.viontech.match.entity.vo.ResponseVo;
import com.viontech.match.runner.ILM;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.http.HttpEntity;
import org.elasticsearch.client.Request;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.Response;
import org.elasticsearch.client.RestClient;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.*;

/**
 * .
 *
 * @author 谢明辉
 * @version 0.0.1
 */
@Service
@Slf4j
public class PoolService {
    @Resource
    private PersonService personService;
    @Resource
    private ElasticsearchClient client;
//    @Resource
//    private ObjectMapper objectMapper;
    @Value("${vion.index.number_of_shards:1}")
    private Integer shards;
    @Value("${vion.index.number_of_replicas:0}")
    private Integer replicas;
    @Value("${vion.index.translog.durability:async}")
    private String translogDurability;

//    @Value("${vion.index.refresh_interval}")
//    private String refreshInterval;
    @Value("${vion.index.translog.sync_interval}")
    private String translogSyncInterval;
    @Value("${vion.index.merge.scheduler.max_thread_count}")
    private Integer mergeThreadCount;

    /**
     * 添加特征池
     *
     * @param requestVo rid, poolId, personPool
     *
     * @return ResponseVo
     * @throws Exception --
     */
    public ResponseVo createPool(RequestVo requestVo) throws Exception {
        return createPool(requestVo, true);
    }

    public ResponseVo createPool(RequestVo requestVo, boolean addPerson) throws Exception {
        String rid = requestVo.getRid();
        String poolId = requestVo.getPoolId();
        log.info("特征池创建操作开始:[{}},IML:{}", poolId, requestVo.isUseILMPolicy());
        try {

            IndexSettings.Builder settings = new IndexSettings.Builder()
                    .numberOfShards(String.valueOf(shards))
                    .numberOfReplicas(String.valueOf(replicas));
            if (StringUtils.isNotEmpty(translogDurability)) {
                settings.translog(t -> t.durability(TranslogDurability.Async));
//                setting.put("index.translog.durability", translogDurability);
                if (StringUtils.isNotEmpty(translogSyncInterval)) {
                    settings.translog(t -> t.syncInterval(i -> i.time(translogSyncInterval)));
//                    setting.put("index.translog.sync_interval", translogSyncInterval);
                }
            }
            if (mergeThreadCount != null) {
//                setting.put("index.merge.scheduler.max_thread_count", mergeThreadCount);
                settings.merge(m -> m.scheduler(s -> s.maxThreadCount(mergeThreadCount)));
            }

            if (requestVo.isUseILMPolicy()) {
//                setting.put("index.lifecycle.name", ILM.LIFECYCLE_NAME);
                settings.lifecycle(l -> l.name(ILM.LIFECYCLE_NAME));
            }

            TypeMapping.Builder builder = getCreateIndexContentBuilder();
            CreateIndexRequest createIndexRequest = new CreateIndexRequest.Builder().index(poolId).mappings(builder.build()).settings(settings.build()).build();
            CreateIndexResponse createIndexResponse = client.indices().create(createIndexRequest);

            if (addPerson) {
                List<Person> personPool = requestVo.getPersonPool();
                if (CollectionUtils.isNotEmpty(personPool)) {
                    BulkResponse bulkItemResponses = personService.addPerson(poolId, personPool);
                    if (bulkItemResponses != null) {
                        log.info(bulkItemResponses.toString());
                    }
                }
            }

            log.info("特征池创建操作完成:[{}],ILM:{}", poolId, requestVo.isUseILMPolicy());
            return ResponseVo.success(rid);
        } catch (ElasticsearchException e) {
            if (e.status() == 400 && e.getMessage().contains(Constant.STR_POOL_ID_ALREADY_EXISTS)) {
                return ResponseVo.error(rid, 20, "already exit poolId");
            } else {
                throw e;
            }
        }
    }


    /**
     * 删除特征池
     *
     * @param requestVo rid,flushPool,poolId,personIds,
     *
     * @return ResponseVo
     * @throws Exception --
     */
    public ResponseVo deletePool(RequestVo requestVo) throws Exception {
        String rid = requestVo.getRid();
        Integer flushPool = requestVo.getFlushPool();
        String poolId = requestVo.getPoolId();
        log.info("特征池删除操作开始:[{}],FlushPool:[{}]", poolId, flushPool);
        try {
            if (0 == flushPool) {
                List<HashMap<String, String>> personIds = requestVo.getPersonIds();
                for (HashMap<String, String> item : personIds) {
                    String personId = item.get("personId");
                    if (personId != null) {
                        personService.deletePerson(poolId, personId);
                    }
                }
            } else if (flushPool == 1) {
                //删除索引
                DeleteIndexResponse delete = client.indices().delete(new DeleteIndexRequest.Builder().index(poolId).build());
            }
            log.info("特征池删除操作完成:[{}],FlushPool:[{}]", poolId, flushPool);
            return ResponseVo.success(rid);
        } catch (ElasticsearchException e) {
            if (e.status() == 404) {
                return ResponseVo.poolIdNotExists(rid);
            } else {
                return ResponseVo.error(rid, e.getMessage());
            }
        }
    }

    /**
     * 修改特征池(添加人员)
     *
     * @param requestVo rid,poolId,personPool,updateType
     *
     * @return ResponseVo
     * @throws Exception --
     */
    public ResponseVo modifyPool(RequestVo requestVo) throws Exception {
        String rid = requestVo.getRid();
        String poolId = requestVo.getPoolId();
        List<Person> personPool = requestVo.getPersonPool();
        Integer updateType = requestVo.getUpdateType();
        log.info("特征池修改操作开始:[{}],updateType:[{}]", poolId, updateType);
        co.elastic.clients.elasticsearch.indices.ExistsRequest poolIdExists = new ExistsRequest.Builder().index(poolId).build();
        BooleanResponse poolIdExistsResponse = client.indices().exists(poolIdExists);
        if (!poolIdExistsResponse.value()) {
            return ResponseVo.poolIdNotExists(rid);
        }
        try {
            BulkResponse bulkItemResponses = personService.addPerson(poolId, personPool);
            if (bulkItemResponses != null && bulkItemResponses.errors()) {
                log.info(bulkItemResponses.items().toString());
                return ResponseVo.error(rid, bulkItemResponses.items().toString());
            }

            log.info("特征池修改操作完成:[{}],updateType:[{}]", poolId, updateType);
            return ResponseVo.success(rid, "success");
        } catch (ElasticsearchException e) {
            log.error("modifyPool", e);
            return ResponseVo.error(rid, e.getMessage());
        }
    }

    /**
     * 查询特征池信息
     *
     * @param requestVo rid,listAll,poolId
     *
     * @return ResponseVo
     * @throws Exception --
     */
    public ResponseVo queryPool(RequestVo requestVo) throws Exception {
        String rid = requestVo.getRid();
        Integer listAll = requestVo.getListAll();
        String poolId = requestVo.getPoolId();
        try {
            List<PoolInfo> poolInfos;
            log.info("查询特征池操作开始:[{}],rid:[{}]", poolId, rid);

            if (listAll != 0) {
                poolInfos = queryPoolInfo(poolId);
            } else {
                if (!existPool(poolId)) {
                    return ResponseVo.poolIdNotExists(rid);
                }
                poolInfos = queryPoolInfo(poolId);
            }
            ResponseVo success = ResponseVo.success(rid, "success");
            success.setPoolIds(poolInfos);
            log.info("查询特征池操作完成:[{}],rid:[{}]", poolId, rid);

            return success;
        } catch (ElasticsearchException e) {
            log.error("queryPool", e);
            return ResponseVo.error(rid, e.getMessage());
        }
    }

    public void refreshPool(String... poolIds) throws IOException {
        RefreshRequest refreshRequest = new RefreshRequest.Builder().index(Arrays.asList(poolIds)).build();
        RefreshResponse refresh = client.indices().refresh(refreshRequest);
        log.info("刷新索引:{},成功:{},失败:{}", poolIds, refresh.shards().successful(), refresh.shards().failed());
    }


    public TypeMapping.Builder getCreateIndexContentBuilder() throws IOException {
//        JSONObject content = new JSONObject();
//        JSONObject properties = new JSONObject();
//        JSONObject data = new JSONObject();
//        data.put("type", "dense_vector");
//        data.put("dims", Constant.FACE_FEATURE_DIMS);
//        JSONObject body = new JSONObject();
//        body.put("type", "dense_vector");
//        body.put("dims", Constant.BODY_FEATURE_DIMS_2048);
//        JSONObject personId = new JSONObject();
//        personId.put("type", "keyword");
//
//        JSONObject unid = new JSONObject();
//        unid.put("type", "keyword");
//        JSONObject fid = new JSONObject();
//        fid.put("type", "text");
//        JSONObject age = new JSONObject();
//        age.put("type", "integer");
//        age.put("doc_values", false);
//        age.put("index", false);
//        JSONObject gender = new JSONObject();
//        gender.put("type", "keyword");
//        gender.put("doc_values", false);
//        gender.put("index", false);
//        JSONObject  counttime = new JSONObject();
//        counttime.put("type", "date");
//        counttime.put("format", "yyyy-MM-dd HH:mm:ss");
//        JSONObject channelSerialNum = new JSONObject();
//        channelSerialNum.put("type", "keyword");
//        JSONObject body_type = new JSONObject();
//        body_type.put("type", "integer");
//        body_type.put("doc_values", false);
//        body_type.put("index", false);
//        JSONObject direction = new JSONObject();
//        direction.put("type", "keyword");
//        JSONObject gateId = new JSONObject();
//        gateId.put("type", "long");
//        properties.put("data", data);
//        properties.put("body", body);
//        properties.put("personId", personId);
//        properties.put("unid", unid);
//        properties.put("fid", fid);
//        properties.put("age", age);
//        properties.put("gender", gender);
//        properties.put("counttime", counttime);
//        properties.put("channelSerialNum", channelSerialNum);
//        properties.put("body_type", body_type);
//        properties.put("gateId", gateId);
//        properties.put("direction", direction);
//
//        InputStream is = new ByteArrayInputStream(properties.toJSONString().getBytes());

        Map<String, Property> properties1 = new HashMap<>();
        properties1.put("direction", new Property.Builder().keyword(k -> k.index(true)).build());
        properties1.put("fid", new Property.Builder().text(t -> t.index(true)).build());
        properties1.put("gateId", new Property.Builder().long_(t -> t.index(true)).build());
        properties1.put("personId", new Property.Builder().keyword(k -> k.index(true)).build());
        properties1.put("body", new Property.Builder().denseVector(d -> d.dims(Constant.BODY_FEATURE_DIMS_2048).similarity("cosine")).build());
        properties1.put("data", new Property.Builder().denseVector(d -> d.dims(Constant.FACE_FEATURE_DIMS).similarity("cosine")).build());
        properties1.put("body_type", new Property.Builder().integer(i -> i.index(false).docValues(false)).build());
        properties1.put("channelSerialNum", new Property.Builder().keyword(k -> k.index(true)).build());
        properties1.put("age", new Property.Builder().integer(i -> i.index(false).docValues(false)).build());
        properties1.put("gender", new Property.Builder().keyword(k -> k.index(false).docValues(false)).build());
        properties1.put("counttime", new Property.Builder().date(d -> d.format("yyyy-MM-dd HH:mm:ss")).build());
        properties1.put("unid", new Property.Builder().keyword(k -> k.index(true)).build());

//        builder.properties(properties1);
        TypeMapping.Builder builder = new TypeMapping.Builder().properties(properties1);
        return builder;
    }

    public boolean existPool(String poolId) throws IOException {
        BooleanResponse exists = client.indices().exists(new ExistsRequest.Builder().index(poolId).build());
        return exists.value();
    }

    public List<PoolInfo> queryPoolInfo(String poolId) throws IOException {
        CountResponse response = client.count(new CountRequest.Builder().index(poolId).build());
        LinkedList<PoolInfo> poolInfos = new LinkedList<>();
        PoolInfo poolInfo = new PoolInfo();
        poolInfo.setPersonCount( response.count());
        poolInfo.setPoolId(poolId);
        poolInfos.add(poolInfo);

        return poolInfos;
    }

}