content: /*
   * Copyright (C) 2016-2023 ActionTech.
   * License: http://www.gnu.org/licenses/gpl.html GPL version 2 or higher.
   */
  
  package com.actiontech.dble.statistic.trace;
  
  import com.actiontech.dble.backend.mysql.nio.handler.ResponseHandler;
  import com.actiontech.dble.backend.mysql.nio.handler.builder.BaseHandlerBuilder;
  import com.actiontech.dble.backend.mysql.nio.handler.query.BaseDMLHandler;
  import com.actiontech.dble.backend.mysql.nio.handler.query.DMLResponseHandler;
  import com.actiontech.dble.backend.mysql.nio.handler.query.impl.MultiNodeMergeHandler;
  import com.actiontech.dble.btrace.provider.SlowLogProvider;
  import com.actiontech.dble.route.RouteResultsetNode;
  import com.actiontech.dble.route.parser.util.Pair;
  import com.actiontech.dble.server.NonBlockingSession;
  import com.actiontech.dble.server.status.SlowQueryLog;
  import com.actiontech.dble.services.mysqlsharding.MySQLResponseService;
  import com.actiontech.dble.statistic.sql.StatisticManager;
  import com.actiontech.dble.statistic.sql.entry.StatisticBackendSqlEntry;
  import com.actiontech.dble.statistic.sql.entry.StatisticFrontendSqlEntry;
  import com.google.common.collect.Lists;
  import com.google.common.collect.Maps;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.util.ArrayList;
  import java.util.List;
  import java.util.Map;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.CopyOnWriteArrayList;
  import java.util.concurrent.atomic.AtomicLong;
  import java.util.stream.Collectors;
  
  public class TraceResult implements Cloneable {
      public enum SqlTraceType {
          SINGLE_NODE_QUERY, MULTI_NODE_QUERY, MULTI_NODE_GROUP, COMPLEX_QUERY, SIMPLE_QUERY, COMPLEX_MODIFY
      }
  
      private static final Logger LOGGER = LoggerFactory.getLogger(TraceResult.class);
      protected long requestStart;
      protected long requestStartMs;
      protected long parseStart; // requestEnd
      protected long routeStart; // parseEnd
      protected long preExecuteStart; //r outeEnd
      protected RouteResultsetNode[] shardingNodes;
      protected long preExecuteEnd;
  
      protected long adtCommitBegin; // auto Distributed Transaction commit begin
      protected long adtCommitEnd; // auto Distributed Transaction commit end
  
      protected ResponseHandler simpleHandler = null;
      protected List<BackendRoute> backendRouteList = Lists.newCopyOnWriteArrayList();
      protected BaseHandlerBuilder builder = null; // for complex query
      protected Map<DMLResponseHandler, ComplexHandler> complexHandlerMap = Maps.newConcurrentMap();
      protected SqlTraceType type;
      protected long requestEnd;
      protected long requestEndMs;
      protected boolean subQuery = false;
  
      protected String sql;
      protected int sqlType = -1;
      protected String schema;
      protected ArrayList<String> tableList = new ArrayList<>();
      protected long sqlRows = 0;
      protected long netOutBytes;
      protected long resultSize;
      protected TraceResult previous = null;
      protected NonBlockingSession currentSession;
      /*
       * when 'set trace = 1' or 'enableSlowLog==true', need to record the time spent in each phase;
       *
       * so, into method:
       * endParse/endRoute/setPreExecuteEnd/setSubQuery/
       * setShardingNodes/setSimpleHandler/setBuilder/
       * setAdtCommitBegin/setAdtCommitEnd/addToRecordStartMap/addToRecordEndMap
       *
       */
      protected boolean isDetailTrace = false;
      // only samplingRate=100
      protected boolean pureRecordSql = true;
  
      public TraceResult(NonBlockingSession session0) {
          this.currentSession = session0;
      }
  
      public void setRequestTime(long time, long timeMs) {
          copyToPrevious();
          reset();
          this.isDetailTrace = currentSession.isTraceEnable() || SlowQueryLog.getInstance().isEnableSlowLog();
          this.pureRecordSql = !isDetailTrace && StatisticManager.getInstance().isPureRecordSql();
          this.requestStart = time;
          this.requestStartMs = timeMs;
      }
  
      public void startProcess(long time) {
          if (currentSession.getIsMultiStatement().get()) // multi-Query, need reset
              reset();
          this.parseStart = time;
      }
  
      public void setQuery(String sql0, int sqlType0) {
          this.schema = currentSession.getShardingService().getSchema();
          this.sql = sql0;
          this.sqlType = sqlType0;
      }
  
      public void addTable(List<Pair<String, String>> tables) { // schema.table
          for (Pair<String, String> p : tables) {
              tableList.add(p.getKey() + "." + p.getValue());
          }
      }
  
      public void endParse(long time) {
          if (!isDetailTrace) return;
          this.routeStart = time;
      }
  
      public void endRoute(long time) {
          if (!isDetailTrace) return;
          this.preExecuteStart = time;
      }
  
      public void setPreExecuteEnd(SqlTraceType type0, long time) {
          if (!isDetailTrace) return;
          if (routeStart != 0 && preExecuteStart != 0)
              this.type = type0;
          this.preExecuteEnd = time;
          backendRouteList.clear();
      }
  
      public void setSubQuery() {
          if (!isDetailTrace) return;
          this.subQuery = true;
      }
  
      public void setBackendRequestTime(MySQLResponseService service, long time) {
          final ResponseHandler responseHandler = service.getResponseHandler();
          if (responseHandler != null) {
              RouteResultsetNode node = (RouteResultsetNode) service.getAttachment();
              String key = service.getTraceRouteKey();
              BackendRoute ar = null;
              for (BackendRoute b : backendRouteList) {
                  if (b.handler == responseHandler && b.routeKey.equals(key)) {
                      ar = b;
                      break;
                  }
              }
              if (ar == null) {
                  ar = new BackendRoute(responseHandler, key, node.getName(), node.getStatement(), time);
                  backendRouteList.add(ar);
              }
          }
      }
  
      public void setBackendResponseTime(MySQLResponseService service, long time) {
          final ResponseHandler responseHandler = service.getResponseHandler();
          if (responseHandler != null) {
              String key = service.getTraceRouteKey();
              BackendRoute ar = null;
              for (BackendRoute b : backendRouteList) {
                  if (b.handler == responseHandler && b.routeKey.equals(key) && b.firstRevTime == 0) {
                      ar = b;
                      break;
                  }
              }
              if (ar != null) {
                  ar.setFirstRevTime(time);
                  ar.setMysqlResponseService(service);
              }
          }
      }
  
      public void setBackendSqlAddRows(MySQLResponseService service, Long num) {
          final ResponseHandler responseHandler = service.getResponseHandler();
          if (responseHandler != null) {
              String key = service.getTraceRouteKey();
              BackendRoute ar = null;
              for (BackendRoute b : backendRouteList) {
                  if (b.handler == responseHandler && b.routeKey.equals(key) && b.firstRevTime != 0) {
                      ar = b;
                      break;
                  }
              }
              if (ar != null) {
                  if (num == null) {
                      ar.getRow().incrementAndGet();
                  } else {
                      ar.getRow().set(num);
                  }
              }
          }
      }
  
      public void setBackendResponseEndTime(MySQLResponseService service, long time) {
          ResponseHandler responseHandler = service.getResponseHandler();
          if (responseHandler != null) {
              RouteResultsetNode node = (RouteResultsetNode) service.getAttachment();
              String key = service.getTraceRouteKey();
              BackendRoute ar = null;
              for (BackendRoute b : backendRouteList) {
                  if (b.handler == responseHandler && b.routeKey.equals(key) && b.firstRevTime != 0 && b.finished == 0) {
                      ar = b;
                      break;
                  }
              }
              if (ar != null) {
                  ar.setFinished(time);
                  ar.setAutocommit(service.isAutocommit());
                  if (pureRecordSql) return;
                  StatisticBackendSqlEntry bEntry = new StatisticBackendSqlEntry(
                          currentSession.getTraceFrontendInfo(),
                          service.getConnection().getTraceBackendInfo(), node.getName(),
                          ar.getRequestTime(), ar.getSql(), node.getSqlType(), ar.getRow().get(), ar.getFinished());
                  bEntry.setNeedToTx(ar.isAutocommit());
                  StatisticManager.getInstance().push(bEntry);
              }
          }
      }
  
      public void setBackendTerminateByComplex(MultiNodeMergeHandler mergeHandler, long time) {
          for (BaseDMLHandler handler : mergeHandler.getExeHandlers()) {
              BackendRoute ar = null;
              for (BackendRoute b : backendRouteList) {
                  if (b.handler == handler && b.firstRevTime != 0 && b.finished == 0) {
                      ar = b;
                      break;
                  }
              }
              if (ar != null) {
                  ar.setFinished(time);
                  if (pureRecordSql) return;
                  MySQLResponseService service;
                  if ((service = ar.getMysqlResponseService()) != null) {
                      RouteResultsetNode node = (RouteResultsetNode) service.getAttachment();
                      if (node != null) {
                          ar.setAutocommit(service.isAutocommit());
                          StatisticBackendSqlEntry bEntry = new StatisticBackendSqlEntry(
                                  currentSession.getTraceFrontendInfo(),
                                  service.getConnection().getTraceBackendInfo(), node.getName(),
                                  ar.getRequestTime(), ar.getSql(), node.getSqlType(), ar.getRow().get(), ar.getFinished());
                          bEntry.setNeedToTx(ar.isAutocommit());
                          StatisticManager.getInstance().push(bEntry);
                      }
                  }
              }
          }
      }
  
      // commit、rollback、quit
      public void setBackendResponseTxEnd(MySQLResponseService service, long time) {
          if (pureRecordSql) return;
          if (!service.isAutocommit()) {
              RouteResultsetNode node = (RouteResultsetNode) service.getAttachment();
              StatisticBackendSqlEntry bEntry = new StatisticBackendSqlEntry(
                      currentSession.getTraceFrontendInfo(),
                      service.getConnection().getTraceBackendInfo(), node.getName(),
                      time, "/** txEnd **/", 0, 0, time);
              bEntry.setNeedToTx(true);
              StatisticManager.getInstance().push(bEntry);
          }
      }
  
      public void setFrontendAddRows() {
          this.sqlRows += 1;
      }
  
      public void setFrontendSetRows(long rows) {
          this.sqlRows = rows;
      }
  
      public void setSqlStat(long sqlRows0, long netOutBytes0, long resultSize0) {
          //this.sqlRows = sqlRows0;
          this.netOutBytes = netOutBytes0;
          this.resultSize = resultSize0;
      }
  
      public void setResponseTime(boolean isSuccess, long time, long timeMs) {
          if (this.requestEnd == 0) {
              this.requestEnd = time;
              this.requestEndMs = timeMs;
              if (this.isCompletedV1() && isSuccess) {
                  long examinedRows = getExaminedRows();
                  StatisticFrontendSqlEntry f = new StatisticFrontendSqlEntry(currentSession.getTraceFrontendInfo(), requestStart, requestStartMs,
                          schema, sql, sqlType, currentSession.getShardingService().getTxId(), examinedRows, sqlRows,
                          netOutBytes, resultSize, requestEnd, requestEndMs, new ArrayList<String>(tableList));
                  StatisticManager.getInstance().push(f);
                  if (isDetailTrace && getOverAllMilliSecond() > SlowQueryLog.getInstance().getSlowTime()) {
                      SlowQueryLog.getInstance().putSlowQueryLog(currentSession.getShardingService(), this.clone());
                  }
                  if (LOGGER.isDebugEnabled()) {
                      LOGGER.debug("try to record sql: {}", sql);
                  }
              }
          }
      }
  
      private long getExaminedRows() {
          long examinedRows = 0;
          for (BackendRoute backendRoute : backendRouteList) {
              if (backendRoute.finished != 0) {
                  examinedRows += backendRoute.getRow().get();
              }
          }
          return examinedRows;
      }
  
      public void setExit() {
          reset();
          previous = null;
      }
  
      public void setShardingNodes(RouteResultsetNode[] shardingNodes) {
          if (!isDetailTrace) return;
          if (this.shardingNodes == null) {
              this.shardingNodes = shardingNodes;
          } else {
              RouteResultsetNode[] tempShardingNodes = new RouteResultsetNode[this.shardingNodes.length + shardingNodes.length];
              System.arraycopy(this.shardingNodes, 0, tempShardingNodes, 0, this.shardingNodes.length);
              System.arraycopy(shardingNodes, 0, tempShardingNodes, this.shardingNodes.length, shardingNodes.length);
              this.shardingNodes = tempShardingNodes;
              SlowLogProvider.setShardingNodes();
              if (LOGGER.isDebugEnabled()) {
                  LOGGER.debug("append to existing shardingNodes,current size is " + this.shardingNodes.length);
              }
          }
      }
  
      public void setSimpleHandler(ResponseHandler simpleHandler) {
          if (!isDetailTrace) return;
          this.simpleHandler = simpleHandler;
      }
  
      public void setBuilder(BaseHandlerBuilder builder) {
          if (!isDetailTrace) return;
          this.builder = builder;
      }
  
      public void setAdtCommitBegin(long time) {
          if (!isDetailTrace) return;
          this.adtCommitBegin = time;
      }
  
      public void setAdtCommitEnd(long time) {
          if (!isDetailTrace) return;
          this.adtCommitEnd = time;
      }
  
      public void addToRecordStartMap(DMLResponseHandler handler, long time) {
          if (!isDetailTrace) return;
          complexHandlerMap.putIfAbsent(handler, new ComplexHandler(handler, time));
      }
  
      public void addToRecordEndMap(DMLResponseHandler handler, long time) {
          if (!isDetailTrace) return;
          if (complexHandlerMap.containsKey(handler)) {
              complexHandlerMap.get(handler).setEndTime(time);
          }
      }
  
      private void reset() {
          parseStart = 0;
          routeStart = 0;
          preExecuteStart = 0;
          preExecuteEnd = 0;
          requestEnd = 0;
          requestEndMs = 0;
          shardingNodes = null;
          adtCommitBegin = 0;
          adtCommitEnd = 0;
          sql = null;
          sqlType = -1;
          schema = null;
          tableList.clear();
          sqlRows = 0;
          netOutBytes = 0;
          resultSize = 0;
          type = null;
          subQuery = false;
          simpleHandler = null;
          builder = null;
          backendRouteList.clear();
          complexHandlerMap.clear();
          isDetailTrace = false;
      }
  
      private void copyToPrevious() {
          if (isDetailTrace) {
              this.previous = this.clone();
              if (!previous.isCompletedV1()) {
                  previous = null;
              }
          } else {
              previous = null;
          }
      }
  
      public List<String[]> genRunningSQLStage() {
          return (new SelectTraceResult(this)).genRunningSQLStage();
      }
  
      // show trace
      public List<String[]> genShowTraceResult() {
          try {
              if (this.previous != null) {
                  return (new SelectTraceResult(this.previous)).genTraceResult();
              }
              return null;
          } catch (Exception e) {
              LOGGER.warn("genShowTraceResult exception {}", e);
              return null;
          } finally {
              this.previous = null;
          }
      }
  
      // slow log
      public List<String[]> genLogResult() {
          return (new SelectTraceResult(this)).genLogResult();
      }
  
      protected boolean isCompletedV1() {
          return sql != null && this.requestStart != 0 && this.requestEnd != 0;
      }
  
      protected boolean isCompletedV2() {
          if (isCompletedV1() && this.routeStart != 0) {
              long firstRevCount = this.backendRouteList.stream().filter(f -> f.getFirstRevTime() != 0).count();
              if (firstRevCount != 0) {
                  long finishedCount = this.backendRouteList.stream().filter(f -> f.getFinished() != 0).count();
                  if (firstRevCount == finishedCount) {
                      long recordStartCount = this.complexHandlerMap.size();
                      long recordEndCount = this.complexHandlerMap.values().stream().filter(f -> f.endTime != 0).count();
                      return recordStartCount == recordEndCount;
                  }
              }
          }
          return false;
      }
  
      protected boolean isNonBusinessSql() {
          return type == null; // || routeStart == null;
      }
  
      public RouteResultsetNode[] getShardingNodes() {
          return shardingNodes;
      }
  
      public SqlTraceType getType() {
          if (this.type == null) {
              return SqlTraceType.SIMPLE_QUERY;
          }
          return this.type;
      }
  
      public double getOverAllMilliSecond() {
          return (double) (this.requestEnd - this.requestStart) / 1000000;
      }
  
      public String getOverAllSecond() {
          double milliSecond = (double) (this.requestEnd - this.requestStart) / 1000000000;
          return String.format("%.6f", milliSecond);
      }
  
      @Override
      public TraceResult clone() {
          TraceResult tr;
          try {
              tr = (TraceResult) super.clone();
              tr.previous = null;
              tr.simpleHandler = this.simpleHandler;
              tr.builder = this.builder;
              tr.backendRouteList = new CopyOnWriteArrayList(this.backendRouteList);
              tr.complexHandlerMap = new ConcurrentHashMap<>(this.complexHandlerMap);
              tr.isDetailTrace = this.isDetailTrace;
              return tr;
          } catch (Exception e) {
              LOGGER.warn("clone TraceResult error", e);
              throw new AssertionError(e.getMessage());
          }
      }
  
      // find handler
      protected List<BackendRoute> findByHandler(ResponseHandler handler0) {
          return this.backendRouteList.stream().filter(f -> f.handler == handler0).collect(Collectors.toList());
      }
  
      protected ComplexHandler findByComplexHandler(ResponseHandler handler0) {
          return complexHandlerMap.get(handler0);
      }
  
      protected static class BackendRoute {
          ResponseHandler handler;
          MySQLResponseService mysqlResponseService;
          String routeKey;
          String shardingNode;
          String sql;
          long requestTime;
          long firstRevTime;
          long finished;
          AtomicLong row;
          boolean autocommit;
  
          public BackendRoute(ResponseHandler handler, String routeKey, String shardingNode, String sql, long requestTime) {
              this.handler = handler;
              this.routeKey = routeKey;
              this.requestTime = requestTime;
              this.shardingNode = shardingNode;
              this.sql = sql;
              this.row = new AtomicLong(0);
          }
  
          public long getRequestTime() {
              return requestTime;
          }
  
          public long getFirstRevTime() {
              return firstRevTime;
          }
  
          public void setFirstRevTime(long firstRevTime) {
              this.firstRevTime = firstRevTime;
          }
  
          public long getFinished() {
              return finished;
          }
  
          public void setFinished(long finished) {
              this.finished = finished;
          }
  
          public AtomicLong getRow() {
              return row;
          }
  
          public String getShardingNode() {
              return shardingNode;
          }
  
          public String getSql() {
              return sql;
          }
  
          public boolean isAutocommit() {
              return autocommit;
          }
  
          public void setAutocommit(boolean autocommit) {
              this.autocommit = autocommit;
          }
  
          public MySQLResponseService getMysqlResponseService() {
              return mysqlResponseService;
          }
  
          public void setMysqlResponseService(MySQLResponseService mysqlResponseService) {
              this.mysqlResponseService = mysqlResponseService;
          }
      }
  
      protected static class ComplexHandler {
          DMLResponseHandler handler;
          long startTime;
          long endTime;
  
          public ComplexHandler(DMLResponseHandler handler, long startTime) {
              this.handler = handler;
              this.startTime = startTime;
          }
  
          public long getStartTime() {
              return startTime;
          }
  
          public long getEndTime() {
              return endTime;
          }
  
          public void setEndTime(long endTime) {
              this.endTime = endTime;
          }
  
          public DMLResponseHandler getHandler() {
              return handler;
          }
      }
  }
  