content: /*
   * Copyright (C) 2016-2023 ActionTech.
   * License: http://www.gnu.org/licenses/gpl.html GPL version 2 or higher.
   */
  
  package com.actiontech.dble.backend.mysql.nio.handler.query.impl;
  
  import com.actiontech.dble.DbleServer;
  import com.actiontech.dble.backend.mysql.nio.handler.query.BaseDMLHandler;
  import com.actiontech.dble.backend.mysql.nio.handler.util.ArrayMinHeap;
  import com.actiontech.dble.backend.mysql.nio.handler.util.HeapItem;
  import com.actiontech.dble.backend.mysql.nio.handler.util.RowDataComparator;
  import com.actiontech.dble.config.model.SystemConfig;
  import com.actiontech.dble.net.connection.BackendConnection;
  import com.actiontech.dble.net.mysql.FieldPacket;
  import com.actiontech.dble.net.mysql.RowDataPacket;
  import com.actiontech.dble.net.service.AbstractService;
  import com.actiontech.dble.plan.Order;
  import com.actiontech.dble.plan.common.exception.MySQLOutPutException;
  import com.actiontech.dble.route.RouteResultsetNode;
  import com.actiontech.dble.server.NonBlockingSession;
  import com.actiontech.dble.services.mysqlsharding.MySQLResponseService;
  import org.jetbrains.annotations.NotNull;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.util.*;
  import java.util.Map.Entry;
  import java.util.concurrent.BlockingQueue;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.LinkedBlockingQueue;
  
  /**
   * mergeHandler will merge data,if contains aggregate function,use group by handler
   *
   * @author ActionTech
   */
  public class MultiNodeMergeAndOrderHandler extends MultiNodeMergeHandler {
      private static final Logger LOGGER = LoggerFactory.getLogger(MultiNodeMergeAndOrderHandler.class);
  
      private final int queueSize;
      // map;conn->blocking queue.if receive row packet, add to the queue,if receive rowEof packet, add NullHeapItem into queue;
      private Map<MySQLResponseService, BlockingQueue<HeapItem>> queues;
      private List<Order> orderBys;
      private RowDataComparator rowComparator;
      private volatile boolean noNeedRows = false;
      private boolean nestLoopDependOn;
  
      public MultiNodeMergeAndOrderHandler(long id, RouteResultsetNode[] route, boolean autocommit, NonBlockingSession session,
                                           List<Order> orderBys, boolean nestLoopDependOn) {
          super(id, route, autocommit, session, true);
          this.orderBys = orderBys;
          this.queueSize = SystemConfig.getInstance().getMergeQueueSize();
          this.queues = new ConcurrentHashMap<>();
          this.merges.add(this);
          this.nestLoopDependOn = nestLoopDependOn;
  
      }
  
      @Override
      public void execute() {
          synchronized (exeHandlers) {
              if (terminate.get())
                  return;
  
              if (Thread.currentThread().getName().contains("complexQueryExecutor")) {
                  doExecute();
              } else {
                  DbleServer.getInstance().getComplexQueryExecutor().execute(new Runnable() {
                      @Override
                      public void run() {
                          doExecute();
                      }
                  });
              }
          }
      }
  
      private void doExecute() {
          for (BaseDMLHandler exeHandler : exeHandlers) {
              if (exeHandler instanceof BaseSelectHandler) {
                  BaseSelectHandler baseSelectHandler = (BaseSelectHandler) exeHandler;
                  session.setHandlerStart(baseSelectHandler); //base start execute
                  try {
                      BackendConnection exeConn = baseSelectHandler.initConnection();
                      exeConn.getBackendService().setComplexQuery(true);
                      queues.put(exeConn.getBackendService(), new LinkedBlockingQueue<>(queueSize));
                      baseSelectHandler.execute(exeConn.getBackendService());
                  } catch (Exception e) {
                      baseSelectHandler.connectionError(e, baseSelectHandler.getRrss());
                      return;
                  }
              }
          }
      }
  
      @Override
      public void fieldEofResponse(byte[] header, List<byte[]> fields, List<FieldPacket> fieldPackets, byte[] eof,
                                   boolean isLeft, @NotNull AbstractService service) {
          if (LOGGER.isDebugEnabled()) {
              LOGGER.debug(service.toString() + "'s field is reached.");
          }
          session.setHandlerStart(this);
          // if terminated
          if (terminate.get()) {
              return;
          }
          lock.lock(); // for combine
          try {
              if (this.fieldPackets.isEmpty()) {
                  this.fieldPackets = fieldPackets;
                  rowComparator = new RowDataComparator(this.fieldPackets, orderBys, this.isAllPushDown(), this.type());
                  nextHandler.fieldEofResponse(null, null, fieldPackets, null, this.isLeft, service);
              }
              if (++reachedConCount == route.length) {
                  session.trace(t -> t.allBackendConnReceive());
                  startOwnThread();
              }
          } finally {
              lock.unlock();
          }
      }
  
      @Override
      public boolean rowResponse(byte[] row, RowDataPacket rowPacket, boolean isLeft, @NotNull AbstractService service) {
          if (terminate.get() || noNeedRows)
              return true;
  
          MySQLResponseService mySQLConn = (MySQLResponseService) service;
          BlockingQueue<HeapItem> queue = queues.get(mySQLConn);
          if (queue == null)
              return true;
          HeapItem item = new HeapItem(row, rowPacket, mySQLConn);
          try {
              queue.put(item);
          } catch (InterruptedException e) {
              //ignore error
          }
          return false;
      }
  
      @Override
      public void rowEofResponse(byte[] data, boolean isLeft, @NotNull AbstractService service) {
          MySQLResponseService responseService = (MySQLResponseService) service;
          if (LOGGER.isDebugEnabled()) {
              LOGGER.debug(responseService.toString() + " 's rowEof is reached.");
          }
  
          if (this.terminate.get())
              return;
          BlockingQueue<HeapItem> queue = queues.get(responseService);
          if (queue == null)
              return;
          try {
              queue.put(HeapItem.nullItem());
          } catch (InterruptedException e) {
              //ignore error
          }
      }
  
      @Override
      protected void ownThreadJob(Object... objects) {
          try {
              ArrayMinHeap<HeapItem> heap = new ArrayMinHeap<>((o1, o2) -> {
                  RowDataPacket row1 = o1.getRowPacket();
                  RowDataPacket row2 = o2.getRowPacket();
                  if (row1 == null || row2 == null) {
                      if (row1 == row2)
                          return 0;
                      if (row1 == null)
                          return -1;
                      return 1;
                  }
                  return rowComparator.compare(row1, row2);
              });
              // init heap
              for (Entry<MySQLResponseService, BlockingQueue<HeapItem>> entry : queues.entrySet()) {
                  HeapItem firstItem = entry.getValue().take();
                  heap.add(firstItem);
              }
              while (!heap.isEmpty()) {
                  if (terminate.get() && !nestLoopDependOn)
                      return;
                  HeapItem top = heap.peak();
                  if (top.isNullItem()) {
                      heap.poll();
                  } else {
                      BlockingQueue<HeapItem> topItemQueue = queues.get(top.getIndex());
                      HeapItem item = topItemQueue.take();
                      heap.replaceTop(item);
                      if (nextHandler.rowResponse(top.getRowData(), top.getRowPacket(), this.isLeft, top.getIndex())) {
                          noNeedRows = true;
                          while (!heap.isEmpty()) {
                              HeapItem itemToDiscard = heap.poll();
                              if (!itemToDiscard.isNullItem()) {
                                  BlockingQueue<HeapItem> discardQueue = queues.get(itemToDiscard.getIndex());
                                  while (true) {
                                      if (Thread.currentThread().isInterrupted()) {
                                          throw new InterruptedException("manual interrupted");
                                      }
                                      if (discardQueue.take().isNullItem() || terminate.get()) {
                                          break;
                                      }
                                  }
                              }
                          }
                      }
                  }
              }
              if (LOGGER.isDebugEnabled()) {
                  String executeQueries = getRoutesSql(route);
                  LOGGER.debug(executeQueries + " heap send eof: ");
              }
              session.setHandlerEnd(this);
              nextHandler.rowEofResponse(null, this.isLeft, queues.keySet().iterator().next());
          } catch (MySQLOutPutException e) {
              String msg = e.getLocalizedMessage();
              LOGGER.info(msg, e);
              session.onQueryError(msg.getBytes());
          } catch (Exception e) {
              String msg = "Merge thread error, " + e.getLocalizedMessage();
              LOGGER.info(msg, e);
              session.onQueryError(msg.getBytes());
          }
      }
  
      @Override
      protected void terminateThread() throws Exception {
          for (Entry<MySQLResponseService, BlockingQueue<HeapItem>> entry : this.queues.entrySet()) {
              // add EOF to signal atoMerge thread
              entry.getValue().clear();
              entry.getValue().put(HeapItem.nullItem());
          }
          recycleConn();
      }
  
      @Override
      protected void recycleResources() {
          Iterator<Entry<MySQLResponseService, BlockingQueue<HeapItem>>> iterator = this.queues.entrySet().iterator();
          while (iterator.hasNext()) {
              Entry<MySQLResponseService, BlockingQueue<HeapItem>> entry = iterator.next();
              // fair lock queue,poll for clear
              while (true) {
                  if (Thread.currentThread().isInterrupted()) {
                      LOGGER.info("manual interrupted");
                      break;
                  }
                  if (entry.getValue().poll() == null) {
                      break;
                  }
              }
              iterator.remove();
          }
      }
  
      private String getRoutesSql(RouteResultsetNode[] nodes) {
          StringBuilder sb = new StringBuilder();
          sb.append('{');
          Map<String, List<RouteResultsetNode>> sqlMap = new HashMap<>();
          for (RouteResultsetNode rrss : nodes) {
              String sql = rrss.getStatement();
              if (!sqlMap.containsKey(sql)) {
                  List<RouteResultsetNode> rrssList = new ArrayList<>();
                  rrssList.add(rrss);
                  sqlMap.put(sql, rrssList);
              } else {
                  List<RouteResultsetNode> rrssList = sqlMap.get(sql);
                  rrssList.add(rrss);
              }
          }
          for (Entry<String, List<RouteResultsetNode>> entry : sqlMap.entrySet()) {
              sb.append(entry.getKey()).append(entry.getValue()).append(';');
          }
          sb.append('}');
          return sb.toString();
      }
  
      @Override
      public HandlerType type() {
          return HandlerType.MERGE_AND_ORDER;
      }
  }
  