package org.apache.rocketmq.client.trace.hook;

import java.util.ArrayList;
import org.apache.rocketmq.client.hook.SendMessageContext;
import org.apache.rocketmq.client.hook.SendMessageHook;
import org.apache.rocketmq.client.producer.SendStatus;
import org.apache.rocketmq.client.trace.AsyncTraceDispatcher;
import org.apache.rocketmq.client.trace.TraceBean;
import org.apache.rocketmq.client.trace.TraceContext;
import org.apache.rocketmq.client.trace.TraceDispatcher;
import org.apache.rocketmq.client.trace.TraceType;
import org.apache.rocketmq.common.protocol.NamespaceUtil;

public class SendMessageTraceHookImpl implements SendMessageHook {

    private TraceDispatcher localDispatcher;

    public SendMessageTraceHookImpl(TraceDispatcher localDispatcher) {
        this.localDispatcher = localDispatcher;
    }

    @Override
    public String hookName() {
        return "SendMessageTraceHook";
    }

    /**
     * 这里并没有发送消息轨道数据，只是先准备好一部分数据，然后存储在发送上下文环境中
     */
    @Override
    public void sendMessageBefore(SendMessageContext context) {
        //if it is message trace data,then it doesn't recorded
        if (context == null || context.getMessage().getTopic().startsWith(((AsyncTraceDispatcher) localDispatcher).getTraceTopicName())) {
            return;
        }
        //build the context content of TuxeTraceContext
        TraceContext tuxeContext = new TraceContext();
        tuxeContext.setTraceBeans(new ArrayList<TraceBean>(1));
        context.setMqTraceContext(tuxeContext);
        tuxeContext.setTraceType(TraceType.Pub);
        tuxeContext.setGroupName(NamespaceUtil.withoutNamespace(context.getProducerGroup()));
        //build the data bean object of message trace
        TraceBean traceBean = new TraceBean();
        traceBean.setTopic(NamespaceUtil.withoutNamespace(context.getMessage().getTopic()));
        traceBean.setTags(context.getMessage().getTags());
        traceBean.setKeys(context.getMessage().getKeys());
        traceBean.setStoreHost(context.getBrokerAddr());
        traceBean.setBodyLength(context.getMessage().getBody().length);
        traceBean.setMsgType(context.getMsgType());
        tuxeContext.getTraceBeans().add(traceBean);
    }

    @Override
    public void sendMessageAfter(SendMessageContext context) {
        //if it is message trace data,then it doesn't recorded
        if (context == null || context.getMessage().getTopic().startsWith(((AsyncTraceDispatcher) localDispatcher).getTraceTopicName())
            || context.getMqTraceContext() == null) {
            return;
        }
        if (context.getSendResult() == null) {
            return;
        }

        if (context.getSendResult().getRegionId() == null
            || !context.getSendResult().isTraceOn()) {
            // if switch is false,skip it
            return;
        }

        TraceContext tuxeContext = (TraceContext) context.getMqTraceContext();
        TraceBean traceBean = tuxeContext.getTraceBeans().get(0);
        int costTime = (int) ((System.currentTimeMillis() - tuxeContext.getTimeStamp()) / tuxeContext.getTraceBeans().size());
        tuxeContext.setCostTime(costTime);
        if (context.getSendResult().getSendStatus().equals(SendStatus.SEND_OK)) {
            tuxeContext.setSuccess(true);
        } else {
            tuxeContext.setSuccess(false);
        }
        tuxeContext.setRegionId(context.getSendResult().getRegionId());
        traceBean.setMsgId(context.getSendResult().getMsgId());
        traceBean.setOffsetMsgId(context.getSendResult().getOffsetMsgId());
        traceBean.setStoreTime(tuxeContext.getTimeStamp() + costTime / 2);
        // 这一步才真正去发送轨迹数据消息   这一步存在数据丢失的可能
        localDispatcher.append(tuxeContext);
    }
}

package org.apache.rocketmq.client.trace;

import org.apache.rocketmq.client.AccessChannel;
import org.apache.rocketmq.client.exception.MQClientException;
import java.io.IOException;

/**
 * Interface of asynchronous transfer data
 */
public interface TraceDispatcher {
    enum Type {
        PRODUCE,
        CONSUME
    }
    /**
     * Initialize asynchronous transfer data module
     */
    void start(String nameSrvAddr, AccessChannel accessChannel) throws MQClientException;

    /**
     * Append the transfering data
     * @param ctx data information
     * @return
     */
    boolean append(Object ctx);

    /**
     * Write flush action
     *
     * @throws IOException
     */
    void flush() throws IOException;

    /**
     * Close the trace Hook
     */
    void shutdown();
}

