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.store;
  
  import com.actiontech.dble.backend.mysql.store.memalloc.MemSizeController;
  import com.actiontech.dble.backend.mysql.store.result.ResultExternal;
  import com.actiontech.dble.buffer.BufferPool;
  import com.actiontech.dble.buffer.BufferPoolRecord;
  import com.actiontech.dble.net.mysql.RowDataPacket;
  import com.actiontech.dble.plan.common.external.ResultStore;
  
  import java.util.ArrayList;
  import java.util.List;
  import java.util.concurrent.locks.Lock;
  import java.util.concurrent.locks.ReentrantLock;
  
  public abstract class LocalResult implements ResultStore {
  
      protected static final int DEFAULT_INITIAL_CAPACITY = 1024;
      protected final int fieldsCount;
      protected volatile int maxMemory = 262144;
  
      protected BufferPool pool;
      protected List<RowDataPacket> rows;
      protected ResultExternal external;
      protected int rowId, rowCount;
      protected int currentMemory;
      protected RowDataPacket currentRow;
      protected RowDataPacket lastRow;
      protected boolean isClosed;
      protected Lock lock;
      /* @bug 1208 */
      protected String charset = "UTF-8";
      protected MemSizeController bufferMC;
      protected volatile int maxReadMemorySize = -1;
      protected BufferPoolRecord.Builder bufferRecordBuilder;
  
      public LocalResult(int initialCapacity, int fieldsCount, BufferPool pool, String charset, BufferPoolRecord.Builder bufferRecordBuilder) {
          this.rows = new ArrayList<>(initialCapacity);
          this.fieldsCount = fieldsCount;
          this.pool = pool;
          this.bufferRecordBuilder = bufferRecordBuilder;
          init();
          this.isClosed = false;
          this.lock = new ReentrantLock();
          this.charset = charset;
      }
  
      public void setMaxMemory(int maxMemory) {
          this.maxMemory = maxMemory;
      }
  
      /**
       * add a row into localresult
       *
       * @param row
       */
      public void add(RowDataPacket row) {
          lock.lock();
          try {
              if (this.isClosed)
                  return;
              lastRow = row;
              rows.add(row);
              rowCount++;
              int incrementSize = getRowMemory(row);
              currentMemory += incrementSize;
              boolean needFlush = false;
              if (bufferMC != null) {
                  if (!bufferMC.addSize(incrementSize)) {
                      needFlush = true;
                  }
              } else if (currentMemory > maxMemory) {
                  needFlush = true;
              }
              if (needFlush) {
                  if (external == null)
                      external = makeExternal();
                  addRowsToDisk();
              }
          } finally {
              lock.unlock();
          }
      }
  
      protected abstract ResultExternal makeExternal();
  
      public RowDataPacket currentRow() {
          return currentRow;
      }
  
      public RowDataPacket getLastRow() {
          return lastRow;
      }
  
      public int getRowCount() {
          return rowCount;
      }
  
      public int getRowId() {
          return rowId;
      }
  
      /**
       * @return next row
       */
      public RowDataPacket next() {
          lock.lock();
          try {
              if (this.isClosed)
                  return null;
              if (++rowId < rowCount) {
                  if (external != null) {
                      currentRow = external.next();
                  } else {
                      currentRow = rows.get(rowId);
                  }
              } else {
                  currentRow = null;
              }
              return currentRow;
          } finally {
              lock.unlock();
          }
      }
  
      /**
       * This method is called after all rows have been added.
       */
      public void done() {
          lock.lock();
          try {
              if (this.isClosed)
                  return;
              if (external == null)
                  doneOnlyMemory();
              else {
                  if (!rows.isEmpty())
                      addRowsToDisk();
                  external.done();
              }
              reset();
          } finally {
              lock.unlock();
          }
      }
  
      protected abstract void doneOnlyMemory();
  
      public void reset() {
          lock.lock();
          try {
              rowId = -1;
              if (external != null) {
                  external.reset();
              }
          } finally {
              lock.unlock();
          }
      }
  
      @Override
      public void close() {
          lock.lock();
          try {
              if (this.isClosed)
                  return;
              this.isClosed = true;
              rows.clear();
              if (bufferMC != null)
                  bufferMC.subSize(currentMemory);
              rows = null;
              if (external != null) {
                  external.close();
                  external = null;
              }
          } finally {
              lock.unlock();
          }
      }
  
      @Override
      public void clear() {
          lock.lock();
          try {
              rows.clear();
              if (bufferMC != null)
                  bufferMC.subSize(currentMemory);
              init();
              if (external != null) {
                  external.close();
                  external = null;
              }
          } finally {
              lock.unlock();
          }
      }
  
      protected final void addRowsToDisk() {
          beforeFlushRows();
          rowCount = external.addRows(rows);
          rows.clear();
          if (bufferMC != null)
              bufferMC.subSize(currentMemory);
          currentMemory = 0;
      }
  
      /**
       * job to do before flush rows into disk
       */
      protected abstract void beforeFlushRows();
  
      protected int getRowMemory(RowDataPacket row) {
          return row.calcPacketSize();
      }
  
      private void init() {
          this.rowId = -1;
          this.rowCount = 0;
          this.currentMemory = 0;
          this.currentRow = null;
          this.lastRow = null;
      }
  
      public LocalResult setMemSizeController(MemSizeController memSizeController) {
          this.bufferMC = memSizeController;
          return this;
      }
  
      public void setMaxReadMemorySize(int maxReadMemorySize) {
          this.maxReadMemorySize = maxReadMemorySize;
      }
  }
  