/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.celeborn.service.deploy.worker.memory;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.atomic.LongAdder;

import com.google.common.annotations.VisibleForTesting;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.ByteBufAllocator;
import io.netty.buffer.PooledByteBufAllocator;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.metrics.source.AbstractSource;
import org.apache.celeborn.common.network.util.NettyUtils;
import org.apache.celeborn.common.network.util.TransportConf;
import org.apache.celeborn.common.util.ThreadUtils;

public class ReadBufferDispatcher {
  private final Logger logger = LoggerFactory.getLogger(ReadBufferDispatcher.class);
  private final LinkedBlockingQueue<ReadBufferRequest> requests = new LinkedBlockingQueue<>();
  private final MemoryManager memoryManager;
  private final ByteBufAllocator readBufferAllocator;
  private final LongAdder allocatedReadBuffers = new LongAdder();
  private final long readBufferAllocationWait;
  private final long readBufferProcessTimeout;
  @VisibleForTesting public volatile boolean stopFlag = false;
  @VisibleForTesting public AtomicReference<Thread> dispatcherThread;
  private ScheduledExecutorService checkAliveThread;

  public ReadBufferDispatcher(
      MemoryManager memoryManager, CelebornConf conf, AbstractSource source) {
    readBufferAllocationWait = conf.readBufferAllocationWait();
    readBufferProcessTimeout = conf.readBufferProcessTimeout();
    long checkThreadInterval = conf.readBufferDispatcherCheckThreadInterval();
    // readBuffer is not a module name, it's a placeholder.
    readBufferAllocator =
        NettyUtils.getByteBufAllocator(new TransportConf("readBuffer", conf), source, true);
    this.memoryManager = memoryManager;
    dispatcherThread =
        new AtomicReference<>(
            ThreadUtils.newThread(new DispatcherRunnable(), "read-buffer-dispatcher"));
    dispatcherThread.get().start();

    if (checkThreadInterval > 0) {
      checkAliveThread =
          ThreadUtils.newDaemonSingleThreadScheduledExecutor("read-buffer-dispatcher-checker");
      checkAliveThread.scheduleWithFixedDelay(
          () -> {
            if (!dispatcherThread.get().isAlive()) {
              dispatcherThread.set(
                  ThreadUtils.newThread(new DispatcherRunnable(), "read-buffer-dispatcher"));
              dispatcherThread.get().start();
            }
          },
          checkThreadInterval,
          checkThreadInterval,
          TimeUnit.MILLISECONDS);
    }
  }

  public void addBufferRequest(ReadBufferRequest request) {
    requests.add(request);
  }

  public void recycle(ByteBuf buf) {
    int bufferSize = buf.capacity();
    int refCnt = buf.refCnt();
    // If a reader failed, related read buffers will have more than one reference count
    if (refCnt > 0) {
      buf.release(refCnt);
    }
    allocatedReadBuffers.decrement();
    memoryManager.changeReadBufferCounter(-1 * bufferSize);
  }

  public int requestsLength() {
    return requests.size();
  }

  public long getAllocatedReadBuffers() {
    return allocatedReadBuffers.sum();
  }

  public void close() {
    stopFlag = true;
    requests.clear();
    if (dispatcherThread != null) {
      dispatcherThread.get().interrupt();
      dispatcherThread = null;
    }
    if (checkAliveThread != null) {
      ThreadUtils.shutdown(checkAliveThread);
      checkAliveThread = null;
    }
  }

  private class DispatcherRunnable implements Runnable {

    public DispatcherRunnable() {}

    @Override
    public void run() {
      while (!stopFlag) {
        try {
          ReadBufferRequest request;
          request = requests.poll(1000, TimeUnit.MILLISECONDS);
          List<ByteBuf> buffers = new ArrayList<>();
          try {
            if (request != null) {
              processBufferRequest(request, buffers);
            } else {
              if (readBufferAllocator instanceof PooledByteBufAllocator) {
                // Free buffer pool memory to main direct memory when dispatcher is idle.
                ((PooledByteBufAllocator) readBufferAllocator).trimCurrentThreadCache();
              }
            }
          } catch (Throwable e) {
            logger.error(e.getMessage(), e);
            try {
              // recycle all allocated buffers
              for (ByteBuf buffer : buffers) {
                recycle(buffer);
              }
            } catch (Throwable e1) {
              logger.error("Recycle read buffer failed.", e1);
            }
            request.getBufferListener().notifyBuffers(null, e);
          }
        } catch (Throwable e) {
          logger.error("Read buffer dispatcher encountered error: {}", e.getMessage(), e);
        }
      }
    }

    void processBufferRequest(ReadBufferRequest request, List<ByteBuf> buffers) throws Exception {
      long start = System.nanoTime();
      int number = request.getNumber();
      int bufferSize = request.getBufferSize();
      while (buffers.size() < number && System.nanoTime() - start <= readBufferProcessTimeout) {
        if (memoryManager.readBufferAvailable(bufferSize)) {
          ByteBuf buf = readBufferAllocator.buffer(bufferSize, bufferSize);
          buffers.add(buf);
          memoryManager.changeReadBufferCounter(bufferSize);
          allocatedReadBuffers.increment();
        } else {
          try {
            // If dispatcher can not allocate requested buffers, it will wait here until
            // necessary buffers are get.
            Thread.sleep(readBufferAllocationWait);
          } catch (InterruptedException e) {
            logger.warn("Buffer dispatcher is closing");
          }
        }
      }
      if (buffers.size() == number) {
        logger.debug(
            "Processing {} read buffer (size: {}) request uses {} ms.",
            number,
            bufferSize,
            TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start));
        request.getBufferListener().notifyBuffers(buffers, null);
      } else {
        // The buffer dispatcher should add timeout constraints to fast fail in case of timeout,
        // avoiding long wait times for client.
        throw new TimeoutException(
            String.format(
                "Process %d read buffer (size: %d) request timeout after %d ms.",
                number, bufferSize, TimeUnit.NANOSECONDS.toMillis(readBufferProcessTimeout)));
      }
    }
  }
}
