/*
 * 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.storage;

import static org.apache.celeborn.common.util.JavaUtils.timeOutOrMeetCondition;

import java.io.File;
import java.io.IOException;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;

import io.netty.channel.Channel;
import org.junit.AfterClass;
import org.junit.Assert;
import org.junit.BeforeClass;
import org.junit.Test;
import org.mockito.Mockito;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.identity.UserIdentifier;
import org.apache.celeborn.common.meta.DiskFileInfo;
import org.apache.celeborn.common.meta.MapFileMeta;
import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.common.util.Utils;
import org.apache.celeborn.service.deploy.worker.memory.MemoryManager;

public class CreditStreamManagerSuiteJ {
  private static final Logger LOG = LoggerFactory.getLogger(CreditStreamManagerSuiteJ.class);
  private static File tempDir =
      Utils.createTempDir(System.getProperty("java.io.tmpdir"), "celeborn");
  private static CelebornConf conf = new CelebornConf();

  @BeforeClass
  public static void beforeAll() {
    conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_PAUSE_RECEIVE().key(), "0.8");
    conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_PAUSE_REPLICATE().key(), "0.9");
    conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_RESUME().key(), "0.5");
    conf.set(CelebornConf.WORKER_PARTITION_SORTER_DIRECT_MEMORY_RATIO_THRESHOLD().key(), "0.6");
    conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_FOR_READ_BUFFER().key(), "0.1");
    conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_FOR_MEMORY_FILE_STORAGE().key(), "0.1");
    conf.set(CelebornConf.WORKER_DIRECT_MEMORY_CHECK_INTERVAL().key(), "10");
    conf.set(CelebornConf.WORKER_DIRECT_MEMORY_REPORT_INTERVAL().key(), "10");
    conf.set(CelebornConf.WORKER_READBUFFER_ALLOCATIONWAIT().key(), "10ms");
    MemoryManager.initialize(conf);
  }

  private File createTemporaryFileWithIndexFile() throws IOException {
    String filename = UUID.randomUUID().toString();
    File temporaryFile = new File(tempDir, filename);
    File indexFile = new File(tempDir, filename + ".index");
    temporaryFile.createNewFile();
    indexFile.createNewFile();
    return temporaryFile;
  }

  @Test
  public void testStreamRegisterAndCleanup() throws Exception {
    CreditStreamManager creditStreamManager = new CreditStreamManager(10, 10, 1, 32);
    Channel channel = Mockito.mock(Channel.class);
    DiskFileInfo diskFileInfo =
        new DiskFileInfo(
            createTemporaryFileWithIndexFile(), new UserIdentifier("default", "default"), conf);
    MapFileMeta mapFileMeta = new MapFileMeta(1024, 10);
    mapFileMeta.setMountPoint("/tmp");
    diskFileInfo.replaceFileMeta(mapFileMeta);
    Consumer<Long> streamIdConsumer = streamId -> Assert.assertTrue(streamId > 0);

    String shuffleKey = "application_1694674023293_0003-0";
    long registerStream1 =
        creditStreamManager.registerStream(
            streamIdConsumer, channel, shuffleKey, 0, 1, 1, diskFileInfo);
    Assert.assertTrue(registerStream1 > 0);
    Assert.assertEquals(1, creditStreamManager.getStreamsCount());

    long registerStream2 =
        creditStreamManager.registerStream(
            streamIdConsumer, channel, shuffleKey, 0, 1, 1, diskFileInfo);
    Assert.assertNotEquals(registerStream1, registerStream2);
    Assert.assertEquals(2, creditStreamManager.getStreamsCount());

    creditStreamManager.registerStream(
        streamIdConsumer, channel, shuffleKey, 0, 1, 1, diskFileInfo);
    creditStreamManager.registerStream(
        streamIdConsumer, channel, shuffleKey, 0, 1, 1, diskFileInfo);

    MapPartitionData mapPartitionData1 =
        creditStreamManager.getStreams().get(registerStream1).getMapPartitionData();
    MapPartitionData mapPartitionData2 =
        creditStreamManager.getStreams().get(registerStream2).getMapPartitionData();
    Assert.assertEquals(mapPartitionData1, mapPartitionData2);

    mapPartitionData1.getStreamReader(registerStream1).recycle();

    timeOutOrMeetCondition(() -> creditStreamManager.getStreamsCount() == 3);
    Assert.assertEquals(creditStreamManager.numRecycleStreams(), 0);

    // registerStream2 can't be cleaned as registerStream2 is not finished
    AtomicInteger numInFlightRequests =
        mapPartitionData2.getStreamReader(registerStream2).getNumInUseBuffers();
    numInFlightRequests.incrementAndGet();

    creditStreamManager.cleanResource(registerStream2);
    Assert.assertEquals(creditStreamManager.numRecycleStreams(), 1);
    Assert.assertEquals(3, creditStreamManager.getStreamsCount());

    // recycle all channel
    numInFlightRequests.decrementAndGet();
    creditStreamManager.connectionTerminated(channel);
    timeOutOrMeetCondition(() -> creditStreamManager.getStreamsCount() == 0);
    // when cpu is busy, even through that timeOutOrMeetCondition is true,
    // creditStreamManager.getStreamsCount are still not be removed
    Assert.assertTrue(creditStreamManager.numRecycleStreams() >= 0);
  }

  @AfterClass
  public static void afterAll() {
    if (tempDir != null) {
      try {
        JavaUtils.deleteRecursively(tempDir);
        tempDir = null;
      } catch (IOException e) {
        LOG.error("Failed to delete temp dir.", e);
      }
    }
  }
}
