org.bdware.doip.cluster.util.ResultCollector Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of doip-audit-tool Show documentation
doip audit tool developed by bdware
There is a newer version: 1.5.4
package org.bdware.doip.cluster.util;

import com.google.gson.Gson;
import com.google.gson.JsonArray;
import com.google.gson.JsonElement;
import com.google.gson.JsonObject;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.bdware.doip.codec.JsonDoipMessage;
import org.bdware.doip.codec.doipMessage.DoipMessage;
import org.bdware.doip.codec.doipMessage.DoipMessageFactory;
import org.bdware.doip.codec.doipMessage.DoipResponseCode;
import org.bdware.doip.endpoint.client.DoipMessageCallback;
import org.bdware.sc.bean.JoinInfo;
import org.bdware.sc.util.JsonUtil;
import wrp.jdk.nashorn.api.scripting.NashornScriptEngineUtil;

import java.io.ByteArrayOutputStream;
import java.io.PrintStream;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;

public class ResultCollector implements DoipMessageCallback {
    static Logger LOGGER = LogManager.getLogger(ResultCollector.class);
    private final DoipMessage request;
    // joinCount记录开始Merge的上限
    int joinCount;
    AtomicInteger count = new AtomicInteger(0);
    DoipMessageCallback originalCallback;
    Logger Logger = LogManager.getLogger(ResultCollector.class);
    Map successResMap = new HashMap<>();
    Set problemResult = new HashSet<>();
    // 记录所有收到结果列表
    Set nodeIDs = new HashSet<>();
    NashornScriptEngineUtil engineUtil;
    JoinInfo joinInfo;


    public ResultCollector(DoipMessage msg, DoipMessageCallback cb, int count, NashornScriptEngineUtil engineUtil, JoinInfo joinInfo) {
        this.request = msg;
        this.originalCallback = cb;
        this.joinCount = count;
        this.engineUtil = engineUtil;
        this.joinInfo = joinInfo;
    }

    @Override
    public void onResult(DoipMessage msg) {
        logDoipMsgInfo(msg);
        // 如果nodeID不存在，或者该节点的返回值已经被处理过了，不需要再次调用onResult了
        String nodeID = getNodeID(msg);
        if (nodeID == null) {
            // 如果当前nodeID == null, 记录到problemResult中
            problemResult.add(msg);
        } else if (nodeIDs.contains(nodeID)) {
            LOGGER.info("Receive Duplicated Msg");
        } else {
            nodeIDs.add(nodeID);
            addDoipResult(nodeID, msg);
        }

        int curCount = count.incrementAndGet();
        // 小于joinCount，意味着还没达到聚合的标准，需要继续处理
        if (curCount <= joinCount) {
            // 达到聚合标准
            if (curCount == joinCount) {
                // 完成聚合 & 结果统计
                DoipMessage finalResult = executeJoinFunc(msg);
                originalCallback.onResult(finalResult);
            }
        }
    }

    public DoipMessageCallback getOriginalCallback() {
        return originalCallback;
    }

    public void logDoipMsgInfo(DoipMessage msg) {
        try {
            String nodeID = getNodeID(msg);
            if (nodeID != null) {
            } else {
                Logger.error("  收到错误的节点回复 : " + "Something wrong happens in the msg received from doip server, msg is : " + new Gson().toJson(msg.header.parameters) + " body:" + msg.body.getDataAsJsonString());
            }
        }catch (Exception e){
            e.printStackTrace();
        }
    }

    public synchronized void addDoipResult(String nodeID, DoipMessage msg) {
        successResMap.put(nodeID, msg);
    }


    public String getNodeID(DoipMessage msg) {
        String nodeID = null;
        if (msg != null && msg.header != null && msg.header.parameters != null && msg.header.parameters.attributes != null) {
            if (msg.header.parameters.attributes.get("nodeID") != null) {
                nodeID = msg.header.parameters.attributes.get("nodeID").getAsString();
            }
        }

        return nodeID;
    }

    public DoipMessage executeJoinFunc(DoipMessage request) {
        DoipMessage resultMessage;
        if (joinInfo != null && problemResult.size() == 0) {
            resultMessage = handleJoinInfo(joinInfo);
            return resultMessage;
        }
        if (successResMap.size() + problemResult.size() == 1) {
            problemResult.addAll(successResMap.values());
            for (DoipMessage message : problemResult)
                return message;
        }

        DoipMessageFactory.DoipMessageBuilder builder = new DoipMessageFactory.DoipMessageBuilder();
        if (problemResult.size() > 0)
            builder.createResponse(DoipResponseCode.UnKnownError, request);
        else
            builder.createResponse(DoipResponseCode.Success, request);
        JsonObject successResponses = new JsonObject();
        for (String key : successResMap.keySet()) {
            successResponses.add(key, JsonUtil.parseObject(JsonDoipMessage.fromDoipMessage(successResMap.get(key))));
        }
        JsonArray problemResponses = new JsonArray();
        for (DoipMessage msg : problemResult) {
            problemResponses.add(JsonUtil.parseObject(JsonDoipMessage.fromDoipMessage(msg)));
        }
        builder.addAttributes("successResponses", successResponses);
        builder.addAttributes("problemResponses", problemResponses);
        return builder.create();
    }

    private DoipMessage handleJoinInfo(JoinInfo joinInfo) {
        try {
            if (joinInfo.useDefault == null) {
                if (joinInfo.joinFuncName != null) {
                    JsonObject jo = new JsonObject();
                    for (String key : successResMap.keySet()) {
                        jo.add(key, JsonUtil.parseObject(JsonDoipMessage.fromDoipMessage(successResMap.get(key))));
                    }
                    JsonArray array = new JsonArray();
                    for (DoipMessage resp : problemResult) {
                        array.add(JsonUtil.parseObject(JsonDoipMessage.fromDoipMessage(resp)));
                    }
                    JsonElement returnVal = engineUtil.invokeFunction(joinInfo.joinFuncName, JsonElement.class, jo, array);
                    JsonDoipMessage result = JsonUtil.fromJson(returnVal, JsonDoipMessage.class);
                    return result.toResponseDoipMessage(request);
                } else throw new IllegalArgumentException("missing joinFuncName");
            }
            DoipMessage returnedMsg = null;
            switch (joinInfo.useDefault) {
                case add:
                    double val = 0;
                    for (DoipMessage msg : successResMap.values()) {
                        val += Double.valueOf(msg.body.getDataAsJsonString());
                        returnedMsg = msg;
                    }
                    returnedMsg.body.encodedData = (val + "").getBytes();
                    return returnedMsg;
                case multiply:
                    val = 1;
                    for (DoipMessage msg : successResMap.values()) {
                        val *= Double.valueOf(msg.body.getDataAsJsonString());
                        returnedMsg = msg;
                    }
                    return returnedMsg;
                default:
                    throw new IllegalArgumentException("missing default merge rule");
            }
        } catch (Exception e) {
            DoipMessageFactory.DoipMessageBuilder builder = new DoipMessageFactory.DoipMessageBuilder();
            builder.createResponse(DoipResponseCode.Declined, request);
            ByteArrayOutputStream bo = new ByteArrayOutputStream();
            e.printStackTrace(new PrintStream(bo));
            builder.setBody(bo.toByteArray());
            e.printStackTrace();
            return builder.create();
        }
    }
}