content: /*
   * 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;
  
  import static org.junit.Assert.*;
  import static org.mockito.Mockito.*;
  
  import java.io.File;
  import java.io.FileOutputStream;
  import java.io.IOException;
  import java.nio.ByteBuffer;
  import java.nio.channels.FileChannel;
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.Map;
  import java.util.Random;
  import java.util.UUID;
  
  import com.google.common.base.Throwables;
  import io.netty.channel.Channel;
  import io.netty.channel.embedded.EmbeddedChannel;
  import org.junit.AfterClass;
  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.FileInfo;
  import org.apache.celeborn.common.network.buffer.NioManagedBuffer;
  import org.apache.celeborn.common.network.client.RpcResponseCallback;
  import org.apache.celeborn.common.network.client.TransportClient;
  import org.apache.celeborn.common.network.client.TransportResponseHandler;
  import org.apache.celeborn.common.network.protocol.ChunkFetchSuccess;
  import org.apache.celeborn.common.network.protocol.Message;
  import org.apache.celeborn.common.network.protocol.OpenStream;
  import org.apache.celeborn.common.network.protocol.RpcFailure;
  import org.apache.celeborn.common.network.protocol.RpcRequest;
  import org.apache.celeborn.common.network.protocol.RpcResponse;
  import org.apache.celeborn.common.network.protocol.StreamHandle;
  import org.apache.celeborn.common.network.protocol.TransportMessage;
  import org.apache.celeborn.common.network.util.TransportConf;
  import org.apache.celeborn.common.protocol.MessageType;
  import org.apache.celeborn.common.protocol.PbBufferStreamEnd;
  import org.apache.celeborn.common.protocol.PbChunkFetchRequest;
  import org.apache.celeborn.common.protocol.PbOpenStream;
  import org.apache.celeborn.common.protocol.PbStreamChunkSlice;
  import org.apache.celeborn.common.protocol.PbStreamHandler;
  import org.apache.celeborn.common.protocol.StreamType;
  import org.apache.celeborn.common.protocol.TransportModuleConstants;
  import org.apache.celeborn.common.unsafe.Platform;
  import org.apache.celeborn.common.util.JavaUtils;
  import org.apache.celeborn.common.util.Utils;
  import org.apache.celeborn.service.deploy.worker.memory.MemoryManager;
  import org.apache.celeborn.service.deploy.worker.storage.PartitionFilesSorter;
  import org.apache.celeborn.service.deploy.worker.storage.StorageManager;
  
  public class FetchHandlerSuiteJ {
  
    private static final Logger LOG = LoggerFactory.getLogger(FetchHandlerSuiteJ.class);
  
    private static final CelebornConf conf = new CelebornConf();
  
    private static final Random random = new Random();
  
    private static final UserIdentifier userIdentifier =
        new UserIdentifier("mock-tenantId", "mock-name");
    private static final int MAX_MAP_ID = 50;
    private static final int DATA_SIZE_PER_BATCH = 256 * 1024 - 16; // 256k - 16byte
  
    public FileInfo prepare(int batchCountPerMap) throws IOException {
      byte[] batchHeader = new byte[16];
      File shuffleFile = File.createTempFile("celeborn", UUID.randomUUID().toString());
  
      DiskFileInfo fileInfo = new DiskFileInfo(shuffleFile, userIdentifier, conf);
      FileOutputStream fileOutputStream = new FileOutputStream(shuffleFile);
      FileChannel channel = fileOutputStream.getChannel();
      Map<Integer, Integer> batchIds = new HashMap<>();
      ArrayList<Integer> mapIds = generateMapIds(batchCountPerMap);
  
      for (int mapId : mapIds) {
        int currentAttemptId = 0;
        int batchId =
            batchIds.compute(
                mapId,
                (k, v) -> {
                  if (v == null) {
                    v = 0;
                  } else {
                    v++;
                  }
                  return v;
                });
        byte[] mockedData = new byte[DATA_SIZE_PER_BATCH];
        Platform.putInt(batchHeader, Platform.BYTE_ARRAY_OFFSET, mapId);
        Platform.putInt(batchHeader, Platform.BYTE_ARRAY_OFFSET + 4, currentAttemptId);
        Platform.putInt(batchHeader, Platform.BYTE_ARRAY_OFFSET + 8, batchId);
        Platform.putInt(batchHeader, Platform.BYTE_ARRAY_OFFSET + 12, DATA_SIZE_PER_BATCH);
        ByteBuffer buf1 = ByteBuffer.wrap(batchHeader);
        while (buf1.hasRemaining()) {
          channel.write(buf1);
        }
        random.nextBytes(mockedData);
        ByteBuffer buf2 = ByteBuffer.wrap(mockedData);
        while (buf2.hasRemaining()) {
          channel.write(buf2);
        }
      }
      long originFileLen = channel.size();
      // update original fileInfo chunk offsets
      for (long offset = conf.shuffleChunkSize();
          offset <= originFileLen;
          offset += conf.shuffleChunkSize()) {
        (fileInfo.getReduceFileMeta()).addChunkOffset(offset);
      }
      // update sorted fileInfo chunk offsets
      fileInfo.updateBytesFlushed(originFileLen);
      return fileInfo;
    }
  
    public void cleanup(FileInfo fileInfo) throws IOException {
      if (fileInfo != null) {
        // origin file
        JavaUtils.deleteRecursively(((DiskFileInfo) fileInfo).getFile());
        // sorted file
        JavaUtils.deleteRecursively(
            new File(((DiskFileInfo) fileInfo).getFile().getPath() + ".sorted"));
        // index file
        JavaUtils.deleteRecursively(
            new File(((DiskFileInfo) fileInfo).getFile().getPath() + ".index"));
      }
    }
  
    @BeforeClass
    public static void beforeAll() {
      MemoryManager.initialize(conf);
    }
  
    @AfterClass
    public static void afterAll() {
      MemoryManager.reset();
    }
  
    @Test
    public void testFetchOriginFile() throws IOException {
      FileInfo fileInfo = null;
      try {
        // total write: 32 * 50 * 256k = 400m
        fileInfo = prepare(32);
        EmbeddedChannel channel = new EmbeddedChannel();
        TransportClient client = new TransportClient(channel, mock(TransportResponseHandler.class));
        FetchHandler fetchHandler = mockFetchHandler(fileInfo);
  
        PbStreamHandler streamHandler =
            openStreamAndCheck(client, channel, fetchHandler, 0, Integer.MAX_VALUE);
  
        fetchChunkAndCheck(client, channel, fetchHandler, streamHandler);
      } finally {
        cleanup(fileInfo);
      }
    }
  
    @Test
    public void testFetchSortFile() throws IOException {
      FileInfo fileInfo = null;
      try {
        // total write size: 32 * 50 * 256k = 400m
        fileInfo = prepare(32);
        EmbeddedChannel channel = new EmbeddedChannel();
        TransportClient client = new TransportClient(channel, mock(TransportResponseHandler.class));
        FetchHandler fetchHandler = mockFetchHandler(fileInfo);
  
        PbStreamHandler streamHandler = openStreamAndCheck(client, channel, fetchHandler, 5, 10);
  
        fetchChunkAndCheck(client, channel, fetchHandler, streamHandler);
      } finally {
        cleanup(fileInfo);
      }
    }
  
    @Test
    public void testLegacyOpenStream() throws IOException {
      FileInfo fileInfo = null;
      try {
        // total write size: 32 * 50 * 256k = 400m
        fileInfo = prepare(32);
        EmbeddedChannel channel = new EmbeddedChannel();
        TransportClient client = new TransportClient(channel, mock(TransportResponseHandler.class));
        FetchHandler fetchHandler = mockFetchHandler(fileInfo);
  
        legacyOpenStreamAndCheck(client, channel, fetchHandler, 0, Integer.MAX_VALUE);
  
      } finally {
        cleanup(fileInfo);
      }
    }
  
    @Test
    public void testWorkerReadSortFileOnceOriginalFileBeDeleted() throws IOException {
      FileInfo fileInfo = null;
      try {
        // total write size: 32 * 50 * 256k = 400m
        fileInfo = prepare(32);
        EmbeddedChannel channel = new EmbeddedChannel();
        TransportClient client = new TransportClient(channel, mock(TransportResponseHandler.class));
        FetchHandler fetchHandler = mockFetchHandler(fileInfo);
  
        PbStreamHandler rangeReadStreamHandler =
            openStreamAndCheck(client, channel, fetchHandler, 5, 10);
        checkOriginFileBeDeleted(fileInfo);
        PbStreamHandler nonRangeReadStreamHandler =
            openStreamAndCheck(client, channel, fetchHandler, 0, Integer.MAX_VALUE);
        fetchChunkAndCheck(client, channel, fetchHandler, nonRangeReadStreamHandler);
        fetchChunkAndCheck(client, channel, fetchHandler, rangeReadStreamHandler);
      } finally {
        cleanup(fileInfo);
      }
    }
  
    @Test
    public void testLocalReadSortFileOnceOriginalFileBeDeleted() throws IOException {
      FileInfo fileInfo = null;
      try {
        // total write size: 32 * 50 * 256k = 400m
        fileInfo = prepare(32);
        EmbeddedChannel channel = new EmbeddedChannel();
        TransportClient client = new TransportClient(channel, mock(TransportResponseHandler.class));
        FetchHandler fetchHandler = mockFetchHandler(fileInfo);
  
        // read local shuffle
        openStreamAndCheck(client, channel, fetchHandler, 5, 10, true);
        checkOriginFileBeDeleted(fileInfo);
      } finally {
        cleanup(fileInfo);
      }
    }
  
    @Test
    public void testDoNotDeleteOriginalFileWhenNonRangeWorkerReadWorkInProgress() throws IOException {
      FileInfo fileInfo = null;
      try {
        // total write size: 32 * 50 * 256k = 400m
        fileInfo = prepare(32);
        EmbeddedChannel channel = new EmbeddedChannel();
        TransportClient client = new TransportClient(channel, mock(TransportResponseHandler.class));
        FetchHandler fetchHandler = mockFetchHandler(fileInfo);
  
        PbStreamHandler nonRangeReadStreamHandler =
            openStreamAndCheck(client, channel, fetchHandler, 0, Integer.MAX_VALUE);
        PbStreamHandler rangeReadStreamHandler =
            openStreamAndCheck(client, channel, fetchHandler, 5, 10);
        fetchChunkAndCheck(client, channel, fetchHandler, nonRangeReadStreamHandler);
        fetchChunkAndCheck(client, channel, fetchHandler, rangeReadStreamHandler);
  
        // non-range fetch finished.
        bufferStreamEnd(client, fetchHandler, nonRangeReadStreamHandler.getStreamId());
        checkOriginFileBeDeleted(fileInfo);
      } finally {
        cleanup(fileInfo);
      }
    }
  
    @Test
    public void testDoNotDeleteOriginalFileWhenNonRangeLocalReadWorkInProgress() throws IOException {
      FileInfo fileInfo = null;
      try {
        // total write size: 32 * 50 * 256k = 400m
        fileInfo = prepare(32);
        EmbeddedChannel channel = new EmbeddedChannel();
        TransportClient client = new TransportClient(channel, mock(TransportResponseHandler.class));
        FetchHandler fetchHandler = mockFetchHandler(fileInfo);
  
        // read local shuffle
        PbStreamHandler nonRangeReadStreamHandler =
            openStreamAndCheck(client, channel, fetchHandler, 0, Integer.MAX_VALUE, true);
        openStreamAndCheck(client, channel, fetchHandler, 5, 10);
  
        // non-range fetch finished.
        bufferStreamEnd(client, fetchHandler, nonRangeReadStreamHandler.getStreamId());
        checkOriginFileBeDeleted(fileInfo);
      } finally {
        cleanup(fileInfo);
      }
    }
  
    private FetchHandler mockFetchHandler(FileInfo fileInfo) {
      WorkerSource workerSource = mock(WorkerSource.class);
      TransportConf transportConf =
          Utils.fromCelebornConf(conf, TransportModuleConstants.FETCH_MODULE, 4);
      FetchHandler fetchHandler0 = new FetchHandler(conf, transportConf, workerSource);
      Worker worker = mock(Worker.class);
      PartitionFilesSorter partitionFilesSorter =
          new PartitionFilesSorter(MemoryManager.instance(), conf, workerSource);
  
      StorageManager storageManager = mock(StorageManager.class);
      Mockito.doReturn(storageManager).when(worker).storageManager();
      Mockito.doReturn(workerSource).when(worker).workerSource();
      Mockito.doReturn(partitionFilesSorter).when(worker).partitionsSorter();
      fetchHandler0.init(worker);
      FetchHandler fetchHandler = spy(fetchHandler0);
      Mockito.doReturn(fileInfo).when(fetchHandler).getRawFileInfo(anyString(), anyString());
      return fetchHandler;
    }
  
    private final String shuffleKey = "dummyShuffleKey-123";
    private final String fileName = "dummyFileName";
    private final long dummyRequestId = 0;
  
    @Deprecated
    private void legacyOpenStreamAndCheck(
        TransportClient client,
        EmbeddedChannel channel,
        FetchHandler fetchHandler,
        int startIndex,
        int endIndex)
        throws IOException {
      ByteBuffer openStreamByteBuffer =
          new OpenStream(shuffleKey, fileName, startIndex, endIndex).toByteBuffer();
      fetchHandler.receive(
          client,
          new RpcRequest(dummyRequestId, new NioManagedBuffer(openStreamByteBuffer)),
          createRpcResponseCallback(channel));
      RpcResponse result = channel.readOutbound();
      StreamHandle streamHandler = (StreamHandle) Message.decode(result.body().nioByteBuffer());
      if (endIndex == Integer.MAX_VALUE) {
        assertEquals(50, streamHandler.numChunks);
      } else {
        assertEquals(endIndex - startIndex, streamHandler.numChunks);
      }
    }
  
    private PbStreamHandler openStreamAndCheck(
        TransportClient client,
        EmbeddedChannel channel,
        FetchHandler fetchHandler,
        int startIndex,
        int endIndex)
        throws IOException {
      return openStreamAndCheck(client, channel, fetchHandler, startIndex, endIndex, false);
    }
  
    private PbStreamHandler openStreamAndCheck(
        TransportClient client,
        EmbeddedChannel channel,
        FetchHandler fetchHandler,
        int startIndex,
        int endIndex,
        Boolean readLocalShuffle)
        throws IOException {
      ByteBuffer openStreamByteBuffer =
          new TransportMessage(
                  MessageType.OPEN_STREAM,
                  PbOpenStream.newBuilder()
                      .setShuffleKey(shuffleKey)
                      .setFileName(fileName)
                      .setStartIndex(startIndex)
                      .setEndIndex(endIndex)
                      .setReadLocalShuffle(readLocalShuffle)
                      .build()
                      .toByteArray())
              .toByteBuffer();
      fetchHandler.receive(
          client,
          new RpcRequest(dummyRequestId, new NioManagedBuffer(openStreamByteBuffer)),
          createRpcResponseCallback(channel));
      RpcResponse result = channel.readOutbound();
      PbStreamHandler streamHandler =
          TransportMessage.fromByteBuffer(result.body().nioByteBuffer()).getParsedPayload();
      if (endIndex == Integer.MAX_VALUE) {
        assertEquals(50, streamHandler.getNumChunks());
      } else {
        assertEquals(endIndex - startIndex, streamHandler.getNumChunks());
      }
      return streamHandler;
    }
  
    private void fetchChunkAndCheck(
        TransportClient client,
        EmbeddedChannel channel,
        FetchHandler fetchHandler,
        PbStreamHandler streamHandler) {
      for (int chunkIndex = 0; chunkIndex < streamHandler.getNumChunks(); chunkIndex++) {
        fetchHandler.receive(
            client,
            new RpcRequest(
                TransportClient.requestId(),
                new NioManagedBuffer(
                    new TransportMessage(
                            MessageType.CHUNK_FETCH_REQUEST,
                            PbChunkFetchRequest.newBuilder()
                                .setStreamChunkSlice(
                                    PbStreamChunkSlice.newBuilder()
                                        .setStreamId(streamHandler.getStreamId())
                                        .setChunkIndex(chunkIndex)
                                        .setOffset(0)
                                        .setLen(Integer.MAX_VALUE))
                                .build()
                                .toByteArray())
                        .toByteBuffer())),
            createRpcResponseCallback(channel));
        ChunkFetchSuccess chunkFetchSuccess = channel.readOutbound();
        chunkFetchSuccess.body().retain();
        // chunk size 8m
        assertEquals(chunkFetchSuccess.body().size(), 8 * 1024 * 1024);
        chunkFetchSuccess.body().release();
      }
    }
  
    private void bufferStreamEnd(TransportClient client, FetchHandler fetchHandler, long streamId) {
      TransportMessage bufferStreamEnd =
          new TransportMessage(
              MessageType.BUFFER_STREAM_END,
              PbBufferStreamEnd.newBuilder()
                  .setStreamId(streamId)
                  .setStreamType(StreamType.ChunkStream)
                  .build()
                  .toByteArray());
      fetchHandler.receive(
          client,
          new RpcRequest(dummyRequestId, new NioManagedBuffer(bufferStreamEnd.toByteBuffer())),
          createRpcResponseCallback(client.getChannel()));
    }
  
    private void checkOriginFileBeDeleted(FileInfo fileInfo) {
      assertFalse(((DiskFileInfo) fileInfo).getFilePath().endsWith(".sorted"));
      long startTs = System.currentTimeMillis();
      boolean deleted = false;
      long timeout = 5 * 1000; // 5s
      while (!deleted) {
        deleted = !((DiskFileInfo) fileInfo).getFile().exists();
        if (System.currentTimeMillis() - startTs > timeout) {
          fail("Origin file was not deleted within the expected timeout of 5 seconds.");
        }
        try {
          Thread.sleep(100);
        } catch (InterruptedException e) {
          LOG.error("catch Exception when checking origin file states.", e);
        }
      }
      assertTrue(deleted);
    }
  
    private ArrayList<Integer> generateMapIds(int batchCountPerMap) {
      int mapCount = batchCountPerMap * MAX_MAP_ID;
      ArrayList<Integer> ids = new ArrayList<>(mapCount);
      for (int i = 0; i < mapCount; i++) {
        ids.add(i / batchCountPerMap);
      }
      Collections.shuffle(ids);
      return ids;
    }
  
    private RpcResponseCallback createRpcResponseCallback(Channel channel) {
      return new RpcResponseCallback() {
        @Override
        public void onSuccess(ByteBuffer response) {
          channel.writeAndFlush(new RpcResponse(dummyRequestId, new NioManagedBuffer(response)));
        }
  
        @Override
        public void onFailure(Throwable e) {
          channel.writeAndFlush(new RpcFailure(dummyRequestId, Throwables.getStackTraceAsString(e)));
        }
      };
    }
  }
  