public class AsyncTraceDispatcher implements TraceDispatcher {

private final ArrayBlockingQueue<TraceContext> traceContextQueue;

public AsyncTraceDispatcher(String group, Type type, String traceTopicName, RPCHook rpcHook) {
    // queueSize is greater than or equal to the n power of 2 of value
    this.queueSize = 2048;
    this.batchSize = 100;
    this.maxMsgSize = 128000;
    this.pollingTimeMil = 100;
    this.waitTimeThresholdMil = 500;
    this.discardCount = new AtomicLong(0L);
    this.traceContextQueue = new ArrayBlockingQueue<TraceContext>(1024);
    this.taskQueueByTopic = new HashMap();
    this.group = group;
    this.type = type;

    this.appenderQueue = new ArrayBlockingQueue<Runnable>(queueSize);
    if (!UtilAll.isBlank(traceTopicName)) {
        this.traceTopicName = traceTopicName;
    } else {
        this.traceTopicName = TopicValidator.RMQ_SYS_TRACE_TOPIC;
    }
    this.traceExecutor = new ThreadPoolExecutor(//
            10, //
            20, //
            1000 * 60, //
            TimeUnit.MILLISECONDS, //
            this.appenderQueue, //
            new ThreadFactoryImpl("MQTraceSendThread_"));
    traceProducer = getAndCreateTraceProducer(rpcHook);
}

public void start(String nameSrvAddr, AccessChannel accessChannel) throws MQClientException {
    if (isStarted.compareAndSet(false, true)) {
        traceProducer.setNamesrvAddr(nameSrvAddr);
        traceProducer.setInstanceName(TRACE_INSTANCE_NAME + "_" + nameSrvAddr);
        traceProducer.start();
    }
    this.accessChannel = accessChannel;
    this.worker = new Thread(new AsyncRunnable(), "MQ-AsyncTraceDispatcher-Thread-" + dispatcherId);
    this.worker.setDaemon(true);
    this.worker.start();
    this.registerShutDownHook();
}

@Override
public boolean append(final Object ctx) {
    boolean result = traceContextQueue.offer((TraceContext) ctx);
    if (!result) {
        log.info("buffer full" + discardCount.incrementAndGet() + " ,context is " + ctx);
    }
    return result;
}

}

class AsyncRunnable implements Runnable {
    private boolean stopped;

    @Override
    public void run() {
        while (!stopped) {
            synchronized (traceContextQueue) {
                long endTime = System.currentTimeMillis() + pollingTimeMil;
                while (System.currentTimeMillis() < endTime) {
                    try {
                        TraceContext traceContext = traceContextQueue.poll(
                                endTime - System.currentTimeMillis(), TimeUnit.MILLISECONDS
                        );

                        if (traceContext != null && !traceContext.getTraceBeans().isEmpty()) {
                            // get the topic which the trace message will send to
                            String traceTopicName = this.getTraceTopicName(traceContext.getRegionId());

                            // get the traceDataSegment which will save this trace message, create if null
                            TraceDataSegment traceDataSegment = taskQueueByTopic.get(traceTopicName);
                            if (traceDataSegment == null) {
                                traceDataSegment = new TraceDataSegment(traceTopicName, traceContext.getRegionId());
                                taskQueueByTopic.put(traceTopicName, traceDataSegment);
                            }

                            // encode traceContext and save it into traceDataSegment
                            // NOTE if data size in traceDataSegment more than maxMsgSize,
                            //  a AsyncDataSendTask will be created and submitted
                            TraceTransferBean traceTransferBean = TraceDataEncoder.encoderFromContextBean(traceContext);
                            traceDataSegment.addTraceTransferBean(traceTransferBean);
                        }
                    } catch (InterruptedException ignore) {
                        log.debug("traceContextQueue#poll exception");
                    }
                }

                // NOTE send the data in traceDataSegment which the first TraceTransferBean
                //  is longer than waitTimeThreshold
                sendDataByTimeThreshold();

                if (AsyncTraceDispatcher.this.stopped) {
                    this.stopped = true;
                }
            }
        }

    }

public class TraceContext implements Comparable<TraceContext> {
    private TraceType traceType;
    private long timeStamp = System.currentTimeMillis();
    private String regionId = "";
    private String regionName = "";
    private String groupName = "";
    private int costTime = 0;
    private boolean isSuccess = true;
    private String requestId = MessageClientIDSetter.createUniqID();
    private int contextCode = 0;
    private List<TraceBean> traceBeans;
}

public class TraceBean {
    private static final String LOCAL_ADDRESS = UtilAll.ipToIPv4Str(UtilAll.getIP());
    private String topic = "";
    private String msgId = "";
    private String offsetMsgId = "";
    private String tags = "";
    private String keys = "";
    private String storeHost = LOCAL_ADDRESS;
    private String clientHost = LOCAL_ADDRESS;
    private long storeTime;
    private int retryTimes;
    private int bodyLength;
    private MessageType msgType;
    private LocalTransactionState transactionState;
    private String transactionId;
    private boolean fromTransactionCheck;
}

public class TraceTransferBean {
    private String transData;
    private Set<String> transKey = new HashSet<String>();
}

class TraceDataSegment {
    private long firstBeanAddTime;
    private int currentMsgSize;
    private final String traceTopicName;
    private final String regionId;
    private final List<TraceTransferBean> traceTransferBeanList = new ArrayList();

    TraceDataSegment(String traceTopicName, String regionId) {
        this.traceTopicName = traceTopicName;
        this.regionId = regionId;
    }


}

    /**
     * 计算第一次添加的时间（用于超时500ms发送）
     * 如果累计消息大小超过消息最大大小，4m
     * 封装成AsyncDataSendTask提交到线程池执行
     * 清空firstBeanAddTime，currentMsgSize，traceTransferBeanList，方便重复利用
     */
    public void addTraceTransferBean(TraceTransferBean traceTransferBean) {
        initFirstBeanAddTime();
        this.traceTransferBeanList.add(traceTransferBean);
        this.currentMsgSize += traceTransferBean.getTransData().length();
        if (currentMsgSize >= traceProducer.getMaxMessageSize()) {
            List<TraceTransferBean> dataToSend = new ArrayList(traceTransferBeanList);
            AsyncDataSendTask asyncDataSendTask = new AsyncDataSendTask(traceTopicName, regionId, dataToSend);
            traceExecutor.submit(asyncDataSendTask);

            this.clear();

        }
    }